add benchmarks using Benchmark framework

1 files changed, 42 insertions, 19 deletions

diff --git a/cmd/bolt/generate.go b/cmd/bolt/generate.go
index ecd391e..15edb27 100644
--- a/cmd/bolt/generate.go
+++ b/cmd/bolt/generate.go
@@ -1,32 +1,55 @@
 package main
 
 import (
-	"bufio"
 	"fmt"
-	"os"
-	"strings"
+
+	"github.com/boltdb/bolt"
 )
 
 // Generate data for benchmarks.
-func Generate(numEvents int, destPath string) {
-	f, err := os.Create(destPath)
+func Generate(destPath string, numBuckets, numItems int) {
+
+	// Open the database.
+	db, err := bolt.Open(destPath, 0600)
 	if err != nil {
-		fatal(err)
-	}
-	defer func() {
-		if err := f.Close(); err != nil {
-			fatal(err)
-		}
-	}()
-	w := bufio.NewWriter(f)
-
-	for i := 0; i < numEvents; i++ {
-		if _, err := w.Write([]byte(fmt.Sprintf("key%d:%s\n", i, strings.Repeat("0", 64)))); err != nil {
-			fatal(err)
-		}
+		fatalf("open db:", err)
+		return
 	}
+	defer db.Close()
+
+	for bucketIndex := 0; bucketIndex < numBuckets; bucketIndex++ {
+		bucketName := fmt.Sprintf("bucket%03d", bucketIndex)
+
+		err = db.Update(func(tx *bolt.Tx) error {
 
-	if err = w.Flush(); err != nil {
+			// Create the bucket if it doesn't exist.
+			if err := tx.CreateBucketIfNotExists([]byte(bucketName)); err != nil {
+				fatalf("create bucket: %s", err)
+				return nil
+			}
+
+			// Find bucket.
+			b := tx.Bucket([]byte(bucketName))
+			if b == nil {
+				fatalf("bucket not found: %s", bucketName)
+				return nil
+			}
+
+			for i := 0; i < numItems; i++ {
+				key := fmt.Sprintf("key%03d", i)
+				value := fmt.Sprintf("value%03d", i)
+
+				// Set value for a given key.
+				if err := b.Put([]byte(key), []byte(value)); err != nil {
+					return err
+				}
+			}
+
+			return nil
+		})
+	}
+	if err != nil {
 		fatal(err)
+		return
 	}
 }