diff --git a/btreeg.go b/btreeg.go index 4028698..218a2db 100644 --- a/btreeg.go +++ b/btreeg.go @@ -1690,8 +1690,31 @@ func (iter *IterG[T]) Release() { iter.tr.unlock(iter.mut) iter.locked = false } - iter.stack = nil - iter.tr = nil + + // Preserve the backing memory for the stack, so that the iterator can be re-used without + // allocating. + stack := iter.stack[:0] + *iter = IterG[T]{} + iter.stack = stack +} + +// 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 + + 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. diff --git a/btreeg_test.go b/btreeg_test.go index 26f8c5f..686758b 100644 --- a/btreeg_test.go +++ b/btreeg_test.go @@ -1872,3 +1872,113 @@ 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.Release() + + // 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.Release() + + 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.Release() + + 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.Release() + } +}