add 'pilosactl bench' command

This commit adds a simple benchmarking utility to the `pilosactl`
binary. It currently only supports individual `SetBit()` commands
but it's a good start towards making a generic benchmarking
framework at the integration level.

The subcommands and usage/help messages were also cleaned up to
output correctly.

attr.go

@@ -56,7 +56,12 @@ func (s *AttrStore) Open() error {
 }
 
 // Close closes the store.
-func (s *AttrStore) Close() error { return s.db.Close() }
+func (s *AttrStore) Close() error {
+	if s.db != nil {
+		s.db.Close()
+	}
+	return nil
+}
 
 // Attrs returns a set of attributes by ID.
 func (s *AttrStore) Attrs(id uint64) (m map[string]interface{}, err error) {

client.go

@@ -92,6 +92,58 @@ func (c *Client) SliceNodes(slice uint64) ([]*Node, error) {
 	return a, nil
 }
 
+// ExecuteQuery executes query against db on the server.
+func (c *Client) ExecuteQuery(db, query string) (result interface{}, err error) {
+	if db == "" {
+		return nil, ErrDatabaseRequired
+	} else if query == "" {
+		return nil, ErrQueryRequired
+	}
+
+	// Encode query request.
+	buf, err := proto.Marshal(&internal.QueryRequest{
+		DB:    proto.String(db),
+		Query: proto.String(query),
+	})
+	if err != nil {
+		return nil, fmt.Errorf("marshal: %s", err)
+	}
+
+	// Create URL & HTTP request.
+	u := url.URL{Scheme: "http", Host: c.host, Path: "/query"}
+	req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf))
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Content-Length", strconv.Itoa(len(buf)))
+	req.Header.Set("Content-Type", "application/x-protobuf")
+	req.Header.Set("Accept", "application/x-protobuf")
+
+	// Execute request against the host.
+	resp, err := c.HTTPClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	// Read body and unmarshal response.
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	} else if resp.StatusCode != http.StatusOK {
+		return nil, errors.New(string(body))
+	}
+
+	var qresp internal.QueryResponse
+	if err := proto.Unmarshal(body, &qresp); err != nil {
+		return nil, fmt.Errorf("unmarshal response: %s", err)
+	} else if s := qresp.GetErr(); s != "" {
+		return nil, errors.New(s)
+	}
+
+	return nil, nil
+}
+
 // Import bulk imports bits for a single slice to a host.
 func (c *Client) Import(db, frame string, slice uint64, bits []Bit) error {
 	if db == "" {

cmd/pilosactl/main.go

@@ -6,25 +6,21 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"log"
+	"math/rand"
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/umbel/pilosa"
 )
 
 var (
-	// ErrUsage is returned when usage should be displayed for the program.
-	ErrUsage = errors.New("usage")
-
 	// ErrUnknownCommand is returned when specifying an unknown command.
 	ErrUnknownCommand = errors.New("unknown command")
 
-	// ErrQuit is returned when the program should simply quit.
-	// This is used when the error message has already been printed.
-	ErrQuit = errors.New("quit")
-
 	// ErrPathRequired is returned when executing a command without a required path.
 	ErrPathRequired = errors.New("path required")
 )
@@ -33,15 +29,15 @@ func main() {
 	m := NewMain()
 
 	// Parse command line arguments.
-	if err := m.ParseFlags(os.Args[1:]); err != nil {
+	if err := m.ParseFlags(os.Args[1:]); err == flag.ErrHelp {
+		os.Exit(2)
+	} else if err != nil {
 		fmt.Fprintln(m.Stderr, err)
 		os.Exit(2)
 	}
 
 	// Execute the program.
-	if err := m.Run(); err == ErrQuit {
-		os.Exit(1)
-	} else if err != nil {
+	if err := m.Run(); err != nil {
 		fmt.Fprintln(m.Stderr, err)
 		os.Exit(1)
 	}
@@ -49,9 +45,8 @@ func main() {
 
 // Main represents the main program execution.
 type Main struct {
-	// Command name and arguments passed into the CLI.
-	Command string
-	Args    []string
+	// Subcommand to execute.
+	Cmd Command
 
 	// Standard input/output
 	Stdin  io.Reader
@@ -68,47 +63,66 @@ func NewMain() *Main {
 	}
 }
 
+// Usage returns the usage message to be printed.
+func (m *Main) Usage() string {
+	return strings.TrimSpace(`
+Pilosactl is a tool for interacting with a pilosa server.
+
+Usage:
+
+	pilosactl command [arguments]
+
+The commands are:
+
+	config     prints the default configuration
+	import     imports data from a CSV file
+	backup     backs up a frame to an archive file
+	restore    restores a frame from an archive file
+	bench      benchmarks operations
+
+Use the "-h" flag with any command for more information.
+`)
+}
+
 // Run executes the main program execution.
-func (m *Main) Run() error {
-	var cmd Command
-	switch m.Command {
+func (m *Main) Run() error { return m.Cmd.Run() }
+
+// ParseFlags parses command line flags from args.
+func (m *Main) ParseFlags(args []string) error {
+	var command string
+	if len(args) > 0 {
+		command = args[0]
+		args = args[1:]
+	}
+
+	switch command {
 	case "", "help", "-h":
-		return ErrUsage
+		fmt.Fprintln(m.Stderr, m.Usage())
+		fmt.Fprintln(m.Stderr, "")
+		return flag.ErrHelp
 	case "config":
-		cmd = NewConfigCommand(m.Stdin, m.Stdout, m.Stderr)
+		m.Cmd = NewConfigCommand(m.Stdin, m.Stdout, m.Stderr)
 	case "import":
-		cmd = NewImportCommand(m.Stdin, m.Stdout, m.Stderr)
+		m.Cmd = NewImportCommand(m.Stdin, m.Stdout, m.Stderr)
 	case "backup":
-		cmd = NewBackupCommand(m.Stdin, m.Stdout, m.Stderr)
+		m.Cmd = NewBackupCommand(m.Stdin, m.Stdout, m.Stderr)
 	case "restore":
-		cmd = NewRestoreCommand(m.Stdin, m.Stdout, m.Stderr)
+		m.Cmd = NewRestoreCommand(m.Stdin, m.Stdout, m.Stderr)
+	case "bench":
+		m.Cmd = NewBenchCommand(m.Stdin, m.Stdout, m.Stderr)
 	default:
 		return ErrUnknownCommand
 	}
 
 	// Parse command's flags.
-	if err := cmd.ParseFlags(m.Args); err == ErrUsage {
-		fmt.Fprintln(m.Stderr, cmd.Usage())
-		return ErrQuit
+	if err := m.Cmd.ParseFlags(args); err == flag.ErrHelp {
+		fmt.Fprintln(m.Stderr, m.Cmd.Usage())
+		fmt.Fprintln(m.Stderr, "")
+		return err
 	} else if err != nil {
 		return err
 	}
 
-	// Execute the command.
-	if err := cmd.Run(); err != nil {
-		return err
-	}
-	return nil
-}
-
-// ParseFlags parses command line flags from args.
-func (m *Main) ParseFlags(args []string) error {
-	if len(args) == 0 {
-		return nil
-	}
-
-	m.Command = args[0]
-	m.Args = args[1:]
 	return nil
 }
 
@@ -139,7 +153,7 @@ func NewConfigCommand(stdin io.Reader, stdout, stderr io.Writer) *ConfigCommand
 // ParseFlags parses command line flags from args.
 func (cmd *ConfigCommand) ParseFlags(args []string) error {
 	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
-	fs.SetOutput(cmd.Stderr)
+	fs.SetOutput(ioutil.Discard)
 	if err := fs.Parse(args); err != nil {
 		return err
 	}
@@ -203,7 +217,7 @@ func NewImportCommand(stdin io.Reader, stdout, stderr io.Writer) *ImportCommand
 // ParseFlags parses command line flags from args.
 func (cmd *ImportCommand) ParseFlags(args []string) error {
 	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
-	fs.SetOutput(cmd.Stderr)
+	fs.SetOutput(ioutil.Discard)
 	fs.StringVar(&cmd.Host, "host", "localhost:15000", "host:port")
 	fs.StringVar(&cmd.Database, "d", "", "database")
 	fs.StringVar(&cmd.Frame, "f", "", "frame")
@@ -359,7 +373,7 @@ func NewBackupCommand(stdin io.Reader, stdout, stderr io.Writer) *BackupCommand
 // ParseFlags parses command line flags from args.
 func (cmd *BackupCommand) ParseFlags(args []string) error {
 	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
-	fs.SetOutput(cmd.Stderr)
+	fs.SetOutput(ioutil.Discard)
 	fs.StringVar(&cmd.Host, "host", "localhost:15000", "host:port")
 	fs.StringVar(&cmd.Database, "d", "", "database")
 	fs.StringVar(&cmd.Frame, "f", "", "frame")
@@ -445,7 +459,7 @@ func NewRestoreCommand(stdin io.Reader, stdout, stderr io.Writer) *RestoreComman
 // ParseFlags parses command line flags from args.
 func (cmd *RestoreCommand) ParseFlags(args []string) error {
 	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
-	fs.SetOutput(cmd.Stderr)
+	fs.SetOutput(ioutil.Discard)
 	fs.StringVar(&cmd.Host, "host", "localhost:15000", "host:port")
 	fs.StringVar(&cmd.Database, "d", "", "database")
 	fs.StringVar(&cmd.Frame, "f", "", "frame")
@@ -500,3 +514,131 @@ func (cmd *RestoreCommand) Run() error {
 
 	return nil
 }
+
+// BenchCommand represents a command for benchmarking database operations.
+type BenchCommand struct {
+	// Destination host and port.
+	Host string
+
+	// Name of the database & frame to execute against.
+	Database string
+	Frame    string
+
+	// Type of operation and number to execute.
+	Op string
+	N  int
+
+	// Standard input/output
+	Stdin  io.Reader
+	Stdout io.Writer
+	Stderr io.Writer
+}
+
+// NewBenchCommand returns a new instance of BenchCommand.
+func NewBenchCommand(stdin io.Reader, stdout, stderr io.Writer) *BenchCommand {
+	return &BenchCommand{
+		Stdin:  stdin,
+		Stdout: stdout,
+		Stderr: stderr,
+	}
+}
+
+// ParseFlags parses command line flags from args.
+func (cmd *BenchCommand) ParseFlags(args []string) error {
+	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
+	fs.SetOutput(ioutil.Discard)
+	fs.StringVar(&cmd.Host, "host", "localhost:15000", "host:port")
+	fs.StringVar(&cmd.Database, "d", "", "database")
+	fs.StringVar(&cmd.Frame, "f", "", "frame")
+	fs.StringVar(&cmd.Op, "op", "", "operation")
+	fs.IntVar(&cmd.N, "n", 0, "op count")
+
+	if err := fs.Parse(args); err != nil {
+		return err
+	}
+	return nil
+}
+
+// Usage returns the usage message to be printed.
+func (cmd *BenchCommand) Usage() string {
+	return strings.TrimSpace(`
+usage: pilosactl bench [args]
+
+Executes a benchmark for a given operation against the database.
+
+The following flags are allowed:
+
+	-host HOSTPORT
+		hostname and port of running pilosa server
+
+	-d DATABASE
+		database to execute operation against
+
+	-f FRAME
+		frame to execute operation against
+
+	-op OP
+		name of operation to execute
+
+	-n COUNT
+		number of iterations to execute
+
+The following operations are available:
+
+	set-bit
+		Sets a single random bit on the frame
+
+`)
+}
+
+// Run executes the main program execution.
+func (cmd *BenchCommand) Run() error {
+	// Create a client to the server.
+	client, err := pilosa.NewClient(cmd.Host)
+	if err != nil {
+		return err
+	}
+
+	switch cmd.Op {
+	case "set-bit":
+		return cmd.runSetBit(client)
+	case "":
+		return errors.New("op required")
+	default:
+		return fmt.Errorf("unknown bench op: %q", cmd.Op)
+	}
+}
+
+// runSetBit executes a benchmark of random SetBit() operations.
+func (cmd *BenchCommand) runSetBit(client *pilosa.Client) error {
+	if cmd.N == 0 {
+		return errors.New("operation count required")
+	} else if cmd.Database == "" {
+		return pilosa.ErrDatabaseRequired
+	} else if cmd.Frame == "" {
+		return pilosa.ErrFrameRequired
+	}
+
+	const maxBitmapID = 1000
+	const maxProfileID = 100000
+
+	startTime := time.Now()
+
+	// Execute operation continuously.
+	for i := 0; i < cmd.N; i++ {
+		bitmapID := rand.Intn(maxBitmapID)
+		profileID := rand.Intn(maxProfileID)
+
+		q := fmt.Sprintf(`SetBit(id=%d, frame="%s", profileID=%d)`, bitmapID, cmd.Frame, profileID)
+
+		if _, err := client.ExecuteQuery(cmd.Database, q); err != nil {
+			return err
+		}
+	}
+
+	// Print results.
+	elapsed := time.Since(startTime)
+	fmt.Fprintf(cmd.Stdout, "Executed %d operations in %s (%0.3f op/sec)\n", cmd.N, elapsed, float64(cmd.N)/elapsed.Seconds())
+
+	return nil
+}

db.go

@@ -1,6 +1,7 @@
 package pilosa
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -153,6 +154,10 @@ func (db *DB) CreateFrameIfNotExists(name string) (*Frame, error) {
 }
 
 func (db *DB) createFrameIfNotExists(name string) (*Frame, error) {
+	if name == "" {
+		return nil, errors.New("frame name required")
+	}
+
 	// Find frame in cache first.
 	if f := db.frames[name]; f != nil {
 		return f, nil

fragment.go

@@ -41,7 +41,7 @@ const (
 	DefaultCacheFlushInterval = 1 * time.Minute
 
 	// DefaultFragmentMaxOpN is the default value for Fragment.MaxOpN.
-	DefaultFragmentMaxOpN = 10000
+	DefaultFragmentMaxOpN = 1000
 )
 
 // Fragment represents the intersection of a frame and slice in a database.
@@ -256,12 +256,12 @@ func (f *Fragment) close() error {
 
 	// Flush cache if closing gracefully.
 	if err := f.flushCache(); err != nil {
-		f.logger().Printf("error flushing cache on close: err=%s, path=%s", err, f.path)
+		f.logger().Printf("fragment: error flushing cache on close: err=%s, path=%s", err, f.path)
 	}
 
 	// Close underlying storage.
 	if err := f.closeStorage(); err != nil {
-		f.logger().Printf("error closing storage: err=%s, path=%s", err, f.path)
+		f.logger().Printf("fragment: error closing storage: err=%s, path=%s", err, f.path)
 	}
 
 	return nil
@@ -634,6 +634,9 @@ func (f *Fragment) Snapshot() error {
 }
 
 func (f *Fragment) snapshot() error {
+	logger := f.logger()
+	logger.Printf("fragment: snapshotting %s/%s/%d", f.db, f.frame, f.slice)
+
 	// Create a temporary file to snapshot to.
 	snapshotPath := f.path + SnapshotExt
 	file, err := os.Create(snapshotPath)

handler.go

@@ -9,6 +9,7 @@ import (
 	"io/ioutil"
 	"log"
 	"net/http"
+	"net/http/pprof"
 	"os"
 	"strconv"
 	"strings"
@@ -49,8 +50,23 @@ func NewHandler() *Handler {
 
 // ServeHTTP handles an HTTP request.
 func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	t := time.Now()
+	// Handle pprof requests separately.
+	if strings.HasPrefix(r.URL.Path, "/debug/pprof") {
+		switch r.URL.Path {
+		case "/debug/pprof/cmdline":
+			pprof.Cmdline(w, r)
+		case "/debug/pprof/profile":
+			pprof.Profile(w, r)
+		case "/debug/pprof/symbol":
+			pprof.Symbol(w, r)
+		default:
+			pprof.Index(w, r)
+		}
+		return
+	}
 
+	// Route API calls to appropriate handler functions.
+	t := time.Now()
 	switch r.URL.Path {
 	case "/schema":
 		switch r.Method {

index.go

@@ -1,6 +1,7 @@
 package pilosa
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -119,6 +120,10 @@ func (i *Index) CreateDBIfNotExists(name string) (*DB, error) {
 }
 
 func (i *Index) createDBIfNotExists(name string) (*DB, error) {
+	if name == "" {
+		return nil, errors.New("database name required")
+	}
+
 	// Return database if it exists.
 	if db := i.db(name); db != nil {
 		return db, nil

pilosa.go

@@ -21,6 +21,9 @@ var (
 
 	// ErrFragmentNotFound is returned when a fragment does not exist.
 	ErrFragmentNotFound = errors.New("fragment not found")
+
+	// ErrQueryRequired is returned when no query is specified.
+	ErrQueryRequired = errors.New("query required")
 )
 
 // Version represents the current running version of Pilosa.