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"),