Merge pull request #146 from benbjohnson/bench-batch

Add --batch-size to 'bolt bench'

2 files changed, 41 insertions, 12 deletions

diff --git a/cmd/bolt/bench.go b/cmd/bolt/bench.go
index 72144b8..10a5599 100644
--- a/cmd/bolt/bench.go
+++ b/cmd/bolt/bench.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
+	"math/rand"
 	"os"
 	"runtime"
 	"runtime/pprof"
@@ -22,6 +23,13 @@ var benchBucketName = []byte("bench")
 func Bench(options *BenchOptions) {
 	var results BenchResults
 
+	// Validate options.
+	if options.BatchSize == 0 {
+		options.BatchSize = options.Iterations
+	} else if options.Iterations%options.BatchSize != 0 {
+		fatal("number of iterations must be divisible by the batch size")
+	}
+
 	// Find temporary location.
 	path := tempfile()
 	defer os.Remove(path)
@@ -78,6 +86,8 @@ func benchWrite(db *bolt.DB, options *BenchOptions, results *BenchResults) error
 	switch options.WriteMode {
 	case "seq":
 		err = benchWriteSequential(db, options, results)
+	case "rnd":
+		err = benchWriteRandom(db, options, results)
 	default:
 		return fmt.Errorf("invalid write mode: %s", options.WriteMode)
 	}
@@ -88,22 +98,38 @@ func benchWrite(db *bolt.DB, options *BenchOptions, results *BenchResults) error
 }
 
 func benchWriteSequential(db *bolt.DB, options *BenchOptions, results *BenchResults) error {
+	var i = uint32(0)
+	return benchWriteWithSource(db, options, results, func() uint32 { i++; return i })
+}
+
+func benchWriteRandom(db *bolt.DB, options *BenchOptions, results *BenchResults) error {
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	return benchWriteWithSource(db, options, results, func() uint32 { return r.Uint32() })
+}
+
+func benchWriteWithSource(db *bolt.DB, options *BenchOptions, results *BenchResults, keySource func() uint32) error {
 	results.WriteOps = options.Iterations
 
-	return db.Update(func(tx *bolt.Tx) error {
-		b, _ := tx.CreateBucketIfNotExists(benchBucketName)
+	for i := 0; i < options.Iterations; i += options.BatchSize {
+		err := db.Update(func(tx *bolt.Tx) error {
+			b, _ := tx.CreateBucketIfNotExists(benchBucketName)
 
-		for i := 0; i < options.Iterations; i++ {
-			var key = make([]byte, options.KeySize)
-			var value = make([]byte, options.ValueSize)
-			binary.BigEndian.PutUint32(key, uint32(i))
-			if err := b.Put(key, value); err != nil {
-				return err
+			for j := 0; j < options.BatchSize; j++ {
+				var key = make([]byte, options.KeySize)
+				var value = make([]byte, options.ValueSize)
+				binary.BigEndian.PutUint32(key, keySource())
+				if err := b.Put(key, value); err != nil {
+					return err
+				}
 			}
-		}
 
-		return nil
-	})
+			return nil
+		})
+		if err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
 // Reads from the database.
@@ -137,7 +163,7 @@ func benchReadSequential(db *bolt.DB, options *BenchOptions, results *BenchResul
 				count++
 			}
 
-			if count != options.Iterations {
+			if options.WriteMode == "seq" && count != options.Iterations {
 				return fmt.Errorf("read seq: iter mismatch: expected %d, got %d", options.Iterations, count)
 			}
 
@@ -213,6 +239,7 @@ type BenchOptions struct {
 	WriteMode    string
 	ReadMode     string
 	Iterations   int
+	BatchSize    int
 	KeySize      int
 	ValueSize    int
 	CPUProfile   string
diff --git a/cmd/bolt/main.go b/cmd/bolt/main.go
index 719bf00..cb2d47a 100644
--- a/cmd/bolt/main.go
+++ b/cmd/bolt/main.go
@@ -98,6 +98,7 @@ func NewApp() *cli.App {
 				&cli.StringFlag{Name: "write-mode", Value: "seq", Usage: "Write mode"},
 				&cli.StringFlag{Name: "read-mode", Value: "seq", Usage: "Read mode"},
 				&cli.IntFlag{Name: "count", Value: 1000, Usage: "Item count"},
+				&cli.IntFlag{Name: "batch-size", Usage: "Write batch size"},
 				&cli.IntFlag{Name: "key-size", Value: 8, Usage: "Key size"},
 				&cli.IntFlag{Name: "value-size", Value: 32, Usage: "Value size"},
 				&cli.StringFlag{Name: "cpuprofile", Usage: "CPU profile output path"},
@@ -110,6 +111,7 @@ func NewApp() *cli.App {
 					WriteMode:    c.String("write-mode"),
 					ReadMode:     c.String("read-mode"),
 					Iterations:   c.Int("count"),
+					BatchSize:    c.Int("batch-size"),
 					KeySize:      c.Int("key-size"),
 					ValueSize:    c.Int("value-size"),
 					CPUProfile:   c.String("cpuprofile"),