Merge pull request #193 from benbjohnson/check-freelist-overflow

Check for freelist overflow

3 files changed, 64 insertions, 11 deletions

diff --git a/bucket_test.go b/bucket_test.go
index 21ddfca..7e959d0 100644
--- a/bucket_test.go
+++ b/bucket_test.go
@@ -190,9 +190,9 @@ func TestBucket_Delete_Large(t *testing.T) {
 	})
 }
 
-// Ensure that deleting a large set of keys will work correctly.
-// Reported by Jordan Sherer: https://github.com/boltdb/bolt/issues/184
-func TestBucket_Delete_Large2(t *testing.T) {
+// Deleting a very large list of keys will overflow the freelist.
+// https://github.com/boltdb/bolt/issues/192
+func TestBucket_Delete_ErrFreelistOverflow(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skipping test in short mode.")
 	}
@@ -223,7 +223,7 @@ func TestBucket_Delete_Large2(t *testing.T) {
 		}
 
 		// Delete all of them in one large transaction
-		db.Update(func(tx *Tx) error {
+		err := db.Update(func(tx *Tx) error {
 			b := tx.Bucket([]byte("0"))
 			c := b.Cursor()
 			for k, _ := c.First(); k != nil; k, _ = c.Next() {
@@ -231,6 +231,9 @@ func TestBucket_Delete_Large2(t *testing.T) {
 			}
 			return nil
 		})
+
+		// Check that a freelist overflow occurred.
+		assert.Equal(t, ErrFreelistOverflow, err)
 	})
 }
 
diff --git a/freelist.go b/freelist.go
index a236079..e27f80a 100644
--- a/freelist.go
+++ b/freelist.go
@@ -1,11 +1,18 @@
 package bolt
 
 import (
+	"errors"
 	"fmt"
 	"sort"
 	"unsafe"
 )
 
+var (
+	// ErrFreelistOverflow is returned when the total number of free pages
+	// exceeds 65,536 and the freelist cannot hold any more.
+	ErrFreelistOverflow = errors.New("freelist overflow")
+)
+
 // freelist represents a list of all pages that are available for allocation.
 // It also tracks pages that have been freed but are still in use by open transactions.
 type freelist struct {
@@ -106,6 +113,11 @@ func (f *freelist) release(txid txid) {
 	sort.Sort(pgids(f.ids))
 }
 
+// rollback removes the pages from a given pending tx.
+func (f *freelist) rollback(txid txid) {
+	delete(f.pending, txid)
+}
+
 // isFree returns whether a given page is in the free list.
 func (f *freelist) isFree(pgid pgid) bool {
 	for _, id := range f.ids {
@@ -134,9 +146,36 @@ func (f *freelist) read(p *page) {
 // write writes the page ids onto a freelist page. All free and pending ids are
 // saved to disk since in the event of a program crash, all pending ids will
 // become free.
-func (f *freelist) write(p *page) {
+func (f *freelist) write(p *page) error {
+	// Combine the old free pgids and pgids waiting on an open transaction.
 	ids := f.all()
+
+	// Make sure that the sum of all free pages is less than the max uint16 size.
+	if len(ids) >= 65565 {
+		return ErrFreelistOverflow
+	}
+
+	// Update the header and write the ids to the page.
 	p.flags |= freelistPageFlag
 	p.count = uint16(len(ids))
 	copy(((*[maxAllocSize]pgid)(unsafe.Pointer(&p.ptr)))[:], ids)
+
+	return nil
+}
+
+// reload reads the freelist from a page and filters out pending items.
+func (f *freelist) reload(p *page) {
+	f.read(p)
+
+	// Filter out pending free pages.
+	ids := f.ids
+	f.ids = nil
+
+	var a []pgid
+	for _, id := range ids {
+		if !f.isFree(id) {
+			a = append(a, id)
+		}
+	}
+	f.ids = a
 }
diff --git a/tx.go b/tx.go
index a9781f2..9ab1178 100644
--- a/tx.go
+++ b/tx.go
@@ -157,7 +157,7 @@ func (tx *Tx) Commit() error {
 	// spill data onto dirty pages.
 	startTime = time.Now()
 	if err := tx.root.spill(); err != nil {
-		tx.close()
+		tx.rollback()
 		return err
 	}
 	tx.stats.SpillTime += time.Since(startTime)
@@ -170,16 +170,19 @@ func (tx *Tx) Commit() error {
 	tx.db.freelist.free(tx.id(), tx.db.page(tx.meta.freelist))
 	p, err := tx.allocate((tx.db.freelist.size() / tx.db.pageSize) + 1)
 	if err != nil {
-		tx.close()
+		tx.rollback()
+		return err
+	}
+	if err := tx.db.freelist.write(p); err != nil {
+		tx.rollback()
 		return err
 	}
-	tx.db.freelist.write(p)
 	tx.meta.freelist = p.id
 
 	// Write dirty pages to disk.
 	startTime = time.Now()
 	if err := tx.write(); err != nil {
-		tx.close()
+		tx.rollback()
 		return err
 	}
 
@@ -193,7 +196,7 @@ func (tx *Tx) Commit() error {
 
 	// Write meta to disk.
 	if err := tx.writeMeta(); err != nil {
-		tx.close()
+		tx.rollback()
 		return err
 	}
 	tx.stats.WriteTime += time.Since(startTime)
@@ -215,10 +218,18 @@ func (tx *Tx) Rollback() error {
 	if tx.db == nil {
 		return ErrTxClosed
 	}
-	tx.close()
+	tx.rollback()
 	return nil
 }
 
+func (tx *Tx) rollback() {
+	if tx.writable {
+		tx.db.freelist.rollback(tx.id())
+		tx.db.freelist.reload(tx.db.page(tx.db.meta().freelist))
+	}
+	tx.close()
+}
+
 func (tx *Tx) close() {
 	if tx.writable {
 		// Remove writer lock.