Merge pull request #196 from jaffee/195-fix-context

thread context through Benchmarks

bench/bench.go

@@ -1,11 +1,13 @@
 package bench
 
 import (
+	"context"
 	"fmt"
-	"golang.org/x/sync/errgroup"
 	"strconv"
 	"sync"
 	"time"
+
+	"golang.org/x/sync/errgroup"
 )
 
 // Benchmark is an interface to guide the creation of new pilosa benchmarks or
@@ -23,7 +25,7 @@ type Benchmark interface {
 	// for each agent, or have each agent set a different set of bits. The return
 	// value of Run is kept generic so that any relevant statistics or metrics
 	// that may be specific to the benchmark in question can be reported.
-	Run(agentNum int) map[string]interface{}
+	Run(ctx context.Context, agentNum int) map[string]interface{}
 }
 
 // Command extends Benchmark by adding methods for configuring via command line flags and returning usage information.
@@ -64,7 +66,7 @@ func (pb *parallelBenchmark) Init(hosts []string, agentNum int) error {
 // Run runs the parallel benchmark and returns it's results in a nested map - the
 // top level keys are the indices of each benchmark in the list of benchmarks,
 // and the values are the results of each benchmark's Run method.
-func (pb *parallelBenchmark) Run(agentNum int) map[string]interface{} {
+func (pb *parallelBenchmark) Run(ctx context.Context, agentNum int) map[string]interface{} {
 	wg := sync.WaitGroup{}
 	results := make(map[string]interface{}, len(pb.benchmarkers))
 	resultsLock := sync.Mutex{}
@@ -72,7 +74,7 @@ func (pb *parallelBenchmark) Run(agentNum int) map[string]interface{} {
 		wg.Add(1)
 		go func(i int, b Benchmark) {
 			defer wg.Done()
-			ret := b.Run(agentNum)
+			ret := b.Run(ctx, agentNum)
 			resultsLock.Lock()
 			results[strconv.Itoa(i)] = ret
 			resultsLock.Unlock()
@@ -114,13 +116,13 @@ func (sb *serialBenchmark) Init(hosts []string, agentNum int) error {
 // Run runs the serial benchmark and returns it's results in a nested map - the
 // top level keys are the indices of each benchmark in the list of benchmarks,
 // and the values are the results of each benchmark's Run method.
-func (sb *serialBenchmark) Run(agentNum int) map[string]interface{} {
+func (sb *serialBenchmark) Run(ctx context.Context, agentNum int) map[string]interface{} {
 	results := make(map[string]interface{}, len(sb.benchmarkers))
 	runtimes := make(map[string]time.Duration)
 	total_start := time.Now()
 	for i, b := range sb.benchmarkers {
 		start := time.Now()
-		ret := b.Run(agentNum)
+		ret := b.Run(ctx, agentNum)
 		end := time.Now()
 		results[strconv.Itoa(i)] = ret
 		runtimes[strconv.Itoa(i)] = end.Sub(start)

bench/diagonal.go

@@ -61,7 +61,7 @@ func (b *DiagonalSetBits) ConsumeFlags(args []string) ([]string, error) {
 }
 
 // Run runs the DiagonalSetBits benchmark
-func (b *DiagonalSetBits) Run(agentNum int) map[string]interface{} {
+func (b *DiagonalSetBits) Run(ctx context.Context, agentNum int) map[string]interface{} {
 	results := make(map[string]interface{})
 	if b.cli == nil {
 		results["error"] = fmt.Errorf("No client set for DiagonalSetBits agent: %v", agentNum)
@@ -73,7 +73,11 @@ func (b *DiagonalSetBits) Run(agentNum int) map[string]interface{} {
 		iterID := agentizeNum(n, b.Iterations, agentNum)
 		query := fmt.Sprintf("SetBit(%d, 'frame.n', %d)", b.BaseBitmapID+iterID, b.BaseProfileID+iterID)
 		start = time.Now()
-		b.cli.ExecuteQuery(context.TODO(), b.DB, query, true)
+		_, err := b.cli.ExecuteQuery(ctx, b.DB, query, true)
+		if err != nil {
+			results["error"] = err
+			return results
+		}
 		s.Add(time.Now().Sub(start))
 	}
 	AddToResults(s, results)

bench/import.go

@@ -139,12 +139,13 @@ func (b *Import) Init(hosts []string, agentNum int) error {
 }
 
 // Run runs the Import benchmark
-func (b *Import) Run(agentNum int) map[string]interface{} {
+func (b *Import) Run(ctx context.Context, agentNum int) map[string]interface{} {
 	results := make(map[string]interface{})
 	results["numbits"] = b.numbits
 	results["db"] = b.Database
 	start := time.Now()
-	err := b.ImportCommand.Run(context.TODO())
+	err := b.ImportCommand.Run(ctx)
+
 	if err != nil {
 		results["error"] = err.Error()
 	}

bench/multidb.go

@@ -55,7 +55,7 @@ func (b *MultiDBSetBits) ConsumeFlags(args []string) ([]string, error) {
 }
 
 // Run runs the MultiDBSetBits benchmark
-func (b *MultiDBSetBits) Run(agentNum int) map[string]interface{} {
+func (b *MultiDBSetBits) Run(ctx context.Context, agentNum int) map[string]interface{} {
 	results := make(map[string]interface{})
 	if b.cli == nil {
 		results["error"] = fmt.Errorf("No client set for MultiDBSetBits agent: %v", agentNum)
@@ -66,7 +66,11 @@ func (b *MultiDBSetBits) Run(agentNum int) map[string]interface{} {
 	for n := 0; n < b.Iterations; n++ {
 		query := fmt.Sprintf("SetBit(%d, 'frame.n', %d)", b.BaseBitmapID+n, b.BaseProfileID+n)
 		start = time.Now()
-		b.cli.ExecuteQuery(context.TODO(), "multidb"+strconv.Itoa(agentNum), query, true)
+		_, err := b.cli.ExecuteQuery(ctx, "multidb"+strconv.Itoa(agentNum), query, true)
+		if err != nil {
+			results["error"] = err
+			return results
+		}
 		s.Add(time.Now().Sub(start))
 	}
 	AddToResults(s, results)

bench/random.go

@@ -77,7 +77,7 @@ func (b *RandomSetBits) ConsumeFlags(args []string) ([]string, error) {
 }
 
 // Run runs the RandomSetBits benchmark
-func (b *RandomSetBits) Run(agentNum int) map[string]interface{} {
+func (b *RandomSetBits) Run(ctx context.Context, agentNum int) map[string]interface{} {
 	src := rand.NewSource(b.Seed + int64(agentNum))
 	rng := rand.New(src)
 	results := make(map[string]interface{})
@@ -92,7 +92,7 @@ func (b *RandomSetBits) Run(agentNum int) map[string]interface{} {
 		profID := rng.Int63n(b.ProfileIDRange)
 		query := fmt.Sprintf("SetBit(%d, 'frame.n', %d)", b.BaseBitmapID+bitmapID, b.BaseProfileID+profID)
 		start = time.Now()
-		b.cli.ExecuteQuery(context.TODO(), b.DB, query, true)
+		b.cli.ExecuteQuery(ctx, b.DB, query, true)
 		s.Add(time.Now().Sub(start))
 	}
 	AddToResults(s, results)

bench/randquery.go

@@ -82,7 +82,7 @@ func (b *RandomQuery) ConsumeFlags(args []string) ([]string, error) {
 }
 
 // Run runs the RandomQuery benchmark
-func (b *RandomQuery) Run(agentNum int) map[string]interface{} {
+func (b *RandomQuery) Run(ctx context.Context, agentNum int) map[string]interface{} {
 	seed := b.Seed + int64(agentNum)
 	results := make(map[string]interface{})
 	if b.cli == nil {
@@ -95,7 +95,7 @@ func (b *RandomQuery) Run(agentNum int) map[string]interface{} {
 	for n := 0; n < b.Iterations; n++ {
 		call := qm.Random(b.MaxN, b.MaxDepth, b.MaxArgs, uint64(b.BaseBitmapID), uint64(b.BitmapIDRange))
 		start = time.Now()
-		b.cli.ExecuteQuery(context.TODO(), b.DBs[n%len(b.DBs)], call.String(), true)
+		b.cli.ExecuteQuery(ctx, b.DBs[n%len(b.DBs)], call.String(), true)
 		s.Add(time.Now().Sub(start))
 	}
 	AddToResults(s, results)

bench/sliceheight.go

@@ -89,7 +89,7 @@ func (b *SliceHeight) Init(hosts []string, agentNum int) error {
 }
 
 // Run runs the SliceHeight benchmark
-func (b *SliceHeight) Run(agentNum int) map[string]interface{} {
+func (b *SliceHeight) Run(ctx context.Context, agentNum int) map[string]interface{} {
 	results := make(map[string]interface{})
 
 	imp := NewImport(b.Stdin, b.Stdout, b.Stderr)
@@ -111,11 +111,11 @@ func (b *SliceHeight) Run(agentNum int) map[string]interface{} {
 		gendur := time.Now().Sub(genstart)
 		iresults["csvgen"] = gendur
 
-		iresults["import"] = imp.Run(agentNum)
+		iresults["import"] = imp.Run(ctx, agentNum)
 
 		qstart := time.Now()
 		q := &pql.TopN{Frame: b.Frame, N: 50}
-		_, err := imp.Client.ExecuteQuery(context.TODO(), b.Database, q.String(), true)
+		_, err := imp.Client.ExecuteQuery(ctx, b.Database, q.String(), true)
 		if err != nil {
 			iresults["query_error"] = err.Error()
 		} else {

cmd/pilosa/main_test.go

@@ -16,7 +16,6 @@ import (
 	"testing"
 	"testing/quick"
 
-	"context"
 	"github.com/BurntSushi/toml"
 	"github.com/pilosa/pilosa"
 	main "github.com/pilosa/pilosa/cmd/pilosa"

cmd/pilosactl/main.go

@@ -99,7 +99,8 @@ The commands are:
 	check      performs a consistency check of data files
 	bench      benchmarks operations
 	create     create pilosa clusters
-  bagent     run a benchmarking agent
+	bagent     run a benchmarking agent
+	bspawn     create a cluster and agents and run benchmarks based on config file
 
 Use the "-h" flag with any command for more information.
 `)
@@ -1222,7 +1223,7 @@ func (cmd *BagentCommand) Run(ctx context.Context) error {
 		return fmt.Errorf("in cmd.Run initialization: %v", err)
 	}
 
-	res := sbm.Run(cmd.AgentNum)
+	res := sbm.Run(ctx, cmd.AgentNum)
 	enc := json.NewEncoder(cmd.Stdout)
 	enc.SetIndent("", "  ")
 	err = enc.Encode(res)
@@ -1230,6 +1231,7 @@ func (cmd *BagentCommand) Run(ctx context.Context) error {
 		fmt.Fprintln(cmd.Stderr, err)
 	}
 	// fmt.Fprintln(cmd.Stdout, res)
+
 	return nil
 }