Merge pull request #135 from benbjohnson/bench

Bench

7 files changed, 308 insertions, 142 deletions

diff --git a/Makefile b/Makefile
index f4964a6..2c0dc97 100644
--- a/Makefile
+++ b/Makefile
@@ -5,8 +5,8 @@ BRANCH=`git rev-parse --abbrev-ref HEAD`
 COMMIT=`git rev-parse --short HEAD`
 GOLDFLAGS="-X main.branch $(BRANCH) -X main.commit $(COMMIT)"
 
-bench: benchpreq
-	go test -v -test.bench=$(BENCH)
+bench:
+	go test -v -test.run=NOTHINCONTAINSTHIS -test.bench=$(BENCH)
 
 # http://cloc.sourceforge.net/
 cloc:
@@ -34,7 +34,7 @@ get:
 
 build: get
 	@mkdir -p bin
-	@go build -ldflags=$(GOLDFLAGS) -a -o bin/bolt-`git rev-parse --short HEAD` ./cmd/bolt
+	@go build -ldflags=$(GOLDFLAGS) -a -o bin/bolt ./cmd/bolt
 
 test: fmt errcheck
 	@go get github.com/stretchr/testify/assert
diff --git a/cmd/bolt/bench.go b/cmd/bolt/bench.go
new file mode 100644
index 0000000..72144b8
--- /dev/null
+++ b/cmd/bolt/bench.go
@@ -0,0 +1,271 @@
+package main
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"runtime"
+	"runtime/pprof"
+	"time"
+
+	"github.com/boltdb/bolt"
+)
+
+// File handlers for the various profiles.
+var cpuprofile, memprofile, blockprofile *os.File
+
+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 {
+		fatal(err)
+		return
+	}
+	defer db.Close()
+
+	// 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)
+}
+
+// tempfile returns a temporary file path.
+func tempfile() string {
+	f, _ := ioutil.TempFile("", "bolt-bench-")
+	f.Close()
+	os.Remove(f.Name())
+	return f.Name()
+}
diff --git a/cmd/bolt/import.go b/cmd/bolt/import.go
index bcb7d2e..0988f32 100644
--- a/cmd/bolt/import.go
+++ b/cmd/bolt/import.go
@@ -23,6 +23,11 @@ func Import(path string, input string) {
 		fatal(err)
 	}
 
+	// Import all of the buckets.
+	importBuckets(path, root)
+}
+
+func importBuckets(path string, root []*rawMessage) {
 	// Open the database.
 	db, err := bolt.Open(path, 0600)
 	if err != nil {
diff --git a/cmd/bolt/main.go b/cmd/bolt/main.go
index 1930e7d..719bf00 100644
--- a/cmd/bolt/main.go
+++ b/cmd/bolt/main.go
@@ -90,7 +90,34 @@ func NewApp() *cli.App {
 				Check(path)
 			},
 		},
-	}
+		{
+			Name:  "bench",
+			Usage: "Performs a synthetic benchmark",
+			Flags: []cli.Flag{
+				&cli.StringFlag{Name: "profile-mode", Value: "rw", Usage: "Profile mode"},
+				&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: "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"},
+				&cli.StringFlag{Name: "memprofile", Usage: "Memory profile output path"},
+				&cli.StringFlag{Name: "blockprofile", Usage: "Block profile output path"},
+			},
+			Action: func(c *cli.Context) {
+				Bench(&BenchOptions{
+					ProfileMode:  c.String("profile-mode"),
+					WriteMode:    c.String("write-mode"),
+					ReadMode:     c.String("read-mode"),
+					Iterations:   c.Int("count"),
+					KeySize:      c.Int("key-size"),
+					ValueSize:    c.Int("value-size"),
+					CPUProfile:   c.String("cpuprofile"),
+					MemProfile:   c.String("memprofile"),
+					BlockProfile: c.String("blockprofile"),
+				})
+			},
+		}}
 	return app
 }
 
diff --git a/cmd/bolt/set.go b/cmd/bolt/set.go
index 9761f44..f4a4696 100644
--- a/cmd/bolt/set.go
+++ b/cmd/bolt/set.go
@@ -21,6 +21,7 @@ func Set(path, name, key, value string) {
 	defer db.Close()
 
 	err = db.Update(func(tx *bolt.Tx) error {
+
 		// Find bucket.
 		b := tx.Bucket([]byte(name))
 		if b == nil {
diff --git a/db_test.go b/db_test.go
index baef98e..57c356e 100644
--- a/db_test.go
+++ b/db_test.go
@@ -5,11 +5,8 @@ import (
 	"flag"
 	"fmt"
 	"io/ioutil"
-	"math/rand"
 	"os"
 	"regexp"
-	"strconv"
-	"strings"
 	"testing"
 	"time"
 	"unsafe"
@@ -356,39 +353,6 @@ func TestDBStats_Sub(t *testing.T) {
 	assert.Equal(t, 7, diff.TxStats.PageCount)
 }
 
-// Benchmark the performance of single put transactions in random order.
-func BenchmarkDB_Put_Sequential(b *testing.B) {
-	value := []byte(strings.Repeat("0", 64))
-	withOpenDB(func(db *DB, path string) {
-		db.Update(func(tx *Tx) error {
-			_, err := tx.CreateBucket([]byte("widgets"))
-			return err
-		})
-		for i := 0; i < b.N; i++ {
-			db.Update(func(tx *Tx) error {
-				return tx.Bucket([]byte("widgets")).Put([]byte(strconv.Itoa(i)), value)
-			})
-		}
-	})
-}
-
-// Benchmark the performance of single put transactions in random order.
-func BenchmarkDB_Put_Random(b *testing.B) {
-	indexes := rand.Perm(b.N)
-	value := []byte(strings.Repeat("0", 64))
-	withOpenDB(func(db *DB, path string) {
-		db.Update(func(tx *Tx) error {
-			_, err := tx.CreateBucket([]byte("widgets"))
-			return err
-		})
-		for i := 0; i < b.N; i++ {
-			db.Update(func(tx *Tx) error {
-				return tx.Bucket([]byte("widgets")).Put([]byte(strconv.Itoa(indexes[i])), value)
-			})
-		}
-	})
-}
-
 func ExampleDB_Update() {
 	// Open the database.
 	db, _ := Open(tempfile(), 0666)
diff --git a/tx_test.go b/tx_test.go
index 53ec53e..f630e97 100644
--- a/tx_test.go
+++ b/tx_test.go
@@ -3,10 +3,7 @@ package bolt
 import (
 	"errors"
 	"fmt"
-	"math/rand"
 	"os"
-	"strconv"
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -266,105 +263,6 @@ func TestTx_OnCommit_Rollback(t *testing.T) {
 	assert.Equal(t, 0, x)
 }
 
-// Benchmark the performance iterating over a cursor.
-func BenchmarkTxCursor1(b *testing.B)     { benchmarkTxCursor(b, 1) }
-func BenchmarkTxCursor10(b *testing.B)    { benchmarkTxCursor(b, 10) }
-func BenchmarkTxCursor100(b *testing.B)   { benchmarkTxCursor(b, 100) }
-func BenchmarkTxCursor1000(b *testing.B)  { benchmarkTxCursor(b, 1000) }
-func BenchmarkTxCursor10000(b *testing.B) { benchmarkTxCursor(b, 10000) }
-
-func benchmarkTxCursor(b *testing.B, total int) {
-	indexes := rand.Perm(total)
-	value := []byte(strings.Repeat("0", 100))
-
-	withOpenDB(func(db *DB, path string) {
-		// Write data to bucket.
-		db.Update(func(tx *Tx) error {
-			tx.CreateBucket([]byte("widgets"))
-			bucket := tx.Bucket([]byte("widgets"))
-			for i := 0; i < total; i++ {
-				bucket.Put([]byte(fmt.Sprintf("%016d", indexes[i])), value)
-			}
-			return nil
-		})
-		b.ResetTimer()
-
-		// Iterate over bucket using cursor.
-		for i := 0; i < b.N; i++ {
-			db.View(func(tx *Tx) error {
-				count := 0
-				c := tx.Bucket([]byte("widgets")).Cursor()
-				for k, _ := c.First(); k != nil; k, _ = c.Next() {
-					count++
-				}
-				if count != total {
-					b.Fatalf("wrong count: %d; expected: %d", count, total)
-				}
-				return nil
-			})
-		}
-	})
-}
-
-// Benchmark the performance of bulk put transactions in random order.
-func BenchmarkTxPutRandom1(b *testing.B)     { benchmarkTxPutRandom(b, 1) }
-func BenchmarkTxPutRandom10(b *testing.B)    { benchmarkTxPutRandom(b, 10) }
-func BenchmarkTxPutRandom100(b *testing.B)   { benchmarkTxPutRandom(b, 100) }
-func BenchmarkTxPutRandom1000(b *testing.B)  { benchmarkTxPutRandom(b, 1000) }
-func BenchmarkTxPutRandom10000(b *testing.B) { benchmarkTxPutRandom(b, 10000) }
-
-func benchmarkTxPutRandom(b *testing.B, total int) {
-	indexes := rand.Perm(total)
-	value := []byte(strings.Repeat("0", 64))
-	withOpenDB(func(db *DB, path string) {
-		db.Update(func(tx *Tx) error {
-			_, err := tx.CreateBucket([]byte("widgets"))
-			return err
-		})
-		var tx *Tx
-		var bucket *Bucket
-		for j := 0; j < b.N; j++ {
-			for i := 0; i < total; i++ {
-				if i%1000 == 0 {
-					if tx != nil {
-						tx.Commit()
-					}
-					tx, _ = db.Begin(true)
-					bucket = tx.Bucket([]byte("widgets"))
-				}
-				bucket.Put([]byte(strconv.Itoa(indexes[i])), value)
-			}
-		}
-		tx.Commit()
-	})
-}
-
-// Benchmark the performance of bulk put transactions in sequential order.
-func BenchmarkTxPutSequential1(b *testing.B)     { benchmarkTxPutSequential(b, 1) }
-func BenchmarkTxPutSequential10(b *testing.B)    { benchmarkTxPutSequential(b, 10) }
-func BenchmarkTxPutSequential100(b *testing.B)   { benchmarkTxPutSequential(b, 100) }
-func BenchmarkTxPutSequential1000(b *testing.B)  { benchmarkTxPutSequential(b, 1000) }
-func BenchmarkTxPutSequential10000(b *testing.B) { benchmarkTxPutSequential(b, 10000) }
-
-func benchmarkTxPutSequential(b *testing.B, total int) {
-	value := []byte(strings.Repeat("0", 64))
-	withOpenDB(func(db *DB, path string) {
-		db.Update(func(tx *Tx) error {
-			_, err := tx.CreateBucket([]byte("widgets"))
-			return err
-		})
-		db.Update(func(tx *Tx) error {
-			bucket := tx.Bucket([]byte("widgets"))
-			for j := 0; j < b.N; j++ {
-				for i := 0; i < total; i++ {
-					bucket.Put([]byte(strconv.Itoa(i)), value)
-				}
-			}
-			return nil
-		})
-	})
-}
-
 func ExampleTx_Rollback() {
 	// Open the database.
 	db, _ := Open(tempfile(), 0666)