Add 'bolt bench'.

This commit adds a flexible benchmarking tool to the 'bolt' CLI. It allows
the user to separately specify the write mode and read mode (e.g. sequential
random, etc). It also allows the user to isolate profiling to either the
read or the writes.

Currently the bench tool only supports "seq" read and write modes. It also
does not support streaming of Bolt counters yet.

Fixes boltdb#95.

/cc @snormore

Makefile

@@ -24,7 +24,7 @@ cpuprofile: fmt
 # go get github.com/kisielk/errcheck
 errcheck:
 	@echo "=== errcheck ==="
-	@.go/bin/errcheck github.com/boltdb/bolt
+	@errcheck github.com/boltdb/bolt
 
 fmt:
 	@go fmt ./...
@@ -34,8 +34,7 @@ get:
 
 build: get
 	@mkdir -p bin
-	@go build -ldflags=$(GOLDFLAGS) -a -o bin/bolt-`git rev-parse --short HEAD` ./cmd/bolt
-	@echo "writing bin/bolt-`git rev-parse --short HEAD`"
+	@go build -ldflags=$(GOLDFLAGS) -a -o bin/bolt ./cmd/bolt
 
 test: fmt errcheck
 	@go get github.com/stretchr/testify/assert

cmd/bolt/bench.go

@@ -1,31 +1,271 @@
 package main
 
 import (
-	"testing"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"runtime"
+	"runtime/pprof"
+	"time"
 
 	"github.com/boltdb/bolt"
-	"github.com/boltdb/bolt/bench"
 )
 
-// Import converts an exported database dump into a new database.
-// readWriteMode: 'read' or 'write'
-// traversalPattern: 'sequentrial' or 'random'
-// parallelism: integer representing number of concurrent reads/writes
-func Bench(inputPath string, readWriteMode string, traversalPattern string, parallelism int) {
+// File handlers for the various profiles.
+var cpuprofile, memprofile, blockprofile *os.File
 
-	// Open the database.
-	db, err := bolt.Open(inputPath, 0600)
+var benchBucketName = []byte("bench")
+
+// Bench executes a customizable, synthetic benchmark against Bolt.
+func Bench(options *BenchOptions) {
+	var results BenchResults
+
+	// Find temporary location.
+	path := tempfile()
+	defer os.Remove(path)
+
+	// Create database.
+	db, err := bolt.Open(path, 0600)
 	if err != nil {
-		fatalf("error: %+v", err)
+		fatal(err)
 		return
 	}
 	defer db.Close()
 
-	b := bench.New(db, &bench.Config{
-		ReadWriteMode:    readWriteMode,
-		TraversalPattern: traversalPattern,
-		Parallelism:      parallelism,
+	// Start profiling for writes.
+	if options.ProfileMode == "rw" || options.ProfileMode == "w" {
+		benchStartProfiling(options)
+	}
+
+	// Write to the database.
+	if err := benchWrite(db, options, &results); err != nil {
+		fatal("bench: write: ", err)
+	}
+
+	// Stop profiling for writes only.
+	if options.ProfileMode == "w" {
+		benchStopProfiling()
+	}
+
+	// Start profiling for reads.
+	if options.ProfileMode == "r" {
+		benchStartProfiling(options)
+	}
+
+	// Read from the database.
+	if err := benchRead(db, options, &results); err != nil {
+		fatal("bench: read: ", err)
+	}
+
+	// Stop profiling for writes only.
+	if options.ProfileMode == "rw" || options.ProfileMode == "r" {
+		benchStopProfiling()
+	}
+
+	// Print results.
+	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("")
+}
+
+// Writes to the database.
+func benchWrite(db *bolt.DB, options *BenchOptions, results *BenchResults) error {
+	var err error
+	var t = time.Now()
+
+	switch options.WriteMode {
+	case "seq":
+		err = benchWriteSequential(db, options, results)
+	default:
+		return fmt.Errorf("invalid write mode: %s", options.WriteMode)
+	}
+
+	results.WriteDuration = time.Since(t)
+
+	return err
+}
+
+func benchWriteSequential(db *bolt.DB, options *BenchOptions, results *BenchResults) error {
+	results.WriteOps = options.Iterations
+
+	return 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
+			}
+		}
+
+		return nil
+	})
+}
+
+// Reads from the database.
+func benchRead(db *bolt.DB, options *BenchOptions, results *BenchResults) error {
+	var err error
+	var t = time.Now()
+
+	switch options.ReadMode {
+	case "seq":
+		err = benchReadSequential(db, options, results)
+	default:
+		return fmt.Errorf("invalid read mode: %s", options.ReadMode)
+	}
+
+	results.ReadDuration = time.Since(t)
+
+	return err
+}
+
+func benchReadSequential(db *bolt.DB, options *BenchOptions, results *BenchResults) error {
+	return db.View(func(tx *bolt.Tx) error {
+		var t = time.Now()
+
+		for {
+			c := tx.Bucket(benchBucketName).Cursor()
+			var count int
+			for k, v := c.First(); k != nil; k, v = c.Next() {
+				if v == nil {
+					return errors.New("invalid value")
+				}
+				count++
+			}
+
+			if count != options.Iterations {
+				return fmt.Errorf("read seq: iter mismatch: expected %d, got %d", options.Iterations, count)
+			}
+
+			results.ReadOps += count
+
+			// Make sure we do this for at least a second.
+			if time.Since(t) >= time.Second {
+				break
+			}
+		}
+
+		return nil
 	})
+}
+
+// Starts all profiles set on the options.
+func benchStartProfiling(options *BenchOptions) {
+	var err error
+
+	// Start CPU profiling.
+	if options.CPUProfile != "" {
+		cpuprofile, err = os.Create(options.CPUProfile)
+		if err != nil {
+			fatal("bench: could not create cpu profile %q: %v", options.CPUProfile, err)
+		}
+		pprof.StartCPUProfile(cpuprofile)
+	}
+
+	// Start memory profiling.
+	if options.MemProfile != "" {
+		memprofile, err = os.Create(options.MemProfile)
+		if err != nil {
+			fatal("bench: could not create memory profile %q: %v", options.MemProfile, err)
+		}
+		runtime.MemProfileRate = 4096
+	}
+
+	// Start fatal profiling.
+	if options.BlockProfile != "" {
+		blockprofile, err = os.Create(options.BlockProfile)
+		if err != nil {
+			fatal("bench: could not create block profile %q: %v", options.BlockProfile, err)
+		}
+		runtime.SetBlockProfileRate(1)
+	}
+}
+
+// Stops all profiles.
+func benchStopProfiling() {
+	if cpuprofile != nil {
+		pprof.StopCPUProfile()
+		cpuprofile.Close()
+		cpuprofile = nil
+	}
+
+	if memprofile != nil {
+		pprof.Lookup("heap").WriteTo(memprofile, 0)
+		memprofile.Close()
+		memprofile = nil
+	}
+
+	if blockprofile != nil {
+		pprof.Lookup("block").WriteTo(blockprofile, 0)
+		blockprofile.Close()
+		blockprofile = nil
+		runtime.SetBlockProfileRate(0)
+	}
+}
+
+// BenchOptions represents the set of options that can be passed to Bench().
+type BenchOptions struct {
+	ProfileMode  string
+	WriteMode    string
+	ReadMode     string
+	Iterations   int
+	KeySize      int
+	ValueSize    int
+	CPUProfile   string
+	MemProfile   string
+	BlockProfile string
+}
+
+// BenchResults represents the performance results of the benchmark.
+type BenchResults struct {
+	WriteOps      int
+	WriteDuration time.Duration
+	ReadOps       int
+	ReadDuration  time.Duration
+}
+
+// Returns the duration for a single write operation.
+func (r *BenchResults) WriteOpDuration() time.Duration {
+	if r.WriteOps == 0 {
+		return 0
+	}
+	return r.WriteDuration / time.Duration(r.WriteOps)
+}
+
+// Returns average number of write operations that can be performed per second.
+func (r *BenchResults) WriteOpsPerSecond() int {
+	var op = r.WriteOpDuration()
+	if op == 0 {
+		return 0
+	}
+	return int(time.Second) / int(op)
+}
+
+// Returns the duration for a single read operation.
+func (r *BenchResults) ReadOpDuration() time.Duration {
+	if r.ReadOps == 0 {
+		return 0
+	}
+	return r.ReadDuration / time.Duration(r.ReadOps)
+}
+
+// Returns average number of read operations that can be performed per second.
+func (r *BenchResults) ReadOpsPerSecond() int {
+	var op = r.ReadOpDuration()
+	if op == 0 {
+		return 0
+	}
+	return int(time.Second) / int(op)
+}
 
-	println(testing.Benchmark(b.Run))
+// tempfile returns a temporary file path.
+func tempfile() string {
+	f, _ := ioutil.TempFile("", "bolt-bench-")
+	f.Close()
+	os.Remove(f.Name())
+	return f.Name()
 }