[draft] do not review (iterator experiments)

btree_test.go

@@ -1,8 +1,10 @@
 package btree
 
 import (
+	"math/rand"
 	"sync"
 	"testing"
+	"time"
 )
 
 func assert(x bool) {
@@ -273,3 +275,113 @@ func TestIter(t *testing.T) {
 	}
 	iter.Release()
 }
+
+type largeItem struct {
+	a uint64
+	b uint64
+	c uint64
+	d uint64
+}
+
+func useIterator(iter IterG[largeItem]) {
+	iter.Seek(largeItem{a: 0})
+	defer iter.ReleaseReuseable()
+
+	// Iterate over 10 items beginning the seeked item.
+	assert(iter.Item().a == 0)
+}
+
+func useIteratorPointer(iter *IterG[largeItem]) {
+	iter.Seek(largeItem{a: 0})
+	defer iter.ReleaseReuseable()
+
+	assert(iter.Item().a == 0)
+}
+
+// This benchmark proves that there exist cases where the iterator creation can
+// cause an allocation
+//
+// Run using: go test -run=^$ -bench ^BenchmarkIteratorCreationAlloc$ github.com/tidwall/btree
+func BenchmarkIteratorCreationAlloc(b *testing.B) {
+	tr := NewBTreeG(func(a, b largeItem) bool {
+		return a.a < b.a
+	})
+
+	for i := 0; i < 1; i++ {
+		tr.Set(largeItem{a: uint64(i * 2), b: uint64(i * 2), c: uint64(i * 2), d: uint64(i * 2)})
+	}
+
+	iter := tr.Iter()
+	iter.Seek(largeItem{a: 0})
+	assert(iter.Item().a == 0)
+
+	// The following will cause 1 allocation per op. Note that this allocation is not due
+	// to iter.stack slice allocating to grow larger. The benchmark is only performing a
+	// single seek, and the btree only has 1 item. The allocation is due to the iterator
+	// escaping to the heap.
+	b.Run("no reuse", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			iter := tr.Iter()
+			useIterator(iter)
+		}
+	})
+
+	reusableIter := tr.Iter()
+	reusableIterPointer := &reusableIter
+
+	// The following will cause 0 allocations per op, since a re-usable iterator is used.
+	b.Run("reuse", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			reusableIterPointer.Init(tr, false)
+			useIteratorPointer(reusableIterPointer)
+		}
+	})
+}
+
+// Ensure that the re-usable iterator works as expected even if the tree
+// is receiving new writes after every iteration.
+func TestBenchmarkIteratorReuseWorks(t *testing.T) {
+	tr := NewBTreeGOptions(func(a, b largeItem) bool {
+		return a.a < b.a
+	}, Options{
+		NoLocks: true,
+	})
+
+	tr.Set(largeItem{a: 0, b: 0, c: 0, d: 0})
+	iter := tr.Iter()
+	reusableIter := &iter
+
+	found := reusableIter.Seek(largeItem{a: 0})
+	assert(found)
+	assert(reusableIter.Item().a == 0)
+	reusableIter.ReleaseReuseable()
+
+	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
+
+	for i := 1; i < 10000; i++ {
+		// Insert a single item. Reset iterator, randomly seek b/w 0 to i, and ensure that the
+		// iterator is seeing the expected items.
+		tr.Set(largeItem{a: uint64(i), b: uint64(i), c: uint64(i), d: uint64(i)})
+
+		seekTo := rng.Intn(i + 1)
+
+		// Reset the iterator.
+		reusableIter.Init(tr, false)
+
+		// Seek to the random position.
+		found = reusableIter.Seek(largeItem{a: uint64(seekTo)})
+		assert(found)
+
+		nextExpectedItem := uint64(seekTo) + 1
+		for iter.Next() {
+			assert(iter.Item().a == nextExpectedItem)
+			nextExpectedItem++
+		}
+
+		assert(nextExpectedItem == uint64(i)+1)
+
+		reusableIter.ReleaseReuseable()
+	}
+}

btreeg.go

@@ -3,7 +3,9 @@
 // license that can be found in the LICENSE file.
 package btree
 
-import "sync"
+import (
+	"sync"
+)
 
 type BTreeG[T any] struct {
 	isoid        uint64
@@ -1694,6 +1696,42 @@ func (iter *IterG[T]) Release() {
 	iter.tr = nil
 }
 
+// ReleaseReuseable is the same as Release, but it preserves the iterator stack
+// so that the iterator can be reused using Init without allocating.
+func (iter *IterG[T]) ReleaseReuseable() {
+	if iter.tr == nil {
+		return
+	}
+	if iter.locked {
+		iter.tr.unlock(iter.mut)
+		iter.locked = false
+	}
+
+	// Preserve the backing memory for the stack, so that the iterator can be re-used without
+	// allocating.
+	iter.stack = iter.stack[:0]
+	iter.tr = nil
+}
+
+// Init is used to initialize an existing iterator with a new tree. ReleaseReusable must've
+// been called on the iterator before re-using it using Init.
+func (iter *IterG[T]) Init(tr *BTreeG[T], mut bool) {
+	iter.tr = tr
+	iter.mut = mut
+
+	iter.locked = tr.lock(iter.mut)
+	if iter.stack == nil {
+		iter.stack = iter.stack0[:0]
+	} else {
+		iter.stack = iter.stack[:0]
+	}
+
+	iter.seeked = false
+	iter.atstart = false
+	iter.atend = false
+	iter.item = tr.empty
+}
+
 // Next moves iterator to the next item in iterator.
 // Returns false if the tree is empty or the iterator is at the end of
 // the tree.