From 7e524ed836c3eddf67944970f7e33e0de671324f Mon Sep 17 00:00:00 2001
From: jaffee <jaffee@umbel.com>
Date: Thu, 10 Nov 2016 14:22:14 -0600
Subject: [PATCH] setup and support BenchmarkCmd interface

---
 bench/bench.go        | 49 +++++++++++++++++++++++++++++++++++++++++++
 cmd/pilosactl/main.go | 43 ++++++++++++++++++++++++-------------
 2 files changed, 78 insertions(+), 14 deletions(-)

diff --git a/bench/bench.go b/bench/bench.go
index cfb1601f1..8bbb78eaa 100644
--- a/bench/bench.go
+++ b/bench/bench.go
@@ -6,6 +6,9 @@ import (
 
 	"strconv"
 
+	"flag"
+	"io/ioutil"
+
 	"github.com/umbel/pilosa"
 )
 
@@ -27,6 +30,12 @@ type Benchmarker interface {
 	Run(agentNum int) map[string]interface{}
 }
 
+type BenchmarkCmd interface {
+	Benchmarker
+	ConsumeFlags(args []string) ([]string, error)
+	Usage() string
+}
+
 // SetBitBenchmark sets a bunch of bits in pilosa, and can be configured in a
 // number of ways.
 type SetBitBenchmark struct {
@@ -44,6 +53,46 @@ type SetBitBenchmark struct {
 	DB string
 }
 
+func (b *SetBitBenchmark) Usage() string {
+	return `
+SetBitBenchmark sets a bunch of bits.
+
+Usage: SetBitBenchmark [arguments]
+
+The following arguments are available:
+
+	-BaseBitmapID int
+		bit num to start from
+
+	-BaseProfileID int
+		profile id num to start from
+
+	-Iterations int
+		number of bits to set
+
+	-NumProfiles int
+		number of profiles to loop through
+
+	-DB string
+		pilosa db to use
+`[1:]
+}
+
+func (b *SetBitBenchmark) ConsumeFlags(args []string) ([]string, error) {
+	fs := flag.NewFlagSet("SetBitBenchmark", flag.ContinueOnError)
+	fs.SetOutput(ioutil.Discard)
+	fs.IntVar(&b.BaseBitmapID, "BaseBitmapID", 0, "bits being set will all be greater than BaseBitmapID")
+	fs.IntVar(&b.BaseProfileID, "BaseProfileID", 0, "profile ids used will all be greater than BaseProfileID")
+	fs.IntVar(&b.Iterations, "Iterations", 100, "Iterations is the number of bits that will be set by this Benchmark")
+	fs.IntVar(&b.NumProfiles, "NumProfiles", 100, "number of profiles to iterate through")
+	fs.StringVar(&b.DB, "DB", "benchdb", "pilosa DB to use")
+
+	if err := fs.Parse(args); err != nil {
+		return nil, err
+	}
+	return fs.Args(), nil
+}
+
 // Init connects to pilosa and sets the client on b.
 func (b *SetBitBenchmark) Init(hosts []string) (err error) {
 	b.cli, err = pilosa.NewClient(hosts[0])
diff --git a/cmd/pilosactl/main.go b/cmd/pilosactl/main.go
index 5d3c80c03..21f5c6fd2 100644
--- a/cmd/pilosactl/main.go
+++ b/cmd/pilosactl/main.go
@@ -1284,24 +1284,39 @@ func NewBagentCommand(stdin io.Reader, stdout, stderr io.Writer) *BagentCommand
 
 // ParseFlags parses command line flags from args.
 func (cmd *BagentCommand) ParseFlags(args []string) error {
-	fs := flag.NewFlagSet("pilosa-bench-agent", flag.ContinueOnError)
+	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
 	fs.SetOutput(ioutil.Discard)
 
-	var benchmarks string
 	var pilosaHosts string
-
-	fs.StringVar(&benchmarks, "benchmarks", "", "Comma separated list of benchmarks to run")
 	fs.StringVar(&pilosaHosts, "hosts", "localhost:15000", "Comma separated list of host:port")
 	fs.IntVar(&cmd.AgentNum, "agentNum", 0, "An integer differentiating this agent from other in the fleet.")
 
 	if err := fs.Parse(args); err != nil {
 		return err
 	}
-	for _, bmName := range strings.Split(benchmarks, ",") {
-		if bm, ok := bench.Benchmarks[bmName]; ok {
-			cmd.Benchmarks = append(cmd.Benchmarks, bm)
-		} else {
-			return fmt.Errorf("%v is not a configured benchmark", bmName)
+	remArgs := fs.Args()
+	if len(remArgs) == 0 {
+		return flag.ErrHelp
+	}
+	for len(remArgs) > 0 {
+		var bm bench.BenchmarkCmd
+		var err error
+		switch remArgs[0] {
+		case "-help", "-h":
+			return flag.ErrHelp
+		case "SetBitBenchmark":
+			bm = &bench.SetBitBenchmark{}
+		default:
+			return fmt.Errorf("Unknown benchmark cmd: %v", remArgs[0])
+		}
+		remArgs, err = bm.ConsumeFlags(remArgs[1:])
+		cmd.Benchmarks = append(cmd.Benchmarks, bm)
+		if err != nil {
+			if err == flag.ErrHelp {
+				fmt.Fprintln(cmd.Stderr, bm.Usage())
+				return fmt.Errorf("")
+			}
+			return fmt.Errorf("BagentCommand.ParseFlags: %v", err)
 		}
 	}
 	cmd.Hosts = strings.Split(pilosaHosts, ",")
@@ -1312,22 +1327,22 @@ func (cmd *BagentCommand) ParseFlags(args []string) error {
 // Usage returns the usage message to be printed.
 func (cmd *BagentCommand) Usage() string {
 	return strings.TrimSpace(`
-pilosa-benchmark-agent is a tool for running a set of benchmarks against a pilosa cluster.
+pilosactl bagent is a tool for running benchmarks against a pilosa cluster.
 
 Usage:
 
-	pilosa-benchmark-agent [arguments]
+pilosactl bagent [options] <subcommand [options]>...
 
 The following arguments are available:
 
-	-benchmarks
-		Comma separated list of benchmarks.
-
 	-hosts
 		Comma separated list of host:port describing all hosts in the cluster.
 
 	-agentNum N
 		An integer differentiating this agent from others in the fleet.
+
+	subcommands:
+		SetBitBenchmark
 `)
 }