diff --git a/roaring/btree.go b/roaring/btree.go index 1051f8985..bbc429704 100644 --- a/roaring/btree.go +++ b/roaring/btree.go @@ -925,6 +925,15 @@ func (e *enumerator) Every(upd func(key uint64, oldV *Container, exists bool) (n if write { if nv == nil { e.t.Delete(i.k) + f, _ := e.t.Seek(e.k) + *e = *f + f.Close() + // we don't want to e.next() here; we'll + // already be on an item with key >= i.k, + // and since we just deleted the item with + // key i.k, that means key is > i.k, which + // makes it the next item. + continue } else { e.q.d[e.i].v = nv } diff --git a/roaring/btree_test.go b/roaring/btree_test.go index 422a7b767..c3c02af3d 100644 --- a/roaring/btree_test.go +++ b/roaring/btree_test.go @@ -11,6 +11,7 @@ import ( "math" "math/rand" "path" + "reflect" "runtime" "runtime/debug" "strings" @@ -998,6 +999,28 @@ func TestBtreeEnumeratorPrevSanity(t *testing.T) { } } +// TestBtreeEnumeratorEveryRegression is a regression test for a "use-after-free" bug. +// Previously, deleting a container would cause some values to be skipped (and sometimes trigger a race condition). +func TestBtreeEnumeratorEveryRegression(t *testing.T) { + r := treeNew() + + r.Set(uint64(10), getDummyC(100)) + r.Set(uint64(20), getDummyC(200)) + r.Set(uint64(30), getDummyC(300)) + + e, _ := r.Seek(0) + expect := []uint64{10, 20, 30} + var found []uint64 + _ = e.Every(func(key uint64, oldV *Container, exists bool) (*Container, bool) { + found = append(found, key) + return nil, true + }) + + if !reflect.DeepEqual(expect, found) { // Before the fix, this skipped the 20. + t.Errorf("had %v in bitmap; only found %v", expect, found) + } +} + func BenchmarkBtreeSeekSeq1e3(b *testing.B) { benchmarkSeekSeq(b, 1e3) }