From a84c6fbeac9a99091b560f8f5599c5a8576c250e Mon Sep 17 00:00:00 2001 From: arjunnair1997 Date: Thu, 30 Oct 2025 15:44:44 -0400 Subject: [PATCH 1/4] add support for re-usable iterator --- btree_test.go | 112 ++++++++++++++++++++++++++++++++++++++++++++++++++ btreeg.go | 34 ++++++++++++++- 2 files changed, 145 insertions(+), 1 deletion(-) diff --git a/btree_test.go b/btree_test.go index 7ef9cad..4d710ff 100644 --- a/btree_test.go +++ b/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) + 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) + + // 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() + } +} diff --git a/btreeg.go b/btreeg.go index 4028698..e25e193 100644 --- a/btreeg.go +++ b/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,36 @@ 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]) { + iter.tr = tr + iter.mut = false + iter.locked = tr.lock(iter.mut) + if iter.stack == nil { + iter.stack = iter.stack0[:0] + } else { + iter.stack = iter.stack[:0] + } +} + // 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. From f8667eb84f0c6cdd12e1bbefd408e33fbff790cd Mon Sep 17 00:00:00 2001 From: arjunnair1997 Date: Thu, 30 Oct 2025 16:02:49 -0400 Subject: [PATCH 2/4] add --- btreeg.go | 1 - 1 file changed, 1 deletion(-) diff --git a/btreeg.go b/btreeg.go index e25e193..6c58af1 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1717,7 +1717,6 @@ func (iter *IterG[T]) ReleaseReuseable() { // been called on the iterator before re-using it using Init. func (iter *IterG[T]) Init(tr *BTreeG[T]) { iter.tr = tr - iter.mut = false iter.locked = tr.lock(iter.mut) if iter.stack == nil { iter.stack = iter.stack0[:0] From 50c6663cb40835b7897e627b4183ce290c0e5cef Mon Sep 17 00:00:00 2001 From: arjunnair1997 Date: Thu, 30 Oct 2025 16:10:42 -0400 Subject: [PATCH 3/4] make sure iterator is truly reset --- btree_test.go | 4 ++-- btreeg.go | 18 +++++++++++++++++- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/btree_test.go b/btree_test.go index 4d710ff..c62eeec 100644 --- a/btree_test.go +++ b/btree_test.go @@ -334,7 +334,7 @@ func BenchmarkIteratorCreationAlloc(b *testing.B) { b.Run("reuse", func(b *testing.B) { b.ReportAllocs() for i := 0; i < b.N; i++ { - reusableIterPointer.Init(tr) + reusableIterPointer.Init(tr, false) useIteratorPointer(reusableIterPointer) } }) @@ -368,7 +368,7 @@ func TestBenchmarkIteratorReuseWorks(t *testing.T) { seekTo := rng.Intn(i + 1) // Reset the iterator. - reusableIter.Init(tr) + reusableIter.Init(tr, false) // Seek to the random position. found = reusableIter.Seek(largeItem{a: uint64(seekTo)}) diff --git a/btreeg.go b/btreeg.go index 6c58af1..abb62fb 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1713,16 +1713,32 @@ func (iter *IterG[T]) ReleaseReuseable() { iter.tr = nil } +func (tr *BTreeG[T]) iter(mut bool) IterG[T] { + var iter IterG[T] + iter.tr = tr + iter.mut = mut + iter.locked = tr.lock(iter.mut) + iter.stack = iter.stack0[:0] + return iter +} + // 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]) { +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. From a2ce4583d5f18f883ccfd4d749fc2f6e63719c08 Mon Sep 17 00:00:00 2001 From: arjunnair1997 Date: Thu, 30 Oct 2025 16:11:11 -0400 Subject: [PATCH 4/4] make sure iterator is truly reset --- btreeg.go | 9 --------- 1 file changed, 9 deletions(-) diff --git a/btreeg.go b/btreeg.go index abb62fb..1e8ea19 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1713,15 +1713,6 @@ func (iter *IterG[T]) ReleaseReuseable() { iter.tr = nil } -func (tr *BTreeG[T]) iter(mut bool) IterG[T] { - var iter IterG[T] - iter.tr = tr - iter.mut = mut - iter.locked = tr.lock(iter.mut) - iter.stack = iter.stack0[:0] - return iter -} - // 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) {