From 5c19b55b8c7ba5a92b644d8ddb9ed07d0dbb2d4c Mon Sep 17 00:00:00 2001 From: Arjun Nair Date: Wed, 6 Aug 2025 20:53:07 -0400 Subject: [PATCH 1/8] add support for no alloc iter --- btreeg.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/btreeg.go b/btreeg.go index 4028698..f8ccbfa 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1580,6 +1580,15 @@ func (tr *BTreeG[T]) Iter() IterG[T] { return tr.iter(false) } +func (tr *BTreeG[T]) IterNoAlloc(iter *IterG[T]) *IterG[T] { + *iter = IterG[T]{} + + iter.tr = tr + iter.mut = false + iter.stack = iter.stack0[:0] + return iter +} + func (tr *BTreeG[T]) IterMut() IterG[T] { return tr.iter(true) } From 23416e2132b30d30994b76062ca38ad9f8f31825 Mon Sep 17 00:00:00 2001 From: Arjun Nair Date: Mon, 11 Aug 2025 21:12:48 -0400 Subject: [PATCH 2/8] add locking back --- btreeg.go | 1 + 1 file changed, 1 insertion(+) diff --git a/btreeg.go b/btreeg.go index f8ccbfa..4d2583e 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1585,6 +1585,7 @@ func (tr *BTreeG[T]) IterNoAlloc(iter *IterG[T]) *IterG[T] { iter.tr = tr iter.mut = false + iter.locked = tr.lock(iter.mut) iter.stack = iter.stack0[:0] return iter } From 5553ba3d63ff197a8e86d68a73a430c0be3822dd Mon Sep 17 00:00:00 2001 From: Arjun Nair Date: Wed, 29 Oct 2025 16:57:22 -0400 Subject: [PATCH 3/8] just create a re-usable iterator --- btreeg.go | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/btreeg.go b/btreeg.go index 4d2583e..a82f818 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1580,16 +1580,6 @@ func (tr *BTreeG[T]) Iter() IterG[T] { return tr.iter(false) } -func (tr *BTreeG[T]) IterNoAlloc(iter *IterG[T]) *IterG[T] { - *iter = IterG[T]{} - - iter.tr = tr - iter.mut = false - iter.locked = tr.lock(iter.mut) - iter.stack = iter.stack0[:0] - return iter -} - func (tr *BTreeG[T]) IterMut() IterG[T] { return tr.iter(true) } @@ -1704,6 +1694,20 @@ 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 without allocating. +func (iter *IterG[T]) ReleaseReuseable() { + if iter.tr == nil { + return + } + if iter.locked { + iter.tr.unlock(iter.mut) + iter.locked = false + } + iter.stack = iter.stack[:0] + iter.tr = nil +} + // 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 6922c1e028d8b9673545de102d3b3c100ee7870a Mon Sep 17 00:00:00 2001 From: Arjun Nair Date: Wed, 29 Oct 2025 17:43:21 -0400 Subject: [PATCH 4/8] add benchmark to show that re-usable iterator causes 0 allocations --- btreeg.go | 18 ++++++++++++++- btreeg_test.go | 59 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/btreeg.go b/btreeg.go index a82f818..9e9a7ed 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1695,7 +1695,7 @@ func (iter *IterG[T]) Release() { } // ReleaseReuseable is the same as Release, but it preserves the iterator stack -// so that the iterator can be reused without allocating. +// so that the iterator can be reused using Init without allocating. func (iter *IterG[T]) ReleaseReuseable() { if iter.tr == nil { return @@ -1704,10 +1704,26 @@ func (iter *IterG[T]) ReleaseReuseable() { 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] + } +} + // 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. diff --git a/btreeg_test.go b/btreeg_test.go index 26f8c5f..cdba5b1 100644 --- a/btreeg_test.go +++ b/btreeg_test.go @@ -1872,3 +1872,62 @@ func TestContiguousDeleteRangeLoad(t *testing.T) { ) } + +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 +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. + b.Run("no reuse", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + iter := tr.Iter() + useIterator(iter) + } + }) + + // The following will cause 0 allocations per op. + reusableIter := tr.Iter() + reusableIterPointer := &reusableIter + + b.Run("reuse", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + reusableIterPointer.Init(tr, false) + useIteratorPointer(reusableIterPointer) + } + }) +} From 66a3b5fae71bba7707092e7eaa255a9a5edc6bcc Mon Sep 17 00:00:00 2001 From: Arjun Nair Date: Wed, 29 Oct 2025 17:44:44 -0400 Subject: [PATCH 5/8] commits --- btreeg_test.go | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/btreeg_test.go b/btreeg_test.go index cdba5b1..0b79293 100644 --- a/btreeg_test.go +++ b/btreeg_test.go @@ -1910,7 +1910,10 @@ func BenchmarkIteratorCreationAlloc(b *testing.B) { iter.Seek(largeItem{a: 0}) assert(iter.Item().a == 0) - // The following will cause 1 allocation per op. + // 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++ { @@ -1919,10 +1922,10 @@ func BenchmarkIteratorCreationAlloc(b *testing.B) { } }) - // The following will cause 0 allocations per op. 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++ { From 386cf47d5cd2e8e244ccb9448510919209390841 Mon Sep 17 00:00:00 2001 From: Arjun Nair Date: Wed, 29 Oct 2025 17:45:14 -0400 Subject: [PATCH 6/8] comment --- btreeg_test.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/btreeg_test.go b/btreeg_test.go index 0b79293..3c747da 100644 --- a/btreeg_test.go +++ b/btreeg_test.go @@ -1897,6 +1897,8 @@ func useIteratorPointer(iter *IterG[largeItem]) { // 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 From 6aa44a672a4f9a9227bc472ea84982b95f03741c Mon Sep 17 00:00:00 2001 From: Arjun Nair Date: Thu, 30 Oct 2025 16:17:37 -0400 Subject: [PATCH 7/8] truly reset iterator --- btreeg.go | 6 ++++++ btreeg_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/btreeg.go b/btreeg.go index 9e9a7ed..2aa3742 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1716,12 +1716,18 @@ func (iter *IterG[T]) ReleaseReuseable() { 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. diff --git a/btreeg_test.go b/btreeg_test.go index 3c747da..bb6f08c 100644 --- a/btreeg_test.go +++ b/btreeg_test.go @@ -1936,3 +1936,49 @@ func BenchmarkIteratorCreationAlloc(b *testing.B) { } }) } + +// 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() + } +} From 421ae8a6004ea6e4677155ddfd1fd68b055cab45 Mon Sep 17 00:00:00 2001 From: Arjun Nair Date: Mon, 3 Nov 2025 07:32:48 -0800 Subject: [PATCH 8/8] re-use Release instead of ReleaseReusable --- btreeg.go | 31 +++++++++---------------------- btreeg_test.go | 8 ++++---- 2 files changed, 13 insertions(+), 26 deletions(-) diff --git a/btreeg.go b/btreeg.go index 2aa3742..218a2db 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1690,30 +1690,22 @@ func (iter *IterG[T]) Release() { iter.tr.unlock(iter.mut) iter.locked = false } - iter.stack = nil - 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 + stack := iter.stack[:0] + *iter = IterG[T]{} + iter.stack = stack } -// Init is used to initialize an existing iterator with a new tree. ReleaseReusable must've +// Init is used to initialize an existing iterator with a new tree. Release must've // been called on the iterator before re-using it using Init. func (iter *IterG[T]) Init(tr *BTreeG[T], mut bool) { + // Re-use the stack, but 0 out the rest of the memory. + stack := iter.stack[:0] + *iter = IterG[T]{} + iter.stack = stack + iter.tr = tr iter.mut = mut @@ -1723,11 +1715,6 @@ func (iter *IterG[T]) Init(tr *BTreeG[T], mut bool) { } 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. diff --git a/btreeg_test.go b/btreeg_test.go index bb6f08c..686758b 100644 --- a/btreeg_test.go +++ b/btreeg_test.go @@ -1882,7 +1882,7 @@ type largeItem struct { func useIterator(iter IterG[largeItem]) { iter.Seek(largeItem{a: 0}) - defer iter.ReleaseReuseable() + defer iter.Release() // Iterate over 10 items beginning the seeked item. assert(iter.Item().a == 0) @@ -1890,7 +1890,7 @@ func useIterator(iter IterG[largeItem]) { func useIteratorPointer(iter *IterG[largeItem]) { iter.Seek(largeItem{a: 0}) - defer iter.ReleaseReuseable() + defer iter.Release() assert(iter.Item().a == 0) } @@ -1953,7 +1953,7 @@ func TestBenchmarkIteratorReuseWorks(t *testing.T) { found := reusableIter.Seek(largeItem{a: 0}) assert(found) assert(reusableIter.Item().a == 0) - reusableIter.ReleaseReuseable() + reusableIter.Release() rng := rand.New(rand.NewSource(time.Now().UnixNano())) @@ -1979,6 +1979,6 @@ func TestBenchmarkIteratorReuseWorks(t *testing.T) { assert(nextExpectedItem == uint64(i)+1) - reusableIter.ReleaseReuseable() + reusableIter.Release() } }