From cabb44e01f70840e5a7057a1d778f5da968ac77e Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Tue, 29 Apr 2014 12:27:38 -0600
Subject: Add --batch-size to 'bolt bench'.

This commit adds a --batch-size CLI argument to the 'bolt bench' tool.
This argument will insert into Bolt in smaller batches which is a more
typical use case.

/cc @snormore
---
 cmd/bolt/bench.go | 42 +++++++++++++++++++++++++++++++-----------
 1 file changed, 31 insertions(+), 11 deletions(-)

(limited to 'cmd/bolt/bench.go')

diff --git a/cmd/bolt/bench.go b/cmd/bolt/bench.go
index 72144b8..193687b 100644
--- a/cmd/bolt/bench.go
+++ b/cmd/bolt/bench.go
@@ -88,22 +88,41 @@ func benchWrite(db *bolt.DB, options *BenchOptions, results *BenchResults) error
 }
 
 func benchWriteSequential(db *bolt.DB, options *BenchOptions, results *BenchResults) error {
-	results.WriteOps = options.Iterations
+	// Default batch size to iteration count, if not specified.
+	var batchSize, iterations = options.BatchSize, options.Iterations
+	if batchSize == 0 {
+		batchSize = iterations
+	}
+
+	// Insert in batches.
+	var count int
+	for i := 0; i < (iterations/batchSize)+1; i++ {
+		err := db.Update(func(tx *bolt.Tx) error {
+			b, _ := tx.CreateBucketIfNotExists(benchBucketName)
+
+			for j := 0; j < batchSize && count < iterations; j++ {
+				var key = make([]byte, options.KeySize)
+				var value = make([]byte, options.ValueSize)
+				binary.BigEndian.PutUint32(key, uint32(count))
 
-	return db.Update(func(tx *bolt.Tx) error {
-		b, _ := tx.CreateBucketIfNotExists(benchBucketName)
+				if err := b.Put(key, value); err != nil {
+					return err
+				}
 
-		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
+				count++
 			}
+
+			return nil
+		})
+		if err != nil {
+			return err
 		}
+	}
 
-		return nil
-	})
+	// Update the write op count.
+	results.WriteOps = count
+
+	return nil
 }
 
 // Reads from the database.
@@ -213,6 +232,7 @@ type BenchOptions struct {
 	WriteMode    string
 	ReadMode     string
 	Iterations   int
+	BatchSize    int
 	KeySize      int
 	ValueSize    int
 	CPUProfile   string
-- 
cgit v1.2.3