Fix bucket free.

5 files changed, 82 insertions, 6 deletions

diff --git a/bucket.go b/bucket.go
index 6672356..4ad62fe 100644
--- a/bucket.go
+++ b/bucket.go
@@ -392,6 +392,40 @@ func (b *Bucket) forEachPage(fn func(*page, int)) {
 	b.tx.forEachPage(b.root, 0, fn)
 }
 
+// forEachPageNode iterates over every page (or node) in a bucket.
+// This also includes inline pages.
+func (b *Bucket) forEachPageNode(fn func(*page, *node, int)) {
+	// If we have an inline page or root node then just use that.
+	if b.page != nil {
+		fn(b.page, nil, 0)
+		return
+	}
+	b._forEachPageNode(b.root, 0, fn)
+}
+
+func (b *Bucket) _forEachPageNode(pgid pgid, depth int, fn func(*page, *node, int)) {
+	var p, n = b.pageNode(pgid)
+
+	// Execute function.
+	fn(p, n, depth)
+
+	// Recursively loop over children.
+	if p != nil {
+		if (p.flags & branchPageFlag) != 0 {
+			for i := 0; i < int(p.count); i++ {
+				elem := p.branchPageElement(uint16(i))
+				b._forEachPageNode(elem.pgid, depth+1, fn)
+			}
+		}
+	} else {
+		if !n.isLeaf {
+			for _, inode := range n.inodes {
+				b._forEachPageNode(inode.pgid, depth+1, fn)
+			}
+		}
+	}
+}
+
 // spill writes all the nodes for this bucket to dirty pages.
 func (b *Bucket) spill() error {
 	// Spill all child buckets first.
@@ -541,8 +575,12 @@ func (b *Bucket) free() {
 	}
 
 	var tx = b.tx
-	tx.forEachPage(b.root, 0, func(p *page, _ int) {
-		tx.db.freelist.free(tx.id(), p)
+	b.forEachPageNode(func(p *page, n *node, _ int) {
+		if p != nil {
+			tx.db.freelist.free(tx.id(), p)
+		} else {
+			n.free()
+		}
 	})
 	b.root = 0
 }
diff --git a/bucket_test.go b/bucket_test.go
index 8fcf84c..b87031b 100644
--- a/bucket_test.go
+++ b/bucket_test.go
@@ -161,6 +161,33 @@ func TestBucket_Delete(t *testing.T) {
 	})
 }
 
+// Ensure that deleting a large set of keys will work correctly.
+func TestBucket_Delete_Large(t *testing.T) {
+	withOpenDB(func(db *DB, path string) {
+		db.Update(func(tx *Tx) error {
+			var b, _ = tx.CreateBucket([]byte("widgets"))
+			for i := 0; i < 100; i++ {
+				assert.NoError(t, b.Put([]byte(strconv.Itoa(i)), []byte(strings.Repeat("*", 1024))))
+			}
+			return nil
+		})
+		db.Update(func(tx *Tx) error {
+			var b = tx.Bucket([]byte("widgets"))
+			for i := 0; i < 100; i++ {
+				assert.NoError(t, b.Delete([]byte(strconv.Itoa(i))))
+			}
+			return nil
+		})
+		db.View(func(tx *Tx) error {
+			var b = tx.Bucket([]byte("widgets"))
+			for i := 0; i < 100; i++ {
+				assert.Nil(t, b.Get([]byte(strconv.Itoa(i))))
+			}
+			return nil
+		})
+	})
+}
+
 // Ensure that accessing and updating nested buckets is ok across transactions.
 func TestBucket_Nested(t *testing.T) {
 	withOpenDB(func(db *DB, path string) {
diff --git a/cursor.go b/cursor.go
index 56948fc..31a3d5f 100644
--- a/cursor.go
+++ b/cursor.go
@@ -148,7 +148,7 @@ func (c *Cursor) seek(seek []byte) (key []byte, value []byte, flags uint32) {
 func (c *Cursor) first() {
 	for {
 		// Exit when we hit a leaf page.
-		ref := &c.stack[len(c.stack)-1]
+		var ref = &c.stack[len(c.stack)-1]
 		if ref.isLeaf() {
 			break
 		}
diff --git a/freelist.go b/freelist.go
index 6c527fe..ebe2810 100644
--- a/freelist.go
+++ b/freelist.go
@@ -1,7 +1,6 @@
 package bolt
 
 import (
-	"fmt"
 	"sort"
 	"unsafe"
 )
@@ -64,7 +63,7 @@ func (f *freelist) free(txid txid, p *page) {
 	}
 	f.pending[txid] = ids
 
-	f.check() // DEBUG ONLY:
+	// DEBUG ONLY: f.check()
 }
 
 // release moves all page ids for a transaction id (or older) to the freelist.
@@ -115,6 +114,7 @@ func (f *freelist) write(p *page) {
 // check verifies there are no double free pages.
 // This is slow so it should only be used while debugging.
 // If errors are found then a panic invoked.
+/*
 func (f *freelist) check() {
 	var lookup = make(map[pgid]txid)
 	for _, id := range f.ids {
@@ -132,6 +132,7 @@ func (f *freelist) check() {
 		}
 	}
 }
+*/
 
 type reverseSortedPgids []pgid
 
diff --git a/node.go b/node.go
index 72145c4..1bc889d 100644
--- a/node.go
+++ b/node.go
@@ -336,7 +336,7 @@ func (n *node) rebalance() {
 		// If root node is a branch and only has one node then collapse it.
 		if !n.isLeaf && len(n.inodes) == 1 {
 			// Move root's child up.
-			child := n.bucket.nodes[n.inodes[0].pgid]
+			child := n.bucket.node(n.inodes[0].pgid, n)
 			n.isLeaf = child.isLeaf
 			n.inodes = child.inodes[:]
 			n.children = child.children
@@ -357,6 +357,16 @@ func (n *node) rebalance() {
 		return
 	}
 
+	// If node has no keys then just remove it.
+	if n.numChildren() == 0 {
+		n.parent.del(n.key)
+		n.parent.removeChild(n)
+		delete(n.bucket.nodes, n.pgid)
+		n.free()
+		n.parent.rebalance()
+		return
+	}
+
 	_assert(n.parent.numChildren() > 1, "parent must have at least 2 children")
 
 	// Destination node is right sibling if idx == 0, otherwise left sibling.