move bench to subcommand

cmd/bench.go

@@ -0,0 +1,43 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"os"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+
+	"github.com/pilosa/pilosa/ctl"
+)
+
+var bencher = ctl.NewBenchCommand(os.Stdin, os.Stdout, os.Stderr)
+
+var benchCmd = &cobra.Command{
+	Use:   "bench",
+	Short: "bench - benchmark operations",
+	Long: `
+Executes a benchmark for a given operation against the database.
+`,
+	Run: func(cmd *cobra.Command, args []string) {
+		if err := bencher.Run(context.Background()); err != nil {
+			fmt.Println(err)
+		}
+	},
+}
+
+func init() {
+	benchCmd.Flags().StringVarP(&bencher.Host, "host", "", "localhost:15000", "host:port of Pilosa.")
+	benchCmd.Flags().StringVarP(&bencher.Database, "database", "d", "", "Pilosa database to benchmark.")
+	benchCmd.Flags().StringVarP(&bencher.Frame, "frame", "f", "", "Frame to benchmark.")
+	benchCmd.Flags().StringVarP(&bencher.Op, "operation", "o", "set-bit", "Operation to perform: choose from [set-bit]")
+	benchCmd.Flags().IntVarP(&bencher.N, "num", "n", 0, "Number of operations to perform.")
+
+	err := viper.BindPFlags(benchCmd.Flags())
+	if err != nil {
+		log.Fatalf("Error binding bench flags: %v", err)
+	}
+
+	RootCmd.AddCommand(benchCmd)
+}

cmd/pilosactl/main.go

@@ -85,7 +85,6 @@ Usage:
 
 The commands are:
 
-	bench      benchmarks operations
 
 Use the "-h" flag with any command for more information.
 `)
@@ -107,8 +106,6 @@ func (m *Main) ParseFlags(args []string) error {
 		fmt.Fprintln(m.Stderr, m.Usage())
 		fmt.Fprintln(m.Stderr, "")
 		return flag.ErrHelp
-	case "bench":
-		m.Cmd = NewBenchCommand(m.Stdin, m.Stdout, m.Stderr)
 	default:
 		return ErrUnknownCommand
 	}

ctl/bench.go

@@ -0,0 +1,143 @@
+package ctl
+
+import (
+	"context"
+	"errors"
+	"flag"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"math/rand"
+	"strings"
+	"time"
+
+	"github.com/pilosa/pilosa"
+)
+
+// 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(ctx context.Context) 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(ctx, 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(ctx context.Context, 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(ctx, cmd.Database, q, true); 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
+}