diff --git a/btree_test.go b/btree_test.go index 7ef9cad..c62eeec 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, 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() + } +} diff --git a/btreeg.go b/btreeg.go index 4028698..1e8ea19 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,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.