Merge pull request #48 from benbjohnson/bucket-stat

Bucket stats

3 files changed, 107 insertions, 0 deletions

diff --git a/bucket.go b/bucket.go
index a973f0e..6653389 100644
--- a/bucket.go
+++ b/bucket.go
@@ -29,3 +29,32 @@ func (b *Bucket) cursor() *Cursor {
 		stack:       make([]pageElementRef, 0),
 	}
 }
+
+// Stat returns stats on a bucket.
+func (b *Bucket) Stat() *BucketStat {
+	s := &BucketStat{}
+	b.transaction.forEachPage(b.root, 0, func(p *page, depth int) {
+		if (p.flags & leafPageFlag) != 0 {
+			s.LeafPageCount++
+			s.KeyCount += int(p.count)
+		} else if (p.flags & branchPageFlag) != 0 {
+			s.BranchPageCount++
+		}
+
+		s.OverflowPageCount += int(p.overflow)
+
+		if depth+1 > s.MaxDepth {
+			s.MaxDepth = (depth + 1)
+		}
+	})
+	return s
+}
+
+// BucketStat represents stats on a bucket such as branch pages and leaf pages.
+type BucketStat struct {
+	BranchPageCount   int
+	LeafPageCount     int
+	OverflowPageCount int
+	KeyCount          int
+	MaxDepth          int
+}
diff --git a/bucket_test.go b/bucket_test.go
new file mode 100644
index 0000000..5e33189
--- /dev/null
+++ b/bucket_test.go
@@ -0,0 +1,62 @@
+package bolt
+
+import (
+	"strconv"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// Ensure a bucket can calculate stats.
+func TestBucketStat(t *testing.T) {
+	withOpenDB(func(db *DB, path string) {
+		db.Do(func(txn *RWTransaction) error {
+			// Add bucket with lots of keys.
+			txn.CreateBucket("widgets")
+			for i := 0; i < 100000; i++ {
+				txn.Put("widgets", []byte(strconv.Itoa(i)), []byte(strconv.Itoa(i)))
+			}
+
+			// Add bucket with fewer keys but one big value.
+			txn.CreateBucket("woojits")
+			for i := 0; i < 500; i++ {
+				txn.Put("woojits", []byte(strconv.Itoa(i)), []byte(strconv.Itoa(i)))
+			}
+			txn.Put("woojits", []byte("really-big-value"), []byte(strings.Repeat("*", 10000)))
+
+			// Add a bucket that fits on a single root leaf.
+			txn.CreateBucket("whozawhats")
+			txn.Put("whozawhats", []byte("foo"), []byte("bar"))
+
+			return nil
+		})
+		db.With(func(txn *Transaction) error {
+			b := txn.Bucket("widgets")
+			stat := b.Stat()
+			assert.Equal(t, stat.BranchPageCount, 15)
+			assert.Equal(t, stat.LeafPageCount, 1281)
+			assert.Equal(t, stat.OverflowPageCount, 0)
+			assert.Equal(t, stat.KeyCount, 100000)
+			assert.Equal(t, stat.MaxDepth, 3)
+
+			b = txn.Bucket("woojits")
+			stat = b.Stat()
+			assert.Equal(t, stat.BranchPageCount, 1)
+			assert.Equal(t, stat.LeafPageCount, 6)
+			assert.Equal(t, stat.OverflowPageCount, 2)
+			assert.Equal(t, stat.KeyCount, 501)
+			assert.Equal(t, stat.MaxDepth, 2)
+
+			b = txn.Bucket("whozawhats")
+			stat = b.Stat()
+			assert.Equal(t, stat.BranchPageCount, 0)
+			assert.Equal(t, stat.LeafPageCount, 1)
+			assert.Equal(t, stat.OverflowPageCount, 0)
+			assert.Equal(t, stat.KeyCount, 1)
+			assert.Equal(t, stat.MaxDepth, 1)
+
+			return nil
+		})
+	})
+}
diff --git a/transaction.go b/transaction.go
index 3951cd8..6e9ca8f 100644
--- a/transaction.go
+++ b/transaction.go
@@ -134,3 +134,19 @@ func (t *Transaction) page(id pgid) *page {
 	// Otherwise return directly from the mmap.
 	return t.db.page(id)
 }
+
+// forEachPage iterates over every page within a given page and executes a function.
+func (t *Transaction) forEachPage(pgid pgid, depth int, fn func(*page, int)) {
+	p := t.page(pgid)
+
+	// Execute function.
+	fn(p, depth)
+
+	// Recursively loop over children.
+	if (p.flags & branchPageFlag) != 0 {
+		for i := 0; i < int(p.count); i++ {
+			elem := p.branchPageElement(uint16(i))
+			t.forEachPage(elem.pgid, depth+1, fn)
+		}
+	}
+}