integration testing

This commit adds integration testing for the `cmd/pilosa` package.
It uses randomly generated `set()` data and verifies the data using `get()`
both immediately after the insert and after restart.

bitmap.go

@@ -309,17 +309,7 @@ func (b *Bitmap) ReadFrom(r io.Reader) (n int64, err error) {
 }
 
 // MarshalJSON returns a JSON-encoded byte slice of b.
-func (b *Bitmap) MarshalJSON() ([]byte, error) {
-	o := bitmapJSON{
-		Chunks: make([]chunkJSON, 0, b.tree.Len()),
-	}
-
-	for itr := b.ChunkIterator(); !itr.Limit(); itr = itr.Next() {
-		o.Chunks = append(o.Chunks, chunkJSON{Key: itr.Item().Key, Value: itr.Item().Value})
-	}
-
-	return json.Marshal(&o)
-}
+func (b *Bitmap) MarshalJSON() ([]byte, error) { return json.Marshal(b.Bits()) }
 
 // MarshalBinary returns a gob-encoded byte slice of b.
 func (b *Bitmap) MarshalBinary() ([]byte, error) {
@@ -338,9 +328,8 @@ func (b *Bitmap) UnmarshalBinary(data []byte) error {
 
 // Bits returns the bits in b as a slice of ints.
 func (b *Bitmap) Bits() []uint64 {
-	result := make([]uint64, b.Count())
+	result := make([]uint64, 0, b.Count())
 
-	x := 0
 	for i := b.ChunkIterator(); !i.Limit(); i = i.Next() {
 		item := i.Item()
 		chunk := item.Key
@@ -349,8 +338,7 @@ func (b *Bitmap) Bits() []uint64 {
 				if (block & (1 << bit)) != 0 {
 					idx := chunk << 11
 					idx = idx | uint64((uint(bi)<<6)|bit)
-					result[x] = idx
-					x++
+					result = append(result, idx)
 				}
 			}
 		}
@@ -440,11 +428,6 @@ func Union(bitmaps []*Bitmap) *Bitmap {
 	return other
 }
 
-// bitmapJSON is the JSON representation of Bitmap.
-type bitmapJSON struct {
-	Chunks []chunkJSON `json:"chunks"`
-}
-
 // Chunk represents a set of blocks in a Bitmap.
 type Chunk struct {
 	Key   uint64
@@ -475,12 +458,6 @@ func decodeChunk(pb *internal.Chunk) *Chunk {
 	}
 }
 
-// chunkJSON is the JSON representation of Chunk.
-type chunkJSON struct {
-	Key   uint64
-	Value []uint64
-}
-
 // ChunkIterator represents an object for iterating over chunks in a bitmap.
 type ChunkIterator struct {
 	itr rbtree.Iterator

cmd/pilosa/config.go

@@ -1,69 +0,0 @@
-package main
-
-import (
-	"time"
-
-	"github.com/umbel/pilosa"
-)
-
-const (
-	// DefaultHost is the default hostname and port to use.
-	DefaultHost = "localhost:15000"
-)
-
-// Config represents the configuration for the command.
-type Config struct {
-	Host string `toml:"host"`
-
-	Cluster struct {
-		ReplicaN int           `toml:"replicas"`
-		Nodes    []*ConfigNode `toml:"nodes"`
-	} `toml:"cluster"`
-
-	Plugins struct {
-		Path string `toml:"path"`
-	} `toml:"plugins"`
-}
-
-type ConfigNode struct {
-	Host string `toml:"host"`
-}
-
-// NewConfig returns an instance of Config with default options.
-func NewConfig() *Config {
-	c := &Config{
-		Host: DefaultHost,
-	}
-	c.Cluster.ReplicaN = pilosa.DefaultReplicaN
-	c.Cluster.Nodes = []*ConfigNode{{Host: DefaultHost}}
-	return c
-}
-
-// PilosaCluster returns a new instance of pilosa.Cluster based on the config.
-func (c *Config) PilosaCluster() *pilosa.Cluster {
-	cluster := pilosa.NewCluster()
-	cluster.ReplicaN = c.Cluster.ReplicaN
-
-	for _, n := range c.Cluster.Nodes {
-		cluster.Nodes = append(cluster.Nodes, &pilosa.Node{Host: n.Host})
-	}
-
-	return cluster
-}
-
-// Duration is a TOML wrapper type for time.Duration.
-type Duration time.Duration
-
-// String returns the string representation of the duration.
-func (d Duration) String() string { return time.Duration(d).String() }
-
-// UnmarshalText parses a TOML value into a duration value.
-func (d *Duration) UnmarshalText(text []byte) error {
-	v, err := time.ParseDuration(string(text))
-	if err != nil {
-		return err
-	}
-
-	*d = Duration(v)
-	return nil
-}

cmd/pilosa/config_test.go

@@ -1,45 +0,0 @@
-package main_test
-
-import (
-	"testing"
-
-	"github.com/BurntSushi/toml"
-	"github.com/umbel/pilosa/cmd/pilosa"
-)
-
-// Ensure the host can be parsed.
-func TestConfig_Parse_Host(t *testing.T) {
-	if c, err := ParseConfig(`host = "local"`); err != nil {
-		t.Fatal(err)
-	} else if c.Host != "local" {
-		t.Fatalf("unexpected host: %s", c.Host)
-	}
-}
-
-// Ensure the addr can be parsed.
-func TestConfig_Parse_Addr(t *testing.T) {
-	if c, err := ParseConfig(`addr = ":80"`); err != nil {
-		t.Fatal(err)
-	} else if c.Addr != ":80" {
-		t.Fatalf("unexpected addr: %s", c.Addr)
-	}
-}
-
-// Ensure the "plugins" config can be parsed.
-func TestConfig_Parse_Plugins(t *testing.T) {
-	if c, err := ParseConfig(`
-[plugins]
-path = "/path/to/plugins"
-`); err != nil {
-		t.Fatal(err)
-	} else if c.Plugins.Path != "/path/to/plugins" {
-		t.Fatalf("unexpected path: %s", c.Plugins.Path)
-	}
-}
-
-// ParseConfig parses s into a config.
-func ParseConfig(s string) (main.Config, error) {
-	var c main.Config
-	_, err := toml.Decode(s, &c)
-	return c, err
-}

cmd/pilosa/main.go

@@ -5,12 +5,14 @@ import (
 	"flag"
 	"fmt"
 	"io"
-	"log"
 	"math/rand"
 	"net"
 	"net/http"
 	"os"
+	"os/user"
+	"path/filepath"
 	"runtime/pprof"
+	"strconv"
 	"time"
 
 	"github.com/BurntSushi/toml"
@@ -28,6 +30,11 @@ func init() {
 	rand.Seed(time.Now().UTC().UnixNano())
 }
 
+const (
+	// DefaultHost is the default hostname and port to use.
+	DefaultHost = "localhost:15000"
+)
+
 func main() {
 	m := NewMain()
 	fmt.Fprintf(m.Stderr, "Pilosa %s\n", Build)
@@ -50,7 +57,8 @@ func main() {
 
 // Main represents the main program execution.
 type Main struct {
-	ln net.Listener
+	index *pilosa.Index
+	ln    net.Listener
 
 	// Path to the configuration file.
 	ConfigPath string
@@ -78,22 +86,27 @@ func NewMain() *Main {
 	}
 }
 
+// Addr returns the address of the listener.
+func (m *Main) Addr() net.Addr {
+	if m.ln == nil {
+		return nil
+	}
+	return m.ln.Addr()
+}
+
 // Run executes the main program execution.
 func (m *Main) Run(args ...string) error {
-	logger := log.New(m.Stderr, "", log.LstdFlags)
-
 	// Notify user of config file.
 	if m.ConfigPath != "" {
 		fmt.Fprintf(m.Stdout, "Using config: %s\n", m.ConfigPath)
 	}
 
 	// Require a port in the hostname.
-	_, addr, err := net.SplitHostPort(m.Config.Host)
+	host, port, err := net.SplitHostPort(m.Config.Host)
 	if err != nil {
 		return err
-	} else if addr == "" {
+	} else if port == "" {
 		return errors.New("port must be specified in config host")
-
 	}
 
 	// Set up profiling.
@@ -107,15 +120,31 @@ func (m *Main) Run(args ...string) error {
 		defer pprof.StopCPUProfile()
 	}
 
-	// Build cluster from config file.
+	// Open HTTP listener to determine port (if specified as :0).
+	ln, err := net.Listen("tcp", ":"+port)
+	if err != nil {
+		return err
+	}
+	m.ln = ln
+
+	// Determine hostname based on listening port.
+	hostname := net.JoinHostPort(host, strconv.Itoa(m.ln.Addr().(*net.TCPAddr).Port))
+
+	// Build cluster from config file. Create local host if none are specified.
 	cluster := m.Config.PilosaCluster()
+	if len(cluster.Nodes) == 0 {
+		cluster.Nodes = []*pilosa.Node{{
+			Host: hostname,
+		}}
+	}
 
 	// Create index to store fragments.
-	index := pilosa.NewIndex("/tmp/")
+	fmt.Fprintf(m.Stderr, "Using data from: %s\n", m.Config.DataDir)
+	m.index = pilosa.NewIndex(m.Config.DataDir)
 
 	// Create executor for executing queries.
-	e := pilosa.NewExecutor(index)
-	e.Host = m.Config.Host
+	e := pilosa.NewExecutor(m.index)
+	e.Host = hostname
 	e.Cluster = cluster
 
 	// Initialize HTTP handler.
@@ -123,17 +152,10 @@ func (m *Main) Run(args ...string) error {
 	h.Executor = e
 	h.LogOutput = m.Stderr
 
-	// Open HTTP listener.
-	ln, err := net.Listen("tcp", ":"+addr)
-	if err != nil {
-		return err
-	}
-	m.ln = ln
-
 	// Serve HTTP.
-	go func() { logger.Print(http.Serve(ln, h)) }()
+	go func() { http.Serve(ln, h) }()
 
-	fmt.Fprintf(m.Stderr, "Listening on http://%s\n", ln.Addr().String())
+	fmt.Fprintf(m.Stderr, "Listening as http://%s\n", hostname)
 
 	return nil
 }
@@ -143,6 +165,11 @@ func (m *Main) Close() error {
 	if m.ln != nil {
 		m.ln.Close()
 	}
+
+	if m.index != nil {
+		m.index.Close()
+	}
+
 	return nil
 }
 
@@ -163,5 +190,73 @@ func (m *Main) ParseFlags(args []string) error {
 		}
 	}
 
+	// If no data directory is specified then use ~/.pilosa
+	if m.Config.DataDir == "" {
+		u, err := user.Current()
+		if err != nil {
+			return err
+		} else if u.HomeDir == "" {
+			return errors.New("data directory not specified and no home dir available")
+		}
+		m.Config.DataDir = filepath.Join(u.HomeDir, ".pilosa")
+	}
+
+	return nil
+}
+
+// Config represents the configuration for the command.
+type Config struct {
+	DataDir string `toml:"data-dir"`
+	Host    string `toml:"host"`
+
+	Cluster struct {
+		ReplicaN int           `toml:"replicas"`
+		Nodes    []*ConfigNode `toml:"nodes"`
+	} `toml:"cluster"`
+
+	Plugins struct {
+		Path string `toml:"path"`
+	} `toml:"plugins"`
+}
+
+type ConfigNode struct {
+	Host string `toml:"host"`
+}
+
+// NewConfig returns an instance of Config with default options.
+func NewConfig() *Config {
+	c := &Config{
+		Host: DefaultHost,
+	}
+	c.Cluster.ReplicaN = pilosa.DefaultReplicaN
+	return c
+}
+
+// PilosaCluster returns a new instance of pilosa.Cluster based on the config.
+func (c *Config) PilosaCluster() *pilosa.Cluster {
+	cluster := pilosa.NewCluster()
+	cluster.ReplicaN = c.Cluster.ReplicaN
+
+	for _, n := range c.Cluster.Nodes {
+		cluster.Nodes = append(cluster.Nodes, &pilosa.Node{Host: n.Host})
+	}
+
+	return cluster
+}
+
+// Duration is a TOML wrapper type for time.Duration.
+type Duration time.Duration
+
+// String returns the string representation of the duration.
+func (d Duration) String() string { return time.Duration(d).String() }
+
+// UnmarshalText parses a TOML value into a duration value.
+func (d *Duration) UnmarshalText(text []byte) error {
+	v, err := time.ParseDuration(string(text))
+	if err != nil {
+		return err
+	}
+
+	*d = Duration(v)
 	return nil
 }

cmd/pilosa/main_test.go

@@ -1 +1,277 @@
 package main_test
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"math/rand"
+	"net/http"
+	"os"
+	"reflect"
+	"sort"
+	"strings"
+	"testing"
+	"testing/quick"
+
+	"github.com/BurntSushi/toml"
+	main "github.com/umbel/pilosa/cmd/pilosa"
+)
+
+// Ensure program can process queries and maintain consistency.
+func TestMain_Set_Quick(t *testing.T) {
+	if err := quick.Check(func(cmds []SetCommand) bool {
+		m := MustRunMain()
+		defer m.Close()
+
+		// Execute set() commands.
+		for _, cmd := range cmds {
+			if res, err := m.Query("d", fmt.Sprintf(`set(id=%d, frame=%q, profile_id=%d)`, cmd.ID, cmd.Frame, cmd.ProfileID)); err != nil {
+				t.Fatal(err)
+			} else if res != `{}`+"\n" {
+				t.Fatalf("unexpected result: %s", res)
+			}
+		}
+
+		// Validate data.
+		for frame, frameSet := range SetCommands(cmds).Frames() {
+			for id, profileIDs := range frameSet {
+				exp := MustMarshalJSON(map[string]interface{}{"result": profileIDs}) + "\n"
+				if res, err := m.Query("d", fmt.Sprintf(`get(id=%d, frame=%q)`, id, frame)); err != nil {
+					t.Fatal(err)
+				} else if res != exp {
+					t.Fatalf("unexpected result:\n\ngot=%s\n\nexp=%s\n\n", res, exp)
+				}
+			}
+		}
+
+		if err := m.Reopen(); err != nil {
+			t.Fatal(err)
+		}
+
+		// Validate data after reopening.
+		for frame, frameSet := range SetCommands(cmds).Frames() {
+			for id, profileIDs := range frameSet {
+				exp := MustMarshalJSON(map[string]interface{}{"result": profileIDs}) + "\n"
+				if res, err := m.Query("d", fmt.Sprintf(`get(id=%d, frame=%q)`, id, frame)); err != nil {
+					t.Fatal(err)
+				} else if res != exp {
+					t.Fatalf("unexpected result:\n\ngot=%s\n\nexp=%s\n\n", res, exp)
+				}
+			}
+		}
+
+		return true
+	}, &quick.Config{
+		Values: func(values []reflect.Value, rand *rand.Rand) {
+			values[0] = reflect.ValueOf(GenerateSetCommands(1000, rand))
+		},
+	}); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// Ensure the host can be parsed.
+func TestConfig_Parse_Host(t *testing.T) {
+	if c, err := ParseConfig(`host = "local"`); err != nil {
+		t.Fatal(err)
+	} else if c.Host != "local" {
+		t.Fatalf("unexpected host: %s", c.Host)
+	}
+}
+
+// Ensure the data directory can be parsed.
+func TestConfig_Parse_DataDir(t *testing.T) {
+	if c, err := ParseConfig(`data-dir = "/tmp/foo"`); err != nil {
+		t.Fatal(err)
+	} else if c.DataDir != "/tmp/foo" {
+		t.Fatalf("unexpected data dir: %s", c.DataDir)
+	}
+}
+
+// Ensure the "plugins" config can be parsed.
+func TestConfig_Parse_Plugins(t *testing.T) {
+	if c, err := ParseConfig(`
+[plugins]
+path = "/path/to/plugins"
+`); err != nil {
+		t.Fatal(err)
+	} else if c.Plugins.Path != "/path/to/plugins" {
+		t.Fatalf("unexpected path: %s", c.Plugins.Path)
+	}
+}
+
+// Main represents a test wrapper for main.Main.
+type Main struct {
+	*main.Main
+
+	Stdin  bytes.Buffer
+	Stdout bytes.Buffer
+	Stderr bytes.Buffer
+}
+
+// NewMain returns a new instance of Main with a temporary data directory and random port.
+func NewMain() *Main {
+	path, err := ioutil.TempDir("", "pilosa-")
+	if err != nil {
+		panic(err)
+	}
+
+	m := &Main{Main: main.NewMain()}
+	m.Config.DataDir = path
+	m.Config.Host = "localhost:0"
+	m.Main.Stdin = &m.Stdin
+	m.Main.Stdout = &m.Stdout
+	m.Main.Stderr = &m.Stderr
+
+	if testing.Verbose() {
+		m.Main.Stdout = io.MultiWriter(os.Stdout, m.Main.Stdout)
+		m.Main.Stderr = io.MultiWriter(os.Stderr, m.Main.Stderr)
+	}
+
+	return m
+}
+
+// MustRunMain returns a new, running Main. Panic on error.
+func MustRunMain() *Main {
+	m := NewMain()
+	if err := m.Run(); err != nil {
+		panic(err)
+	}
+	return m
+}
+
+// Close closes the program and removes the underlying data directory.
+func (m *Main) Close() error {
+	defer os.RemoveAll(m.Config.DataDir)
+	return m.Main.Close()
+}
+
+// Reopen closes the program and reopens it.
+func (m *Main) Reopen() error {
+	if err := m.Main.Close(); err != nil {
+		return err
+	}
+
+	// Create new main with the same config.
+	config := m.Config
+	m.Main = main.NewMain()
+	m.Config = config
+
+	// Run new program.
+	if err := m.Run(); err != nil {
+		return err
+	}
+	return nil
+}
+
+// URL returns the base URL string for accessing the running program.
+func (m *Main) URL() string { return "http://" + m.Addr().String() }
+
+// Query executes a query against the program through the HTTP API.
+func (m *Main) Query(db, query string) (string, error) {
+	resp := MustDo("POST", m.URL()+"/query?db="+db, query)
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("invalid status: %d, body=%s", resp.StatusCode, resp.Body)
+	}
+	return resp.Body, nil
+}
+
+// SetCommand represents a command to set a bit.
+type SetCommand struct {
+	ID        uint64
+	Frame     string
+	ProfileID uint64
+}
+
+type SetCommands []SetCommand
+
+// Frames returns the set of profile ids for each frame/bitmap.
+func (a SetCommands) Frames() map[string]map[uint64][]uint64 {
+	// Create a set of unique commands.
+	m := make(map[SetCommand]struct{})
+	for _, cmd := range a {
+		m[cmd] = struct{}{}
+	}
+
+	// Build unique ids for each frame & bitmap.
+	frames := make(map[string]map[uint64][]uint64)
+	for cmd := range m {
+		if frames[cmd.Frame] == nil {
+			frames[cmd.Frame] = make(map[uint64][]uint64)
+		}
+		frames[cmd.Frame][cmd.ID] = append(frames[cmd.Frame][cmd.ID], cmd.ProfileID)
+	}
+
+	// Sort each set of profile ids.
+	for _, frame := range frames {
+		for id := range frame {
+			sort.Sort(uint64Slice(frame[id]))
+		}
+	}
+
+	return frames
+}
+
+// GenerateSetCommands generates random SetCommand objects.
+func GenerateSetCommands(n int, rand *rand.Rand) []SetCommand {
+	cmds := make([]SetCommand, rand.Intn(n))
+	for i := range cmds {
+		cmds[i] = SetCommand{
+			ID:        uint64(rand.Intn(1000)),
+			Frame:     "x.n",
+			ProfileID: uint64(rand.Intn(10)),
+		}
+	}
+	return cmds
+}
+
+// ParseConfig parses s into a Config.
+func ParseConfig(s string) (main.Config, error) {
+	var c main.Config
+	_, err := toml.Decode(s, &c)
+	return c, err
+}
+
+// MustDo executes http.Do() with an http.NewRequest(). Panic on error.
+func MustDo(method, urlStr string, body string) *httpResponse {
+	req, err := http.NewRequest(method, urlStr, strings.NewReader(body))
+	if err != nil {
+		panic(err)
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		panic(err)
+	}
+	defer resp.Body.Close()
+
+	buf, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		panic(err)
+	}
+
+	return &httpResponse{Response: resp, Body: string(buf)}
+}
+
+// httpResponse is a wrapper for http.Response that holds the Body as a string.
+type httpResponse struct {
+	*http.Response
+	Body string
+}
+
+// MustMarshalJSON marshals v into a string. Panic on error.
+func MustMarshalJSON(v interface{}) string {
+	buf, err := json.Marshal(v)
+	if err != nil {
+		panic(err)
+	}
+	return string(buf)
+}
+
+// uint64Slice represents a sortable slice of uint64 numbers.
+type uint64Slice []uint64
+
+func (p uint64Slice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
+func (p uint64Slice) Len() int           { return len(p) }
+func (p uint64Slice) Less(i, j int) bool { return p[i] < p[j] }

fragment.go

@@ -71,13 +71,13 @@ func (f *Fragment) Open() error {
 		// Open the data file to be mmap'd and used as an ops log.
 		file, err := os.OpenFile(f.path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
 		if err != nil {
-			return err
+			return fmt.Errorf("open file: %s", err)
 		}
 		f.file = file
 
 		// Lock the underlying file.
-		if err := syscall.Flock(int(f.file.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err == nil {
-			return nil
+		if err := syscall.Flock(int(f.file.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err != nil {
+			return fmt.Errorf("flock: %s", err)
 		}
 
 		// If the file is empty then initialize it with an empty bitmap.
@@ -86,6 +86,11 @@ func (f *Fragment) Open() error {
 			return err
 		} else if fi.Size() == 0 {
 			if _, err := f.storage.WriteTo(f.file); err != nil {
+				return fmt.Errorf("init storage file: %s", err)
+			}
+
+			fi, err = f.file.Stat()
+			if err != nil {
 				return err
 			}
 		}
@@ -103,8 +108,9 @@ func (f *Fragment) Open() error {
 		}
 
 		// Attach the mmap file to the bitmap.
-		if err := f.storage.UnmarshalBinary((*[0x7FFFFFFF]byte)(unsafe.Pointer(&f.storageData[0]))[:]); err != nil {
-			return fmt.Errorf("unmarshal storage: %s", err)
+		data := (*[0x7FFFFFFF]byte)(unsafe.Pointer(&f.storageData[0]))[:fi.Size()]
+		if err := f.storage.UnmarshalBinary(data); err != nil {
+			return fmt.Errorf("unmarshal storage: file=%s, err=%s", f.file.Name(), err)
 		}
 
 		// Attach the file to the bitmap to act as a write-ahead log.

handler.go

@@ -69,6 +69,9 @@ func (h *Handler) handleQuery(w http.ResponseWriter, r *http.Request) {
 func (h *Handler) handlePostQuery(w http.ResponseWriter, r *http.Request) {
 	// Parse incoming request.
 	db, query, slices, err := h.readQueryRequest(r)
+
+	// h.logger().Printf("%s %s db=%s q=%s slices=%v", r.Method, r.URL.Path, db, query, slices)
+
 	if err != nil {
 		w.WriteHeader(http.StatusBadRequest)
 		h.writeQueryResponse(w, r, nil, err)