add -stats and -batch-size support to bench cmd

2 files changed, 63 insertions, 12 deletions

diff --git a/cmd/bolt/bench.go b/cmd/bolt/bench.go
index 72144b8..e80f51a 100644
--- a/cmd/bolt/bench.go
+++ b/cmd/bolt/bench.go
@@ -5,7 +5,9 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
+	"math/rand"
 	"os"
+	"reflect"
 	"runtime"
 	"runtime/pprof"
 	"time"
@@ -68,6 +70,19 @@ func Bench(options *BenchOptions) {
 	fmt.Printf("# Write\t%v\t(%v/op)\t(%v op/sec)\n", results.WriteDuration, results.WriteOpDuration(), results.WriteOpsPerSecond())
 	fmt.Printf("# Read\t%v\t(%v/op)\t(%v op/sec)\n", results.ReadDuration, results.ReadOpDuration(), results.ReadOpsPerSecond())
 	fmt.Println("")
+
+	if options.Stats {
+		fmt.Println("Transaction Stats")
+		printStruct(db.Stats().TxStats)
+		fmt.Println("")
+		db.View(func(tx *bolt.Tx) error {
+			b := tx.Bucket(benchBucketName)
+			fmt.Println("Storage Stats")
+			printStruct(b.Stats())
+			fmt.Println("")
+			return nil
+		})
+	}
 }
 
 // Writes to the database.
@@ -78,6 +93,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 +105,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(42))
+	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 +170,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)
 			}
 
@@ -215,6 +248,8 @@ type BenchOptions struct {
 	Iterations   int
 	KeySize      int
 	ValueSize    int
+	BatchSize    int
+	Stats        bool
 	CPUProfile   string
 	MemProfile   string
 	BlockProfile string
@@ -269,3 +304,11 @@ func tempfile() string {
 	os.Remove(f.Name())
 	return f.Name()
 }
+
+func printStruct(s interface{}) {
+	v := reflect.ValueOf(s)
+	t := reflect.TypeOf(s)
+	for i := 0; i < v.NumField(); i++ {
+		fmt.Printf("  %s: %v\n", t.Field(i).Name, v.Field(i).Interface())
+	}
+}
diff --git a/cmd/bolt/main.go b/cmd/bolt/main.go
index 719bf00..cab4e0e 100644
--- a/cmd/bolt/main.go
+++ b/cmd/bolt/main.go
@@ -100,11 +100,17 @@ func NewApp() *cli.App {
 				&cli.IntFlag{Name: "count", Value: 1000, Usage: "Item count"},
 				&cli.IntFlag{Name: "key-size", Value: 8, Usage: "Key size"},
 				&cli.IntFlag{Name: "value-size", Value: 32, Usage: "Value size"},
+				&cli.IntFlag{Name: "batch-size", Value: 0, Usage: "Write batch size"},
 				&cli.StringFlag{Name: "cpuprofile", Usage: "CPU profile output path"},
 				&cli.StringFlag{Name: "memprofile", Usage: "Memory profile output path"},
 				&cli.StringFlag{Name: "blockprofile", Usage: "Block profile output path"},
+				&cli.BoolFlag{Name: "stats", Usage: "Output storage and transaction stats"},
 			},
 			Action: func(c *cli.Context) {
+				bs := c.Int("batch-size")
+				if bs == 0 {
+					bs = c.Int("count")
+				}
 				Bench(&BenchOptions{
 					ProfileMode:  c.String("profile-mode"),
 					WriteMode:    c.String("write-mode"),
@@ -112,9 +118,11 @@ func NewApp() *cli.App {
 					Iterations:   c.Int("count"),
 					KeySize:      c.Int("key-size"),
 					ValueSize:    c.Int("value-size"),
+					BatchSize:    bs,
 					CPUProfile:   c.String("cpuprofile"),
 					MemProfile:   c.String("memprofile"),
 					BlockProfile: c.String("blockprofile"),
+					Stats:        c.Bool("stats"),
 				})
 			},
 		}}