1 files changed, 25 insertions, 18 deletions
diff --git a/cmd/bolt/bench.go b/cmd/bolt/bench.go
index 193687b..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,28 +98,29 @@ func benchWrite(db *bolt.DB, options *BenchOptions, results *BenchResults) error
 }
 
 func benchWriteSequential(db *bolt.DB, options *BenchOptions, results *BenchResults) error {
-	// Default batch size to iteration count, if not specified.
-	var batchSize, iterations = options.BatchSize, options.Iterations
-	if batchSize == 0 {
-		batchSize = iterations
-	}
+	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() })
+}
 
-	// Insert in batches.
-	var count int
-	for i := 0; i < (iterations/batchSize)+1; i++ {
+func benchWriteWithSource(db *bolt.DB, options *BenchOptions, results *BenchResults, keySource func() uint32) error {
+	results.WriteOps = options.Iterations
+
+	for i := 0; i < options.Iterations; i += options.BatchSize {
 		err := db.Update(func(tx *bolt.Tx) error {
 			b, _ := tx.CreateBucketIfNotExists(benchBucketName)
 
-			for j := 0; j < batchSize && count < iterations; j++ {
+			for j := 0; j < options.BatchSize; j++ {
 				var key = make([]byte, options.KeySize)
 				var value = make([]byte, options.ValueSize)
-				binary.BigEndian.PutUint32(key, uint32(count))
-
+				binary.BigEndian.PutUint32(key, keySource())
 				if err := b.Put(key, value); err != nil {
 					return err
 				}
-
-				count++
 			}
 
 			return nil
@@ -118,10 +129,6 @@ func benchWriteSequential(db *bolt.DB, options *BenchOptions, results *BenchResu
 			return err
 		}
 	}
-
-	// Update the write op count.
-	results.WriteOps = count
-
 	return nil
 }
 
@@ -156,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)
 			}