diff --git a/any_table.go b/any_table.go index b73d3907..9dd52e92 100644 --- a/any_table.go +++ b/any_table.go @@ -22,8 +22,13 @@ func (t AnyTable) NumObjects(txn ReadTxn) int { } func (t AnyTable) All(txn ReadTxn) iter.Seq2[any, Revision] { - all, _ := t.AllWatch(txn) - return all + indexTxn := txn.mustIndexReadTxn(t.Meta, PrimaryIndexPos) + all := indexTxn.allNoWatch() + return func(yield func(any, Revision) bool) { + all.All(func(_ []byte, iobj object) bool { + return yield(iobj.data, iobj.revision) + }) + } } func (t AnyTable) AllWatch(txn ReadTxn) (iter.Seq2[any, Revision], <-chan struct{}) { @@ -42,7 +47,7 @@ func (t AnyTable) UnmarshalYAML(data []byte) (any, error) { func (t AnyTable) Insert(txn WriteTxn, obj any) (old any, hadOld bool, err error) { var iobj object - iobj, hadOld, _, err = txn.unwrap().insert(t.Meta, Revision(0), obj) + iobj, hadOld, _, err = txn.unwrap().insert(t.Meta, Revision(0), obj, false) if hadOld { old = iobj.data } @@ -63,7 +68,7 @@ func (t AnyTable) Get(txn ReadTxn, index string, key string) (any, Revision, boo if err != nil { return nil, 0, false, err } - obj, _, found := itxn.get(rawKey) + obj, found := itxn.getNoWatch(rawKey) if found { return obj.data, obj.revision, found, nil } @@ -75,7 +80,7 @@ func (t AnyTable) Prefix(txn ReadTxn, index string, key string) (iter.Seq2[any, if err != nil { return nil, err } - iter, _ := itxn.prefix(rawKey) + iter := itxn.prefixNoWatch(rawKey) return objSeq[any](iter), nil } @@ -84,7 +89,7 @@ func (t AnyTable) LowerBound(txn ReadTxn, index string, key string) (iter.Seq2[a if err != nil { return nil, err } - iter, _ := itxn.lowerBound(rawKey) + iter := itxn.lowerBoundNoWatch(rawKey) return objSeq[any](iter), nil } @@ -93,7 +98,7 @@ func (t AnyTable) List(txn ReadTxn, index string, key string) (iter.Seq2[any, Re if err != nil { return nil, err } - iter, _ := itxn.list(rawKey) + iter := itxn.listNoWatch(rawKey) return objSeq[any](iter), nil } diff --git a/deletetracker.go b/deletetracker.go index b134a48a..eae387df 100644 --- a/deletetracker.go +++ b/deletetracker.go @@ -37,7 +37,7 @@ func (dt *deleteTracker[Obj]) getRevision() uint64 { // called! func (dt *deleteTracker[Obj]) deleted(txn ReadTxn, minRevision Revision) *iterator[Obj] { indexEntry := txn.root()[dt.table.tablePos()].indexes[GraveyardRevisionIndexPos] - objs, _ := indexEntry.lowerBoundNext(index.Uint64(minRevision)) + objs := indexEntry.lowerBoundNextNoWatch(index.Uint64(minRevision)) return &iterator[Obj]{objs} } diff --git a/graveyard.go b/graveyard.go index cc18bae8..44be3752 100644 --- a/graveyard.go +++ b/graveyard.go @@ -63,7 +63,7 @@ func graveyardWorker(db *DB, ctx context.Context, gcRateLimitInterval time.Durat // to the low watermark. indexTree := rtxn.mustIndexReadTxn(table.meta, GraveyardRevisionIndexPos) - iter, _ := indexTree.all() + iter := indexTree.allNoWatch() for key, obj := range iter.All { if obj.revision > lowWatermark { break diff --git a/http.go b/http.go index d2fb8253..16a37c74 100644 --- a/http.go +++ b/http.go @@ -131,9 +131,9 @@ type QueryResponse struct { func runQuery(reader tableIndexReader, lowerbound bool, queryKey index.Key, onObject func(object) error) { var iter tableIndexIterator if lowerbound { - iter, _ = reader.lowerBound(queryKey) + iter = reader.lowerBoundNoWatch(queryKey) } else { - iter, _ = reader.list(queryKey) + iter = reader.listNoWatch(queryKey) } for _, obj := range iter.All { if err := onObject(obj); err != nil { diff --git a/iterator.go b/iterator.go index 50c24436..99737933 100644 --- a/iterator.go +++ b/iterator.go @@ -172,7 +172,7 @@ func (it *changeIterator[Obj]) refresh(txn ReadTxn) { panic(fmt.Sprintf("Table[%T].Changes().Next() called with the target table locked. This is not supported.", obj)) } indexEntry := tableEntry.indexes[RevisionIndexPos] - updated, _ := indexEntry.lowerBoundNext(index.Uint64(it.revision + 1)) + updated := indexEntry.lowerBoundNextNoWatch(index.Uint64(it.revision + 1)) updateIter := &iterator[Obj]{updated} deleteIter := it.dt.deleted(txn, it.deleteRevision+1) it.iter = newDualIterator(deleteIter, updateIter) diff --git a/lpm_index.go b/lpm_index.go index 2b28a4ec..348c7985 100644 --- a/lpm_index.go +++ b/lpm_index.go @@ -203,6 +203,10 @@ func (l lpmIndex) all() (tableIndexIterator, <-chan struct{}) { return newLPMIterator(l.lpm.All()), l.watch } +func (l lpmIndex) allNoWatch() tableIndexIterator { + return newLPMIterator(l.lpm.All()) +} + // get implements tableIndex. func (l lpmIndex) get(ikey index.Key) (object, <-chan struct{}, bool) { entry, found := l.lpm.Lookup(ikey) @@ -215,6 +219,14 @@ func (l lpmIndex) get(ikey index.Key) (object, <-chan struct{}, bool) { return object{}, l.watch, false } +func (l lpmIndex) getNoWatch(ikey index.Key) (object, bool) { + entry, found := l.lpm.Lookup(ikey) + if !found { + return object{}, false + } + return entry.first() +} + // len implements tableIndex. func (l lpmIndex) len() int { return l.size @@ -229,16 +241,32 @@ func (l lpmIndex) list(key index.Key) (tableIndexIterator, <-chan struct{}) { return &entry, l.watch } +func (l lpmIndex) listNoWatch(key index.Key) tableIndexIterator { + entry, found := l.lpm.Lookup(key) + if !found || entry.len() == 0 { + return emptyTableIndexIterator + } + return &entry +} + // lowerBound implements tableIndex. func (l lpmIndex) lowerBound(key index.Key) (tableIndexIterator, <-chan struct{}) { return newLPMIterator(l.lpm.LowerBound(key)), l.watch } +func (l lpmIndex) lowerBoundNoWatch(key index.Key) tableIndexIterator { + return newLPMIterator(l.lpm.LowerBound(key)) +} + // lowerBoundNext implements tableIndexTxn. func (l lpmIndex) lowerBoundNext(key index.Key) (func() ([]byte, object, bool), <-chan struct{}) { return newLPMNextFunc(l.lpm.LowerBound(key)), l.watch } +func (l lpmIndex) lowerBoundNextNoWatch(key index.Key) func() ([]byte, object, bool) { + return newLPMNextFunc(l.lpm.LowerBound(key)) +} + // objectToKey implements tableIndex. func (l lpmIndex) objectToKey(obj object) index.Key { return l.objectToKeys(obj).First() @@ -249,6 +277,10 @@ func (l lpmIndex) prefix(key index.Key) (tableIndexIterator, <-chan struct{}) { return newLPMIterator(l.lpm.Prefix(key)), l.watch } +func (l lpmIndex) prefixNoWatch(key index.Key) tableIndexIterator { + return newLPMIterator(l.lpm.Prefix(key)) +} + // rootWatch implements tableIndex. func (l lpmIndex) rootWatch() <-chan struct{} { return l.watch @@ -287,6 +319,10 @@ func (l *lpmIndexTxn) all() (tableIndexIterator, <-chan struct{}) { return newLPMIterator(l.tx.All()), l.index.watch } +func (l *lpmIndexTxn) allNoWatch() tableIndexIterator { + return newLPMIterator(l.tx.All()) +} + // commit implements tableIndexTxn. func (l *lpmIndexTxn) commit() (tableIndex, tableIndexTxnNotify) { lpm := l.tx.Commit() @@ -311,11 +347,19 @@ func (l *lpmIndexTxn) insert(key index.Key, obj object) (old object, hadOld bool panic("LPM index cannot be the primary index") } +func (l *lpmIndexTxn) insertNoWatch(key index.Key, obj object) (old object, hadOld bool) { + panic("LPM index cannot be the primary index") +} + // modify implements tableIndexTxn. func (l *lpmIndexTxn) modify(key index.Key, obj object, mod func(old, new object) object) (old object, newObj object, hadOld bool, watch <-chan struct{}) { panic("LPM index cannot be the primary index") } +func (l *lpmIndexTxn) modifyNoWatch(key index.Key, obj object, mod func(old, new object) object) (old object, newObj object, hadOld bool) { + panic("LPM index cannot be the primary index") +} + // get implements tableIndexTxn. func (l *lpmIndexTxn) get(key index.Key) (object, <-chan struct{}, bool) { entry, found := l.tx.Lookup(key) @@ -328,6 +372,14 @@ func (l *lpmIndexTxn) get(key index.Key) (object, <-chan struct{}, bool) { return object{}, l.index.watch, false } +func (l *lpmIndexTxn) getNoWatch(key index.Key) (object, bool) { + entry, found := l.tx.Lookup(key) + if !found { + return object{}, false + } + return entry.first() +} + // len implements tableIndexTxn. func (l *lpmIndexTxn) len() int { return l.size @@ -342,16 +394,32 @@ func (l *lpmIndexTxn) list(key index.Key) (tableIndexIterator, <-chan struct{}) return &entry, l.index.watch } +func (l *lpmIndexTxn) listNoWatch(key index.Key) tableIndexIterator { + entry, found := l.tx.Lookup(key) + if !found || entry.len() == 0 { + return emptyTableIndexIterator + } + return &entry +} + // lowerBound implements tableIndexTxn. func (l *lpmIndexTxn) lowerBound(key index.Key) (tableIndexIterator, <-chan struct{}) { return newLPMIterator(l.tx.LowerBound(key)), l.index.watch } +func (l *lpmIndexTxn) lowerBoundNoWatch(key index.Key) tableIndexIterator { + return newLPMIterator(l.tx.LowerBound(key)) +} + // lowerBoundNext implements tableIndexTxn. func (l *lpmIndexTxn) lowerBoundNext(key index.Key) (func() ([]byte, object, bool), <-chan struct{}) { return newLPMNextFunc(l.tx.LowerBound(key)), l.index.watch } +func (l *lpmIndexTxn) lowerBoundNextNoWatch(key index.Key) func() ([]byte, object, bool) { + return newLPMNextFunc(l.tx.LowerBound(key)) +} + // notify implements tableIndexTxn. func (l *lpmIndexTxn) notify() { if l.index.watch != nil { @@ -370,6 +438,10 @@ func (l *lpmIndexTxn) prefix(key index.Key) (tableIndexIterator, <-chan struct{} return newLPMIterator(l.tx.Prefix(key)), l.index.watch } +func (l *lpmIndexTxn) prefixNoWatch(key index.Key) tableIndexIterator { + return newLPMIterator(l.tx.Prefix(key)) +} + // reindex implements tableIndexTxn. func (l *lpmIndexTxn) reindex(primaryKey index.Key, old object, new object) { var newKeys index.KeySet diff --git a/part/iterator.go b/part/iterator.go index afe3a97f..39d8cdd4 100644 --- a/part/iterator.go +++ b/part/iterator.go @@ -138,7 +138,7 @@ func newIterator[T any](start *header[T]) Iterator[T] { return Iterator[T]{start: start} } -func prefixSearch[T any](root *header[T], rootWatch <-chan struct{}, prefix []byte) (Iterator[T], <-chan struct{}) { +func prefixSearch[T any](root *header[T], rootWatch *watchState, prefix []byte) (Iterator[T], *watchState) { if root == nil { return newIterator[T](nil), rootWatch } diff --git a/part/map.go b/part/map.go index e9755639..2153a081 100644 --- a/part/map.go +++ b/part/map.go @@ -75,7 +75,7 @@ func (m Map[K, V]) Get(key K) (value V, found bool) { if !m.hasTree { return } - kv, _, found := m.tree.Get(m.keyToBytes(key)) + kv, found := m.tree.Get(m.keyToBytes(key)) return kv.Value, found } @@ -178,7 +178,7 @@ func (m Map[K, V]) Prefix(prefix K) iter.Seq2[K, V] { if !m.hasTree { return toSeq2[K, V](Iterator[mapKVPair[K, V]]{}) } - iter, _ := m.tree.Prefix(m.keyToBytes(prefix)) + iter := m.tree.Prefix(m.keyToBytes(prefix)) return toSeq2(iter) } @@ -442,14 +442,14 @@ func (txn MapTxn[K, V]) Delete(key K) bool { // Get a value from the map by its key. func (txn MapTxn[K, V]) Get(key K) (value V, found bool) { - kv, _, found := txn.txn.Get(txn.bytesFromKeyFunc(key)) + kv, found := txn.txn.Get(txn.bytesFromKeyFunc(key)) return kv.Value, found } // Prefix iterates in order over all keys that start with // the given prefix. func (txn MapTxn[K, V]) Prefix(prefix K) iter.Seq2[K, V] { - iter, _ := txn.txn.Prefix(txn.bytesFromKeyFunc(prefix)) + iter := txn.txn.Prefix(txn.bytesFromKeyFunc(prefix)) return toSeq2(iter) } diff --git a/part/node.go b/part/node.go index f4e935ba..1f0402cc 100644 --- a/part/node.go +++ b/part/node.go @@ -25,8 +25,8 @@ const ( type header[T any] struct { flags uint16 // kind(4b) | unused(3b) | size(9b) prefixLen uint16 - prefixP *byte // the compressed prefix, [0] is the key - watch chan struct{} // watch channel that is closed when this node mutates + prefixP *byte // the compressed prefix, [0] is the key + watch *watchState // watch that is closed when this node mutates } func (n *header[T]) key() byte { @@ -204,7 +204,7 @@ func (n *header[T]) clone(watch bool) *header[T] { panic(fmt.Sprintf("unknown node kind: %x", n.kind())) } if watch { - nCopy.watch = make(chan struct{}) + nCopy.watch = newWatchState() } else { nCopy.watch = nil } @@ -221,7 +221,7 @@ func (n *header[T]) promote(txnID uint64) *header[T] { node4.txnID = txnID node4.setKind(nodeKind4) if n.watch != nil { - node4.watch = make(chan struct{}) + node4.watch = newWatchState() } return node4.self() case nodeKind4: @@ -234,7 +234,7 @@ func (n *header[T]) promote(txnID uint64) *header[T] { copy(node16.children[:], node4.children[:size]) copy(node16.keys[:], node4.keys[:size]) if n.watch != nil { - node16.watch = make(chan struct{}) + node16.watch = newWatchState() } return node16.self() case nodeKind16: @@ -248,7 +248,7 @@ func (n *header[T]) promote(txnID uint64) *header[T] { node48.index[k] = uint8(i + 1) } if n.watch != nil { - node48.watch = make(chan struct{}) + node48.watch = newWatchState() } return node48.self() case nodeKind48: @@ -264,7 +264,7 @@ func (n *header[T]) promote(txnID uint64) *header[T] { node256.children[child.prefix()[0]] = child } if n.watch != nil { - node256.watch = make(chan struct{}) + node256.watch = newWatchState() } return node256.self() case nodeKind256: @@ -274,15 +274,6 @@ func (n *header[T]) promote(txnID uint64) *header[T] { } } -func isClosedChan(ch <-chan struct{}) bool { - select { - case <-ch: - return true - default: - return false - } -} - func (n *header[T]) printTree(level int) { if n == nil { return @@ -309,9 +300,9 @@ func (n *header[T]) printTree(level int) { panic("unknown node kind") } if leaf := n.getLeaf(); leaf != nil { - fmt.Printf(" %x -> %v (L:%p W:%p %v)", leaf.fullKey(), leaf.value, leaf, leaf.watch, isClosedChan(leaf.watch)) + fmt.Printf(" %x -> %v (L:%p W:%p %v)", leaf.fullKey(), leaf.value, leaf, leaf.watch, leaf.watch.isClosed()) } - fmt.Printf(" (N:%p, W:%p %v)\n", n, n.watch, isClosedChan(n.watch)) + fmt.Printf(" (N:%p, W:%p %v)\n", n, n.watch, n.watch.isClosed()) for _, child := range children { if child != nil { @@ -538,7 +529,7 @@ func newLeaf[T any](o options, prefix, key []byte, value T) *leaf[T] { leaf.setPrefix(prefix) leaf.setKind(nodeKindLeaf) if !o.rootOnlyWatch() { - leaf.watch = make(chan struct{}) + leaf.watch = newWatchState() } return leaf @@ -575,7 +566,7 @@ type node256[T any] struct { children [256]*header[T] } -func search[T any](root *header[T], rootWatch <-chan struct{}, key []byte) (value T, watch <-chan struct{}, ok bool) { +func search[T any](root *header[T], rootWatch *watchState, key []byte) (value T, watch *watchState, ok bool) { this := root watch = rootWatch if root == nil { diff --git a/part/ops.go b/part/ops.go index 7e8b5b99..66ec489d 100644 --- a/part/ops.go +++ b/part/ops.go @@ -10,15 +10,19 @@ type Ops[T any] interface { Len() int // Get fetches the value associated with the given key. - // Returns the value, a watch channel (which is closed on - // modification to the key) and boolean which is true if - // value was found. - Get(key []byte) (T, <-chan struct{}, bool) - - // Prefix returns an iterator for all objects that starts with the - // given prefix, and a channel that closes when any objects matching - // the given prefix are upserted or deleted. - Prefix(key []byte) (Iterator[T], <-chan struct{}) + Get(key []byte) (T, bool) + + // GetWatch fetches the value and returns a channel that closes when the + // key is modified. + GetWatch(key []byte) (T, <-chan struct{}, bool) + + // Prefix returns an iterator for all objects that start with the given + // prefix. + Prefix(key []byte) Iterator[T] + + // PrefixWatch returns matching objects and a channel that closes when any + // matching object is upserted or deleted. + PrefixWatch(key []byte) (Iterator[T], <-chan struct{}) // LowerBound returns an iterator for all objects that have a // key equal or higher than the given 'key'. diff --git a/part/part_test.go b/part/part_test.go index b4802d6a..861243ef 100644 --- a/part/part_test.go +++ b/part/part_test.go @@ -8,6 +8,7 @@ import ( "encoding/binary" "fmt" "math/rand" + "slices" "testing" "time" @@ -53,15 +54,15 @@ func Test_insertion_and_watches(t *testing.T) { tree = txn.CommitAndNotify() assertOpen(t, watch_ab) - _, w, f := tree.Get([]byte("ab")) + _, w, f := tree.GetWatch([]byte("ab")) assert.True(t, f) assertOpen(t, w) - _, w2 := tree.Prefix([]byte("a")) + _, w2 := tree.PrefixWatch([]byte("a")) assertOpen(t, w2) - _, w3, f2 := tree.Get([]byte("abc")) + _, w3, f2 := tree.GetWatch([]byte("abc")) assert.True(t, f2) assertOpen(t, w3) - _, w4 := tree.Prefix([]byte("abc")) + _, w4 := tree.PrefixWatch([]byte("abc")) assertOpen(t, w4) _, _, tree = tree.Insert([]byte("ab"), 42) @@ -88,10 +89,10 @@ func Test_insertion_and_watches(t *testing.T) { _, _, tree = tree.Insert([]byte("a"), 1) - _, w, f := tree.Get([]byte("a")) + _, w, f := tree.GetWatch([]byte("a")) assert.True(t, f) assertOpen(t, w) - _, w2 := tree.Prefix([]byte("a")) + _, w2 := tree.PrefixWatch([]byte("a")) assertOpen(t, w2) _, _, tree = tree.Insert([]byte("b"), 2) @@ -105,10 +106,10 @@ func Test_insertion_and_watches(t *testing.T) { _, _, tree = tree.Insert([]byte("a"), 1) - _, w, f := tree.Get([]byte("a")) + _, w, f := tree.GetWatch([]byte("a")) assert.True(t, f) assertOpen(t, w) - _, w2 := tree.Prefix([]byte("a")) + _, w2 := tree.PrefixWatch([]byte("a")) assertOpen(t, w2) txn := tree.Txn() @@ -137,11 +138,11 @@ func Test_insertion_and_watches(t *testing.T) { txn.Insert([]byte("ab"), 3) tree = txn.CommitAndNotify() - _, w, f := tree.Get([]byte("ab")) + _, w, f := tree.GetWatch([]byte("ab")) assert.True(t, f) assertOpen(t, w) - _, w2 := tree.Prefix([]byte("ab")) + _, w2 := tree.PrefixWatch([]byte("ab")) assertOpen(t, w2) // This should move the L(ab) laterally and insert a N4(ab) with a L(abc, 4) @@ -161,9 +162,9 @@ func Test_insertion_and_watches(t *testing.T) { _, _, tree = tree.Insert([]byte("a"), 1) - _, w, found := tree.Get([]byte("a")) + _, w, found := tree.GetWatch([]byte("a")) assert.True(t, found) - _, w2, found := tree.Get([]byte("aa")) + _, w2, found := tree.GetWatch([]byte("aa")) assert.False(t, found) assert.NotEqual(t, w, w2, "did not expect Get(aa) to return watch channel of Get(a)") @@ -196,7 +197,7 @@ func Test_watchClosingRandom(t *testing.T) { watches := []<-chan struct{}{} for _, key := range keys { - _, watch, found := tree.Get(key) + _, watch, found := tree.GetWatch(key) require.True(t, found) watches = append(watches, watch) } @@ -226,6 +227,68 @@ func Test_watchClosingRandom(t *testing.T) { } } +func TestTxnWatchCollection(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + keys []string + operations func(*Txn[int]) + }{ + { + name: "root leaf deletion", + keys: []string{"a"}, + operations: func(txn *Txn[int]) { + txn.Delete([]byte("a")) + }, + }, + { + name: "child leaf deletion", + keys: []string{"a", "b"}, + operations: func(txn *Txn[int]) { + txn.Delete([]byte("a")) + }, + }, + { + name: "internal leaf deletion", + keys: []string{"a", "ab", "ac"}, + operations: func(txn *Txn[int]) { + txn.Delete([]byte("a")) + }, + }, + { + name: "repeated replacement and deletion", + keys: []string{"a"}, + operations: func(txn *Txn[int]) { + txn.Insert([]byte("a"), 2) + txn.Insert([]byte("a"), 3) + txn.Delete([]byte("a")) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + tree := New[int]() + txn := tree.Txn() + for i, key := range test.keys { + txn.Insert([]byte(key), i) + } + tree = txn.CommitAndNotify() + + txn = tree.Txn() + test.operations(txn) + + watches := slices.Clone(txn.watches) + require.NotEmpty(t, watches) + txn.CommitAndNotify() + for _, watch := range watches { + require.True(t, watch.isClosed()) + } + }) + } +} + func Test_commonPrefix(t *testing.T) { t.Parallel() @@ -258,13 +321,13 @@ func Test_search(t *testing.T) { _, _, tree = tree.Insert([]byte("c"), []byte("c")) _, _, tree = tree.Insert([]byte("ca"), []byte("ca")) - v, _, ok := tree.Get([]byte("nope")) + v, ok := tree.Get([]byte("nope")) if ok { t.Fatalf("found unexpected value: %v", v) } for _, key := range []string{"a", "ba", "bb"} { - v, _, ok = tree.Get([]byte(key)) + v, ok = tree.Get([]byte(key)) if !ok || string(v) != key { t.Fatalf("%q not found (%v) or mismatch %q", key, ok, v) } @@ -288,7 +351,7 @@ func Test_simple_delete(t *testing.T) { tree := New[uint64]() - _, _, found := tree.Get(uint64Key(1)) + _, found := tree.Get(uint64Key(1)) require.False(t, found) txn := tree.Txn() @@ -296,7 +359,7 @@ func Test_simple_delete(t *testing.T) { _, hadOld := txn.Insert(uint64Key(1), 1) require.False(t, hadOld) - _, watch, found := txn.Get(uint64Key(1)) + _, watch, found := txn.GetWatch(uint64Key(1)) require.True(t, found) select { @@ -310,7 +373,7 @@ func Test_simple_delete(t *testing.T) { _, hadOld = txn.Delete(uint64Key(1)) require.False(t, hadOld) - _, _, found = txn.Get(uint64Key(1)) + _, found = txn.Get(uint64Key(1)) require.False(t, found) tree = txn.CommitAndNotify() @@ -339,7 +402,7 @@ func Test_simple_delete(t *testing.T) { _, hadOld = txn.Delete(uint64Key(5)) require.True(t, hadOld) - _, _, ok := txn.Get(uint64Key(5)) + _, ok := txn.Get(uint64Key(5)) require.False(t, ok) } @@ -378,7 +441,7 @@ func Test_delete_compress_regression(t *testing.T) { // Retrieve an unknown key that returns the watch channel // for the internal node 'foo-' - _, watch, _ := tree.Get([]byte("foo-2/ccc")) + _, watch, _ := tree.GetWatch([]byte("foo-2/ccc")) txn := tree.Txn() // Remove the keys under foo-1. This will extend the prefix of the @@ -428,7 +491,7 @@ func Test_delete(t *testing.T) { for _, i := range keys { _, hadOld = txn.Insert(uint64Key(i), i) assert.False(t, hadOld) - v, _, ok := txn.Get(uint64Key(i)) + v, ok := txn.Get(uint64Key(i)) assert.True(t, ok) assert.EqualValues(t, v, i) } @@ -442,20 +505,20 @@ func Test_delete(t *testing.T) { txn = tree.Txn() for _, i := range keys { - v, _, ok := txn.Get(uint64Key(i)) + v, ok := txn.Get(uint64Key(i)) assert.True(t, ok) assert.EqualValues(t, v, i) v, hadOld = txn.Delete(uint64Key(i)) assert.True(t, hadOld) assert.EqualValues(t, v, i) - _, _, ok = txn.Get(uint64Key(i)) + _, ok = txn.Get(uint64Key(i)) assert.False(t, ok) } tree = txn.CommitAndNotify() assert.Equal(t, 0, tree.Len()) for _, i := range keys { - _, _, ok := tree.Get(uint64Key(i)) + _, ok := tree.Get(uint64Key(i)) assert.False(t, ok) } @@ -471,7 +534,7 @@ func Test_delete(t *testing.T) { for _, i := range keys { _, hadOld = txn.Insert(uint64Key(i), i) assert.False(t, hadOld) - v, watch, ok := txn.Get(uint64Key(i)) + v, watch, ok := txn.GetWatch(uint64Key(i)) watches[i] = watch assert.True(t, ok) assert.EqualValues(t, v, i) @@ -484,7 +547,7 @@ func Test_delete(t *testing.T) { // Lookup with a Txn txn = tree.Txn() for _, i := range keys { - v, _, ok := txn.Get(uint64Key(i)) + v, ok := txn.Get(uint64Key(i)) assert.True(t, ok) assert.EqualValues(t, v, i) } @@ -537,7 +600,7 @@ func Test_delete(t *testing.T) { // Test that prefix iteration is ordered and correct prev = 0 - iter, _ = tree.Prefix([]byte{}) + iter = tree.Prefix([]byte{}) for { _, v, ok := iter.Next() if !ok { @@ -561,7 +624,7 @@ func Test_delete(t *testing.T) { _, hadOld = txn.Delete(k) assert.True(t, hadOld) - _, _, ok := txn.Get(k) + _, ok := txn.Get(k) assert.False(t, ok) } @@ -580,7 +643,7 @@ func Test_delete(t *testing.T) { // Check that everything is gone after commit. for _, i := range keys { - _, _, ok := tree.Get(uint64Key(i)) + _, ok := tree.Get(uint64Key(i)) assert.False(t, ok) } @@ -596,7 +659,7 @@ func Test_watch(t *testing.T) { // Insert 'a', get it and check watch channel is not closed. _, _, tree = tree.Insert([]byte("a"), []byte("a")) - _, watchA, ok := tree.Get([]byte("a")) + _, watchA, ok := tree.GetWatch([]byte("a")) if !ok { t.Fatal("expected to find 'a'") } @@ -608,7 +671,7 @@ func Test_watch(t *testing.T) { // Get 'b' that should not exist and the watch channel should // not be closed. - _, watchB, ok := tree.Get([]byte("b")) + _, watchB, ok := tree.GetWatch([]byte("b")) assert.False(t, ok, "Get(b)") select { @@ -626,7 +689,7 @@ func Test_watch(t *testing.T) { t.Fatal("expected watch channel to close") } - v, _, ok := tree.Get([]byte("a")) + v, ok := tree.Get([]byte("a")) if !ok { t.Fatal("expected to find 'a'") } @@ -651,7 +714,7 @@ func Test_insert(t *testing.T) { } for i := range 1000 { key := binary.NativeEndian.AppendUint32(nil, uint32(i)) - _, _, ok := tree.Get(key) + _, ok := tree.Get(key) if !ok { t.Fatalf("%d not found", i) } @@ -666,7 +729,7 @@ func Test_modify(t *testing.T) { // Modify without the value existing inserts it. _, _, tree = tree.Modify(key, 1, func(x, _ int) int { return 123 }) - v, _, ok := tree.Get(key) + v, ok := tree.Get(key) require.True(t, ok) require.Equal(t, 1, v) @@ -678,7 +741,7 @@ func Test_modify(t *testing.T) { } tree = txn.CommitAndNotify() - v, _, ok = tree.Get(key) + v, ok = tree.Get(key) require.True(t, ok) require.Equal(t, 1001, v) } @@ -693,11 +756,11 @@ func Test_replaceRoot(t *testing.T) { _, _, tree = tree.Insert(keyB, 3) _, _, tree = tree.Delete(keyA) _, _, tree = tree.Insert(keyA, 2) - val, _, ok := tree.Get(keyA) + val, ok := tree.Get(keyA) if !ok || val != 2 { t.Fatalf("%v not found", keyA) } - val, _, ok = tree.Get(keyB) + val, ok = tree.Get(keyB) if !ok || val != 3 { t.Fatalf("%v not found", keyB) } @@ -710,7 +773,7 @@ func Test_deleteRoot(t *testing.T) { keyA := []byte{'a'} _, _, tree = tree.Insert(keyA, 1) _, _, tree = tree.Delete(keyA) - _, _, ok := tree.Get(keyA) + _, ok := tree.Get(keyA) if ok { t.Fatal("Root exists") } @@ -727,15 +790,15 @@ func Test_deleteIntermediate(t *testing.T) { _, _, tree = tree.Insert(keyAB, 2) _, _, tree = tree.Insert(keyABC, 3) _, _, tree = tree.Delete(keyAB) - _, _, ok := tree.Get(keyA) + _, ok := tree.Get(keyA) if !ok { t.Fatal("A doesn't exist") } - _, _, ok = tree.Get(keyAB) + _, ok = tree.Get(keyAB) if ok { t.Fatal("AB exists") } - _, _, ok = tree.Get(keyABC) + _, ok = tree.Get(keyABC) if !ok { t.Fatal("ABC doesn't exist") } @@ -750,11 +813,11 @@ func Test_deleteNonExistantIntermediate(t *testing.T) { _, _, tree = tree.Insert(keyAB, 1) _, _, tree = tree.Insert(keyAC, 2) _, _, tree = tree.Delete([]byte{'a'}) - _, _, ok := tree.Get(keyAB) + _, ok := tree.Get(keyAB) if !ok { t.Fatal("AB doesn't exist") } - _, _, ok = tree.Get(keyAC) + _, ok = tree.Get(keyAC) if !ok { t.Fatal("AC doesn't exist") } @@ -767,7 +830,7 @@ func Test_deleteNonExistantCommonPrefix(t *testing.T) { keyAB := []byte{'a', 'b', 'c'} _, _, tree = tree.Insert(keyAB, 1) _, _, tree = tree.Delete([]byte{'a', 'b', 'e'}) - _, _, ok := tree.Get(keyAB) + _, ok := tree.Get(keyAB) if !ok { t.Fatal("AB doesn't exist") } @@ -800,7 +863,7 @@ func Test_prefix(t *testing.T) { ins("abcd") ins("bc") - iter, _ := tree.Prefix([]byte("ab")) + iter := tree.Prefix([]byte("ab")) k, v, ok := iter.Next() assert.True(t, ok) assert.Equal(t, []byte("ab"), k) @@ -827,7 +890,7 @@ func Test_deleteEmptyKey(t *testing.T) { _, _, tree = tree.Insert([]byte{}, "x") - v, watch, ok := tree.Get([]byte{}) + v, watch, ok := tree.GetWatch([]byte{}) assert.True(t, ok) assert.Equal(t, "x", v) select { @@ -838,7 +901,7 @@ func Test_deleteEmptyKey(t *testing.T) { _, _, tree = tree.Delete([]byte{}) - _, _, ok = tree.Get([]byte{}) + _, ok = tree.Get([]byte{}) assert.False(t, ok) select { @@ -883,7 +946,7 @@ func Test_txn(t *testing.T) { // Original tree should be untouched. for i := 1; i <= 3; i++ { - _, _, ok := tree.Get(uint64Key(uint64(i))) + _, ok := tree.Get(uint64Key(uint64(i))) assert.True(t, ok, "Get(%d)", i) } @@ -1071,11 +1134,11 @@ func Test_prefix_regression(t *testing.T) { _, _, tree = tree.Insert([]byte("foobar"), "foobar") _, _, tree = tree.Insert([]byte("foo"), "foo") - s, _, found := tree.Get([]byte("foobar")) + s, found := tree.Get([]byte("foobar")) require.True(t, found) require.Equal(t, s, "foobar") - s, _, found = tree.Get([]byte("foo")) + s, found = tree.Get([]byte("foo")) require.True(t, found) require.Equal(t, s, "foo") } @@ -1099,7 +1162,7 @@ func Test_iterate(t *testing.T) { watches := []<-chan struct{}{} for _, i := range keys { _, _, tree = tree.Insert(hexKey(uint64(i)), uint64(i)) - v, watch, ok := tree.Get(hexKey(uint64(i))) + v, watch, ok := tree.GetWatch(hexKey(uint64(i))) require.True(t, ok, "Get %x", hexKey(uint64(i))) require.Equal(t, v, uint64(i), "values equal") require.NotNil(t, watch, "watch not nil") @@ -1167,7 +1230,7 @@ func Test_iterate(t *testing.T) { // All the rest of the keys can still be found for _, j := range keys[i+1:] { - n, _, found := txn.Get(hexKey(j)) + n, found := txn.Get(hexKey(j)) if !assert.True(t, found) || !assert.Equal(t, n, j) { fmt.Println("--- new tree") txn.PrintTree() @@ -1196,7 +1259,7 @@ func Test_closed_chan_regression(t *testing.T) { // No reachable channel should be closed for _, c := range tree.root.children() { select { - case <-c.watch: + case <-c.watch.channel(): t.Logf("%x %p closed already", c.prefix(), &c.watch) t.FailNow() default: @@ -1258,7 +1321,7 @@ func TestEmptyKey(t *testing.T) { _, _, tree = tree.Insert([]byte(""), 1) require.Equal(t, 1, tree.Len()) - v, _, found := tree.Get([]byte("")) + v, found := tree.Get([]byte("")) require.True(t, found) require.Equal(t, 1, v) @@ -1351,6 +1414,28 @@ func benchmark_Insert(b *testing.B, opts ...Option) { b.ReportMetric(float64(b.N*numObjectsToInsert)/b.Elapsed().Seconds(), "objects/sec") } +func Benchmark_WatchReplace(b *testing.B) { + keys := make([][]byte, numObjectsToInsert) + tree := New[int]() + txn := tree.Txn() + for i := range keys { + keys[i] = binary.BigEndian.AppendUint64(nil, uint64(i)) + txn.Insert(keys[i], i) + } + tree = txn.CommitAndNotify() + + b.ResetTimer() + for b.Loop() { + txn := tree.Txn() + for i, key := range keys { + txn.Insert(key, i) + } + tree = txn.CommitAndNotify() + } + b.StopTimer() + b.ReportMetric(float64(b.N*numObjectsToInsert)/b.Elapsed().Seconds(), "objects/sec") +} + func benchmark_Modify_vs_GetInsert(b *testing.B, doGetInsert bool) { tree := New[int](RootOnlyWatch) keys := [][]byte{} @@ -1364,7 +1449,7 @@ func benchmark_Modify_vs_GetInsert(b *testing.B, doGetInsert bool) { txn := tree.Txn() for _, key := range keys { if doGetInsert { - v, _, _ := txn.Get(key) + v, _ := txn.Get(key) txn.Insert(key, v) } else { txn.Modify(key, 123, func(x, _ int) int { return x }) @@ -1494,7 +1579,7 @@ func Benchmark_Get(b *testing.B) { for b.Loop() { for j := range uint64(numObjectsToInsert) { binary.BigEndian.PutUint64(key[:], j) - v, _, ok := tree.Get(key[:]) + v, ok := tree.Get(key[:]) if v != j { b.Fatalf("impossible: %d != %d || %v", v, j, ok) } @@ -1504,6 +1589,25 @@ func Benchmark_Get(b *testing.B) { b.ReportMetric(float64(numObjectsToInsert*b.N)/b.Elapsed().Seconds(), "objects/sec") } +func Benchmark_GetWatch(b *testing.B) { + tree := New[uint64](RootOnlyWatch) + for j := range uint64(numObjectsToInsert) { + _, _, tree = tree.Insert(uint64Key(j), j) + } + + var key [8]byte // to avoid the allocation + for b.Loop() { + for j := range uint64(numObjectsToInsert) { + binary.BigEndian.PutUint64(key[:], j) + v, _, ok := tree.GetWatch(key[:]) + if v != j { + b.Fatalf("impossible: %d != %d || %v", v, j, ok) + } + } + } + b.ReportMetric(float64(numObjectsToInsert*b.N)/b.Elapsed().Seconds(), "objects/sec") +} + func Benchmark_All(b *testing.B) { tree := New[uint64](RootOnlyWatch) for j := uint64(1); j <= numObjectsToInsert; j++ { diff --git a/part/quick_test.go b/part/quick_test.go index f796b767..21c792da 100644 --- a/part/quick_test.go +++ b/part/quick_test.go @@ -30,13 +30,13 @@ func TestQuick_InsertGetPrefix(t *testing.T) { watchChannels := []<-chan struct{}{} // Add all possible watch channels for prefixes of the key for i := range len(key) { - _, watch := tree.Prefix([]byte(key)[:i]) + _, watch := tree.PrefixWatch([]byte(key)[:i]) watchChannels = append(watchChannels, watch) } // Check that root watch channel always closes on modifications. - watchChannels = append(watchChannels, tree.rootWatch) + watchChannels = append(watchChannels, tree.rootWatch.channel()) - _, watchBefore, _ := tree.Get([]byte(key)) + _, watchBefore, _ := tree.GetWatch([]byte(key)) if watchBefore == nil { return "nil watch from Get()" } @@ -48,7 +48,7 @@ func TestQuick_InsertGetPrefix(t *testing.T) { return "nil watch from InsertWatch()" } - val, watchAfter, found := txn.Get([]byte(key)) + val, watchAfter, found := txn.GetWatch([]byte(key)) if !found { return fmt.Sprintf("inserted value not found for %q", key) } @@ -73,14 +73,14 @@ func TestQuick_InsertGetPrefix(t *testing.T) { } get := func(key, value string) any { - val, watch, found := tree.Get([]byte(key)) + val, watch, found := tree.GetWatch([]byte(key)) if watch == nil { panic("nil watch from Get()") } for i := range len(key) { prefix := []byte(key)[:i] - iter, watch := tree.Prefix(prefix) + iter, watch := tree.PrefixWatch(prefix) if watch == nil { panic("nil watch from Prefix()") } @@ -101,7 +101,7 @@ func TestQuick_InsertGetPrefix(t *testing.T) { return val } - iter, watch := tree.Prefix([]byte(key)) + iter, watch := tree.PrefixWatch([]byte(key)) _, v, _ := iter.Next() return v } @@ -123,12 +123,12 @@ func TestQuick_IteratorReuse(t *testing.T) { iterate := func(key, value string, cloneFirst bool) bool { _, _, tree = tree.Insert([]byte(key), value) - v, _, ok := tree.Get([]byte(key)) + v, ok := tree.Get([]byte(key)) if !ok || value != v { return false } - prefixIter, _ := tree.Prefix([]byte(key)) + prefixIter := tree.Prefix([]byte(key)) iterators := []Iterator[string]{ tree.LowerBound([]byte(key)), prefixIter, @@ -165,7 +165,7 @@ func TestQuick_Delete(t *testing.T) { do := func(key, value string, delete bool) bool { _, _, tree = tree.Insert([]byte(key), value) treeAfterInsert := tree - v, watch, ok := tree.Get([]byte(key)) + v, watch, ok := tree.GetWatch([]byte(key)) if !ok || v != value { t.Logf("value not in tree after insert") return false @@ -174,13 +174,13 @@ func TestQuick_Delete(t *testing.T) { // delete some of the time to construct different variations of trees. if delete { _, _, tree = tree.Delete([]byte(key)) - _, _, ok := tree.Get([]byte(key)) + _, ok := tree.Get([]byte(key)) if ok { t.Logf("value exists after delete") return false } - _, _, ok = treeAfterInsert.Get([]byte(key)) + _, ok = treeAfterInsert.Get([]byte(key)) if !ok { t.Logf("value deleted from original") } @@ -207,7 +207,7 @@ func TestQuick_ClosedWatch(t *testing.T) { _, _, tree = tree.Insert([]byte(key), value) treeAfterInsert := tree - val, watch, ok := tree.Get([]byte(key)) + val, watch, ok := tree.GetWatch([]byte(key)) if !ok { return false } @@ -231,13 +231,13 @@ func TestQuick_ClosedWatch(t *testing.T) { } // Original tree unaffected. - val, _, ok = treeAfterInsert.Get([]byte(key)) + val, ok = treeAfterInsert.Get([]byte(key)) if !ok || val != value { t.Logf("original changed!") return false } - val, _, ok = tree.Get([]byte(key)) + val, ok = tree.Get([]byte(key)) if !ok || val != "x" { t.Logf("new tree does not have x!") return false diff --git a/part/set.go b/part/set.go index 181ffe1f..f634911d 100644 --- a/part/set.go +++ b/part/set.go @@ -79,7 +79,7 @@ func (s Set[T]) Has(v T) bool { if !s.hasTree { return false } - _, _, found := s.tree.Get(s.toBytes(v)) + _, found := s.tree.Get(s.toBytes(v)) return found } diff --git a/part/tree.go b/part/tree.go index 3b188975..fcd149df 100644 --- a/part/tree.go +++ b/part/tree.go @@ -15,7 +15,7 @@ import ( // This allows watching any part of the tree (any prefix) for changes. type Tree[T any] struct { root *header[T] - rootWatch chan struct{} + rootWatch *watchState size int // the number of objects in the tree opts options prevTxn *atomic.Pointer[Txn[T]] // the previous txn for reusing the allocation @@ -30,7 +30,7 @@ func New[T any](opts ...Option) Tree[T] { } t := Tree[T]{ root: nil, - rootWatch: make(chan struct{}), + rootWatch: newWatchState(), size: 0, opts: o, prevTxn: &atomic.Pointer[Txn[T]]{}, @@ -46,9 +46,7 @@ type Option func(*options) var RootOnlyWatch = (*options).setRootOnlyWatch func newTxn[T any](o options) *Txn[T] { - txn := &Txn[T]{ - watches: make(map[chan struct{}]struct{}), - } + txn := &Txn[T]{} txn.deleteParentsCache = make([]deleteParent[T], 0, 32) txn.opts = o return txn @@ -64,6 +62,7 @@ func (t *Tree[T]) Txn() *Txn[T] { if prevTxn := t.prevTxn.Swap(nil); prevTxn != nil { txn = prevTxn clear(txn.watches) + txn.watches = txn.watches[:0] txn.dirty = false } else { txn = newTxn[T](t.opts) @@ -84,26 +83,37 @@ func (t *Tree[T]) Len() int { } // Get fetches the value associated with the given key. -// Returns the value, a watch channel (which is closed on -// modification to the key) and boolean which is true if -// value was found. -func (t *Tree[T]) Get(key []byte) (T, <-chan struct{}, bool) { +func (t *Tree[T]) Get(key []byte) (T, bool) { + value, _, ok := search(t.root, t.rootWatch, key) + return value, ok +} + +// GetWatch fetches the value associated with the given key and returns a watch +// channel that closes when the key is modified. +func (t *Tree[T]) GetWatch(key []byte) (T, <-chan struct{}, bool) { value, watch, ok := search(t.root, t.rootWatch, key) - return value, watch, ok + return value, watch.channel(), ok +} + +// Prefix returns an iterator for all objects that start with the given prefix. +func (t *Tree[T]) Prefix(prefix []byte) Iterator[T] { + iter, _ := prefixSearch(t.root, t.rootWatch, prefix) + return iter } -// Prefix returns an iterator for all objects that starts with the -// given prefix, and a channel that closes when any objects matching -// the given prefix are upserted or deleted. -func (t *Tree[T]) Prefix(prefix []byte) (Iterator[T], <-chan struct{}) { - return prefixSearch(t.root, t.rootWatch, prefix) +// PrefixWatch returns an iterator for all objects that start with the given +// prefix and a channel that closes when any matching object is upserted or +// deleted. +func (t *Tree[T]) PrefixWatch(prefix []byte) (Iterator[T], <-chan struct{}) { + iter, watch := prefixSearch(t.root, t.rootWatch, prefix) + return iter, watch.channel() } // RootWatch returns a watch channel for the root of the tree. // Since this is the channel associated with the root, this closes // when there are any changes to the tree. func (t *Tree[T]) RootWatch() <-chan struct{} { - return t.rootWatch + return t.rootWatch.channel() } // LowerBound returns an iterator for all keys that have a value @@ -153,6 +163,6 @@ func (t *Tree[T]) All(yield func([]byte, T) bool) { // PrintTree to the standard output. For debugging. func (t *Tree[T]) PrintTree() { - fmt.Printf("rootWatch: %p %v\n", t.rootWatch, isClosedChan(t.rootWatch)) + fmt.Printf("rootWatch: %p %v\n", t.rootWatch, t.rootWatch.isClosed()) t.root.printTree(0) } diff --git a/part/txn.go b/part/txn.go index a2404ca4..e309d455 100644 --- a/part/txn.go +++ b/part/txn.go @@ -16,7 +16,7 @@ import ( type Txn[T any] struct { root *header[T] oldRoot *header[T] - rootWatch chan struct{} + rootWatch *watchState prevTxn *atomic.Pointer[Txn[T]] dirty bool @@ -35,7 +35,7 @@ type Txn[T any] struct { // watches contains the channels of cloned nodes that should be closed // when transaction is committed. - watches map[chan struct{}]struct{} + watches []*watchState // deleteParentsCache keeps the last allocated slice to avoid // reallocating it on every deletion. @@ -72,7 +72,11 @@ func (txn *Txn[T]) Clone() Tree[T] { // Insert or update the tree with the given key and value. // Returns the old value if it exists. func (txn *Txn[T]) Insert(key []byte, value T) (old T, hadOld bool) { - old, hadOld, _ = txn.InsertWatch(key, value) + old, _, hadOld, _, txn.root = txn.insert(txn.root, key, value) + validateTree(txn.root, nil, txn.watches, txn.txnID) + if !hadOld { + txn.size++ + } return } @@ -80,14 +84,16 @@ func (txn *Txn[T]) Insert(key []byte, value T) (old T, hadOld bool) { // Returns the old value if it exists and a watch channel that closes when the // key changes again. func (txn *Txn[T]) InsertWatch(key []byte, value T) (old T, hadOld bool, watch <-chan struct{}) { - old, _, hadOld, watch, txn.root = txn.insert(txn.root, key, value) + var state *watchState + old, _, hadOld, state, txn.root = txn.insert(txn.root, key, value) validateTree(txn.root, nil, txn.watches, txn.txnID) if !hadOld { txn.size++ } if txn.opts.rootOnlyWatch() { - watch = txn.rootWatch + state = txn.rootWatch } + watch = state.channel() return } @@ -95,7 +101,11 @@ func (txn *Txn[T]) InsertWatch(key []byte, value T) (old T, hadOld bool, watch < // caller to not mutate the value in-place and to return a clone. // Returns the old value (if it exists) and the new possibly merged value. func (txn *Txn[T]) Modify(key []byte, value T, mod func(T, T) T) (old T, newValue T, hadOld bool) { - old, newValue, hadOld, _ = txn.ModifyWatch(key, value, mod) + old, newValue, hadOld, _, txn.root = txn.modify(txn.root, key, value, mod) + validateTree(txn.root, nil, txn.watches, txn.txnID) + if !hadOld { + txn.size++ + } return } @@ -105,14 +115,16 @@ func (txn *Txn[T]) Modify(key []byte, value T, mod func(T, T) T) (old T, newValu // Returns the old value (if it exists) and the new possibly merged value, // and a watch channel that closes when the key changes again. func (txn *Txn[T]) ModifyWatch(key []byte, value T, mod func(T, T) T) (old T, newValue T, hadOld bool, watch <-chan struct{}) { - old, newValue, hadOld, watch, txn.root = txn.modify(txn.root, key, value, mod) + var state *watchState + old, newValue, hadOld, state, txn.root = txn.modify(txn.root, key, value, mod) validateTree(txn.root, nil, txn.watches, txn.txnID) if !hadOld { txn.size++ } if txn.opts.rootOnlyWatch() { - watch = txn.rootWatch + state = txn.rootWatch } + watch = state.channel() return } @@ -131,25 +143,38 @@ func (txn *Txn[T]) Delete(key []byte) (old T, hadOld bool) { // Since this is the channel associated with the root, this closes // when there are any changes to the tree. func (txn *Txn[T]) RootWatch() <-chan struct{} { - return txn.rootWatch + return txn.rootWatch.channel() } // Get fetches the value associated with the given key. -// Returns the value, a watch channel (which is closed on -// modification to the key) and boolean which is true if -// value was found. -func (txn *Txn[T]) Get(key []byte) (T, <-chan struct{}, bool) { +func (txn *Txn[T]) Get(key []byte) (T, bool) { + value, _, ok := search(txn.root, txn.rootWatch, key) + return value, ok +} + +// GetWatch fetches the value associated with the given key and returns a watch +// channel that closes when the key is modified. +func (txn *Txn[T]) GetWatch(key []byte) (T, <-chan struct{}, bool) { value, watch, ok := search(txn.root, txn.rootWatch, key) - return value, watch, ok + return value, watch.channel(), ok +} + +// Prefix returns an iterator for all objects that start with the given prefix. +func (txn *Txn[T]) Prefix(key []byte) Iterator[T] { + // Bump txnID in order to freeze the current tree. + txn.txnID++ + iter, _ := prefixSearch(txn.root, txn.rootWatch, key) + return iter } -// Prefix returns an iterator for all objects that starts with the -// given prefix, and a channel that closes when any objects matching -// the given prefix are upserted or deleted. -func (txn *Txn[T]) Prefix(key []byte) (Iterator[T], <-chan struct{}) { +// PrefixWatch returns an iterator for all objects that start with the given +// prefix and a channel that closes when any matching object is upserted or +// deleted. +func (txn *Txn[T]) PrefixWatch(key []byte) (Iterator[T], <-chan struct{}) { // Bump txnID in order to freeze the current tree. txn.txnID++ - return prefixSearch(txn.root, txn.rootWatch, key) + iter, watch := prefixSearch(txn.root, txn.rootWatch, key) + return iter, watch.channel() } // LowerBound returns an iterator for all objects that have a @@ -181,7 +206,7 @@ func (txn *Txn[T]) CommitAndNotify() Tree[T] { func (txn *Txn[T]) Commit() Tree[T] { newRootWatch := txn.rootWatch if txn.dirty { - newRootWatch = make(chan struct{}) + newRootWatch = newWatchState() validateTree(txn.oldRoot, nil, nil, txn.txnID) validateTree(txn.root, nil, txn.watches, txn.txnID) } @@ -199,28 +224,31 @@ func (txn *Txn[T]) Commit() Tree[T] { return t } -// watchesReuseThreshold is the threshold at which the [txn.watches] hash -// map is reused for next transaction. -const watchesReuseThreshold = 64 +// watchesReuseThreshold is the largest [txn.watches] backing array retained +// for the next transaction. At the threshold this retains roughly 8KB. +const watchesReuseThreshold = 1024 // Notify closes the watch channels of nodes that were // mutated as part of this transaction. Must be called before // Tree.Txn() is used again. func (txn *Txn[T]) Notify() { - for ch := range txn.watches { - close(ch) + for _, watch := range txn.watches { + watch.close() } if !txn.dirty && len(txn.watches) > 0 { panic("BUG: watch channels marked but txn not dirty") } - // Clear or reallocate the watches hash map for the next transaction. - if len(txn.watches) <= watchesReuseThreshold { - clear(txn.watches) + // Clear the channels to allow them to be collected. Avoid retaining an + // unusually large backing array on the transaction. + numWatches := len(txn.watches) + clear(txn.watches) + if numWatches <= watchesReuseThreshold { + txn.watches = txn.watches[:0] } else { - txn.watches = make(map[chan struct{}]struct{}) + txn.watches = nil } if txn.dirty && txn.rootWatch != nil { - close(txn.rootWatch) + txn.rootWatch.close() txn.rootWatch = nil } if !txn.opts.rootOnlyWatch() { @@ -232,7 +260,7 @@ func (txn *Txn[T]) Notify() { func (txn *Txn[T]) PrintTree() { txn.root.printTree(0) fmt.Printf("watches: ") - for watch := range txn.watches { + for _, watch := range txn.watches { fmt.Printf("%p ", watch) } fmt.Println() @@ -244,19 +272,24 @@ func (txn *Txn[T]) cloneNode(n *header[T]) *header[T] { // be mutated in-place. return n } - if n.watch != nil { - txn.watches[n.watch] = struct{}{} - } + txn.addWatch(n.watch) n = n.clone(!txn.opts.rootOnlyWatch()) n.setTxnID(txn.txnID) return n } -func (txn *Txn[T]) insert(root *header[T], key []byte, value T) (oldValue T, newValue T, hadOld bool, watch <-chan struct{}, newRoot *header[T]) { +func (txn *Txn[T]) addWatch(watch *watchState) { + if watch == nil { + return + } + txn.watches = append(txn.watches, watch) +} + +func (txn *Txn[T]) insert(root *header[T], key []byte, value T) (oldValue T, newValue T, hadOld bool, watch *watchState, newRoot *header[T]) { return txn.modify(root, key, value, nil) } -func (txn *Txn[T]) modify(root *header[T], key []byte, newValue T, mod func(T, T) T) (oldValue T, newValueOut T, hadOld bool, watch <-chan struct{}, newRoot *header[T]) { +func (txn *Txn[T]) modify(root *header[T], key []byte, newValue T, mod func(T, T) T) (oldValue T, newValueOut T, hadOld bool, watch *watchState, newRoot *header[T]) { txn.dirty = true fullKey := key newValueOut = newValue @@ -290,9 +323,7 @@ func (txn *Txn[T]) modify(root *header[T], key []byte, newValue T, mod func(T, T // We've found a free slot where to insert the key. if this.size()+1 > this.cap() { // Node too small, promote it to the next size. - if this.watch != nil { - txn.watches[this.watch] = struct{}{} - } + txn.addWatch(this.watch) this = this.promote(txn.txnID) } else { // Node is big enough, clone it so we can mutate it @@ -368,7 +399,7 @@ func (txn *Txn[T]) modify(root *header[T], key []byte, newValue T, mod func(T, T newNode.setPrefix(common) newNode.setKind(nodeKind4) if !txn.opts.rootOnlyWatch() { - newNode.watch = make(chan struct{}) + newNode.watch = newWatchState() } switch { @@ -461,22 +492,18 @@ func (txn *Txn[T]) delete(root *header[T], key []byte) (oldValue T, hadOld bool, hadOld = true // Mark the watch channel of the target for closing. - if leaf.watch != nil { - txn.watches[leaf.watch] = struct{}{} + if !target.isLeaf() { + txn.addWatch(leaf.watch) } if target == root { switch { case root.isLeaf() || root.size() == 0: - if root.watch != nil { - txn.watches[root.watch] = struct{}{} - } + txn.addWatch(root.watch) // Root is a leaf or node without children newRoot = nil case root.size() == 1: - if root.watch != nil { - txn.watches[root.watch] = struct{}{} - } + txn.addWatch(root.watch) // Root is a non-leaf node with single child. We can replace // the root with the child. child := root.children()[0] @@ -501,9 +528,7 @@ func (txn *Txn[T]) delete(root *header[T], key []byte) (oldValue T, hadOld bool, if this.node.size() == 1 { // The target node is not a leaf node and has only a single // child. Shift the child up. - if this.node.watch != nil { - txn.watches[this.node.watch] = struct{}{} - } + txn.addWatch(this.node.watch) child := this.node.children()[0] childClone := child.clone(false) childClone.watch = child.watch @@ -520,9 +545,7 @@ func (txn *Txn[T]) delete(root *header[T], key []byte) (oldValue T, hadOld bool, } else { // The target node is a leaf node or a non-leaf node without any // children. We can just drop it from the parent. - if this.node.watch != nil { - txn.watches[this.node.watch] = struct{}{} - } + txn.addWatch(this.node.watch) parent.node = txn.removeChild(parent.node, this.index) } @@ -575,9 +598,7 @@ func (txn *Txn[T]) removeChild(parent *header[T], index int) (newParent *header[ remainingIndex = 1 } - if parent.watch != nil { - txn.watches[parent.watch] = struct{}{} - } + txn.addWatch(parent.watch) child := parent.node4().children[remainingIndex] // Clone for prefix adjustment, but leave watch alone. @@ -591,25 +612,25 @@ func (txn *Txn[T]) removeChild(parent *header[T], index int) (newParent *header[ case parent.kind() == nodeKind256 && size <= 49: demoted := (&node48[T]{header: *parent}).self() if parent.watch != nil { - demoted.watch = make(chan struct{}) + demoted.watch = newWatchState() } demoted.setKind(nodeKind48) demoted.setSize(size - 1) demoted.setTxnID(txn.txnID) - n48 := demoted.node48() - n48.leaf = parent.getLeaf() - children := n48.children[:0] - for k, n := range parent.node256().children[:] { - if k != index && n != nil { - n48.index[k] = uint8(len(children) + 1) - children = append(children, n) - } + n48 := demoted.node48() + n48.leaf = parent.getLeaf() + children := n48.children[:0] + for k, n := range parent.node256().children[:] { + if k != index && n != nil { + n48.index[k] = uint8(len(children) + 1) + children = append(children, n) } + } newParent = demoted case parent.kind() == nodeKind48 && size <= 17: demoted := (&node16[T]{header: *parent}).self() if parent.watch != nil { - demoted.watch = make(chan struct{}) + demoted.watch = newWatchState() } demoted.setKind(nodeKind16) demoted.setSize(size - 1) @@ -628,7 +649,7 @@ func (txn *Txn[T]) removeChild(parent *header[T], index int) (newParent *header[ case parent.kind() == nodeKind16 && size <= 5: demoted := (&node4[T]{header: *parent}).self() if parent.watch != nil { - demoted.watch = make(chan struct{}) + demoted.watch = newWatchState() } demoted.setKind(nodeKind4) demoted.setSize(size - 1) @@ -650,9 +671,7 @@ func (txn *Txn[T]) removeChild(parent *header[T], index int) (newParent *header[ newParent.remove(index) return newParent } - if parent.watch != nil { - txn.watches[parent.watch] = struct{}{} - } + txn.addWatch(parent.watch) return newParent } @@ -660,7 +679,7 @@ var runValidation = os.Getenv("STATEDB_VALIDATE") != "" // validateTree checks that the resulting tree is well-formed and panics // if it is not. -func validateTree[T any](node *header[T], parents []*header[T], watches map[chan struct{}]struct{}, maxTxnID uint64) { +func validateTree[T any](node *header[T], parents []*header[T], watches []*watchState, maxTxnID uint64) { if !runValidation { return } @@ -690,8 +709,8 @@ func validateTree[T any](node *header[T], parents []*header[T], watches map[chan // If a leaf's watch channel is to be closed then parent's should be // marked closed too. The case where node is a leaf is handled below. if !node.isLeaf() { - if _, found := watches[leaf.watch]; found { - _, found := watches[node.watch] + if found := slices.Contains(watches, leaf.watch); found { + found := slices.Contains(watches, node.watch) assert(found, "node's watch channel not marked for closing when leaf is") } } @@ -725,14 +744,12 @@ func validateTree[T any](node *header[T], parents []*header[T], watches map[chan // Nodes that have a watch channel that is to be closed must // also have all their parent's watch channels to be closed. - if _, found := watches[node.watch]; found { - select { - case <-node.watch: + if found := slices.Contains(watches, node.watch); found { + if node.watch.isClosed() { panic("node's watch channel marked for closing but is already closed!") - default: } for i, p := range parents { - _, found := watches[p.watch] + found := slices.Contains(watches, p.watch) if !found { p.printTree(0) panic(fmt.Sprintf("parent %p (%d) watch channel (%p) not marked for closing (child %p, watch %p)", p, i, p.watch, node, node.watch)) @@ -743,17 +760,13 @@ func validateTree[T any](node *header[T], parents []*header[T], watches map[chan // If a node's watch channel is closed then all the parents must be // closed as well. - select { - case <-node.watch: + if node.watch.isClosed() { for _, p := range parents { - select { - case <-p.watch: - default: + if !p.watch.isClosed() { p.printTree(0) panic(fmt.Sprintf("parent watch channel (%p) not marked for closing (child %p)", p.watch, node.watch)) } } - default: } parents = append(parents, node) @@ -801,8 +814,8 @@ func validateRemovedWatches[T any](oldRoot *header[T], newRoot *header[T]) { } } - var collectWatches func(depth int, watches map[<-chan struct{}]func() string, node *header[T]) - collectWatches = func(depth int, watches map[<-chan struct{}]func() string, node *header[T]) { + var collectWatches func(depth int, watches map[*watchState]func() string, node *header[T]) + collectWatches = func(depth int, watches map[*watchState]func() string, node *header[T]) { if node == nil { return } @@ -820,9 +833,9 @@ func validateRemovedWatches[T any](oldRoot *header[T], newRoot *header[T]) { } } - oldWatches := map[<-chan struct{}]func() string{} + oldWatches := map[*watchState]func() string{} collectWatches(0, oldWatches, oldRoot) - newWatches := map[<-chan struct{}]func() string{} + newWatches := map[*watchState]func() string{} collectWatches(0, newWatches, newRoot) // Check that any nodes that kept the old watch channel have exactly @@ -845,9 +858,7 @@ func validateRemovedWatches[T any](oldRoot *header[T], newRoot *header[T]) { } for watch, desc := range oldWatches { - select { - case <-watch: - default: + if !watch.isClosed() { oldRoot.printTree(0) fmt.Println("---") newRoot.printTree(0) diff --git a/part/watch.go b/part/watch.go new file mode 100644 index 00000000..9e4a92a3 --- /dev/null +++ b/part/watch.go @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright Authors of Cilium + +package part + +import "sync/atomic" + +// watchState is the identity of a logical watch. The channel itself is +// allocated lazily when channel() is first called. A watchState may be shared +// by multiple physical nodes when an ART rewrite only changes a compressed +// prefix and leaves the watched subtree unchanged. +type watchState struct { + value atomic.Pointer[watchChannel] +} + +type watchChannel struct { + ch chan struct{} +} + +var closedWatchChannel = func() *watchChannel { + w := &watchChannel{ch: make(chan struct{})} + close(w.ch) + return w +}() + +func newWatchState() *watchState { + return &watchState{} +} + +// channel returns the stable channel for this watch. It is safe to race with +// close: either the newly installed channel is closed by close, or this returns +// the shared already-closed channel. +func (w *watchState) channel() <-chan struct{} { + if w == nil { + return nil + } + for { + if value := w.value.Load(); value != nil { + return value.ch + } + + candidate := &watchChannel{ch: make(chan struct{})} + if w.value.CompareAndSwap(nil, candidate) { + return candidate.ch + } + } +} + +// close closes the watch, including when it races with the first call to +// channel. It is idempotent so a shared watchState is safe to encounter more +// than once while rebuilding a tree. +func (w *watchState) close() { + if w == nil { + return + } + old := w.value.Swap(closedWatchChannel) + if old != nil && old != closedWatchChannel { + close(old.ch) + } +} + +func (w *watchState) isClosed() bool { + return w != nil && w.value.Load() == closedWatchChannel +} diff --git a/part/watch_test.go b/part/watch_test.go new file mode 100644 index 00000000..58c59622 --- /dev/null +++ b/part/watch_test.go @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright Authors of Cilium + +package part + +import ( + "sync" + "testing" +) + +func TestWatchStateCloseBeforeChannel(t *testing.T) { + w := newWatchState() + w.close() + + select { + case <-w.channel(): + default: + t.Fatal("watch created after close is open") + } +} + +func TestWatchStateConcurrentChannelAndClose(t *testing.T) { + for range 100 { + w := newWatchState() + channels := make([]<-chan struct{}, 32) + start := make(chan struct{}) + var wg sync.WaitGroup + for i := range channels { + wg.Add(1) + go func() { + defer wg.Done() + <-start + if i%4 == 0 { + w.close() + } + channels[i] = w.channel() + }() + } + close(start) + wg.Wait() + w.close() + + for _, ch := range channels { + select { + case <-ch: + default: + t.Fatal("watch remained open after close") + } + } + } +} + +func TestNoWatchOperationsDoNotCreateChannels(t *testing.T) { + tree := New[int]() + _, _, tree = tree.Insert([]byte("ab"), 1) + _, _, tree = tree.Insert([]byte("ac"), 2) + + if _, ok := tree.Get([]byte("ab")); !ok { + t.Fatal("Get did not find key") + } + iter := tree.Prefix([]byte("a")) + if _, _, ok := iter.Next(); !ok { + t.Fatal("Prefix did not find key") + } + + if tree.rootWatch.value.Load() != nil { + t.Fatal("root watch channel created by a no-watch operation") + } + var check func(*header[int]) + check = func(n *header[int]) { + if n.watch != nil && n.watch.value.Load() != nil { + t.Fatalf("node %p has a channel after no-watch operations", n) + } + if leaf := n.getLeaf(); leaf != nil && !n.isLeaf() && leaf.watch.value.Load() != nil { + t.Fatalf("leaf %p has a channel after no-watch operations", leaf) + } + for _, child := range n.children() { + if child != nil { + check(child) + } + } + } + check(tree.root) +} + +func TestLateWatchOnInvalidatedSnapshot(t *testing.T) { + tree := New[int]() + _, _, tree = tree.Insert([]byte("key"), 1) + oldTree := tree + + txn := tree.Txn() + txn.Insert([]byte("key"), 2) + _ = txn.Commit() + txn.Notify() + + _, watch, ok := oldTree.GetWatch([]byte("key")) + if !ok { + t.Fatal("old snapshot lost key") + } + select { + case <-watch: + default: + t.Fatal("watch requested after notification is open") + } +} + +func TestLateWatchAfterWatchPreservingRewrite(t *testing.T) { + tree := New[int]() + _, _, tree = tree.Insert([]byte("ab"), 1) + oldTree := tree + + // Splitting the leaf changes its compressed prefix but preserves its + // logical watch identity. + _, _, tree = tree.Insert([]byte("abc"), 2) + _, _, tree = tree.Insert([]byte("ab"), 3) + + _, watch, ok := oldTree.GetWatch([]byte("ab")) + if !ok { + t.Fatal("old snapshot lost key") + } + select { + case <-watch: + default: + t.Fatal("watch on rewritten old snapshot is open after key changed") + } +} diff --git a/part_index.go b/part_index.go index 917ffead..aeb73b63 100644 --- a/part_index.go +++ b/part_index.go @@ -118,13 +118,17 @@ func (r *partIndex) list(key index.Key) (tableIndexIterator, <-chan struct{}) { return partList(r.unique, &r.tree, key) } +func (r *partIndex) listNoWatch(key index.Key) tableIndexIterator { + return partListNoWatch(r.unique, &r.tree, key) +} + var emptyTableIndexIterator = &singletonTableIndexIterator{} func partList(unique bool, tree part.Ops[object], key index.Key) (tableIndexIterator, <-chan struct{}) { if unique { // Unique index means that there can be only a single matching object. // Doing a Get() is more efficient than constructing an iterator. - obj, watch, ok := tree.Get(key) + obj, watch, ok := tree.GetWatch(key) if ok { return &singletonTableIndexIterator{key, obj}, watch } @@ -137,10 +141,24 @@ func partList(unique bool, tree part.Ops[object], key index.Key) (tableIndexIter // form , and thus the // iteration will continue until key length mismatches, e.g. we hit a // longer key sharing the same prefix. - iter, watch := tree.Prefix(key) + iter, watch := tree.PrefixWatch(key) return newNonUniquePartIterator(iter, false, key), watch } +func partListNoWatch(unique bool, tree part.Ops[object], key index.Key) tableIndexIterator { + if unique { + obj, ok := tree.Get(key) + if ok { + return &singletonTableIndexIterator{key, obj} + } + return emptyTableIndexIterator + } + + key = encodeNonUniqueBytes(key) + iter := tree.Prefix(key) + return newNonUniquePartIterator(iter, false, key) +} + // rootWatch implements tableIndex. func (r *partIndex) rootWatch() <-chan struct{} { return r.tree.RootWatch() @@ -159,17 +177,21 @@ func (r *partIndex) get(ikey index.Key) (iobj object, watch <-chan struct{}, fou return partGet(r.unique, &r.tree, ikey) } +func (r *partIndex) getNoWatch(ikey index.Key) (iobj object, found bool) { + return partGetNoWatch(r.unique, &r.tree, ikey) +} + func partGet(unique bool, tree part.Ops[object], ikey index.Key) (iobj object, watch <-chan struct{}, found bool) { searchKey := ikey if unique { // On a unique index we can do a direct get rather than a prefix search. - return tree.Get(searchKey) + return tree.GetWatch(searchKey) } searchKey = encodeNonUniqueBytes(searchKey) // For a non-unique index we need to do a prefix search. - iter, watch := tree.Prefix(searchKey) + iter, watch := tree.PrefixWatch(searchKey) for { var key []byte key, iobj, found = iter.Next() @@ -185,6 +207,23 @@ func partGet(unique bool, tree part.Ops[object], ikey index.Key) (iobj object, w return iobj, watch, found } +func partGetNoWatch(unique bool, tree part.Ops[object], ikey index.Key) (iobj object, found bool) { + searchKey := ikey + if unique { + return tree.Get(searchKey) + } + + searchKey = encodeNonUniqueBytes(searchKey) + iter := tree.Prefix(searchKey) + for { + var key []byte + key, iobj, found = iter.Next() + if !found || nonUniqueKey(key).secondaryLen() == len(searchKey) { + return + } + } +} + // len implements tableIndex. func (r *partIndex) len() int { return r.tree.Len() @@ -194,27 +233,50 @@ func (r *partIndex) all() (tableIndexIterator, <-chan struct{}) { return &r.tree, r.rootWatch() } +func (r *partIndex) allNoWatch() tableIndexIterator { + return &r.tree +} + // prefix implements tableIndex. func (r *partIndex) prefix(ikey index.Key) (tableIndexIterator, <-chan struct{}) { return partPrefix(r.unique, &r.tree, ikey) } +func (r *partIndex) prefixNoWatch(ikey index.Key) tableIndexIterator { + return partPrefixNoWatch(r.unique, &r.tree, ikey) +} + func partPrefix(unique bool, tree part.Ops[object], key index.Key) (tableIndexIterator, <-chan struct{}) { if !unique { key = encodeNonUniqueBytes(key) } - iter, watch := tree.Prefix(key) + iter, watch := tree.PrefixWatch(key) if unique { return iter, watch } return newNonUniquePartIterator(iter, true, key), watch } +func partPrefixNoWatch(unique bool, tree part.Ops[object], key index.Key) tableIndexIterator { + if !unique { + key = encodeNonUniqueBytes(key) + } + iter := tree.Prefix(key) + if unique { + return &iter + } + return newNonUniquePartIterator(iter, true, key) +} + // lowerBound implements tableIndexTxn. func (r *partIndex) lowerBound(ikey index.Key) (tableIndexIterator, <-chan struct{}) { return partLowerBound(r.unique, &r.tree, ikey), r.rootWatch() } +func (r *partIndex) lowerBoundNoWatch(ikey index.Key) tableIndexIterator { + return partLowerBound(r.unique, &r.tree, ikey) +} + // lowerBoundNext implements tableIndexTxn. func (r *partIndex) lowerBoundNext(key index.Key) (func() ([]byte, object, bool), <-chan struct{}) { if !r.unique { @@ -227,6 +289,17 @@ func (r *partIndex) lowerBoundNext(key index.Key) (func() ([]byte, object, bool) return newNonUniqueLowerBoundPartIterator(iter, key).Next, r.rootWatch() } +func (r *partIndex) lowerBoundNextNoWatch(key index.Key) func() ([]byte, object, bool) { + if !r.unique { + key = encodeNonUniqueBytes(key) + } + iter := r.tree.LowerBound(key) + if r.unique { + return iter.Next + } + return newNonUniqueLowerBoundPartIterator(iter, key).Next +} + func partLowerBound(unique bool, tree part.Ops[object], key index.Key) tableIndexIterator { if !unique { key = encodeNonUniqueBytes(key) @@ -259,18 +332,33 @@ func (r *partIndexTxn) all() (tableIndexIterator, <-chan struct{}) { return &snapshot, r.rootWatch() } +func (r *partIndexTxn) allNoWatch() tableIndexIterator { + snapshot := r.tx.Clone() + return &snapshot +} + // list implements tableIndexTxn. func (r *partIndexTxn) list(ikey index.Key) (tableIndexIterator, <-chan struct{}) { snapshot := r.tx.Clone() return partList(r.unique, &snapshot, ikey) } +func (r *partIndexTxn) listNoWatch(ikey index.Key) tableIndexIterator { + snapshot := r.tx.Clone() + return partListNoWatch(r.unique, &snapshot, ikey) +} + // lowerBound implements tableIndexTxn. func (r *partIndexTxn) lowerBound(ikey index.Key) (tableIndexIterator, <-chan struct{}) { snapshot := r.tx.Clone() return partLowerBound(r.unique, &snapshot, ikey), r.rootWatch() } +func (r *partIndexTxn) lowerBoundNoWatch(ikey index.Key) tableIndexIterator { + snapshot := r.tx.Clone() + return partLowerBound(r.unique, &snapshot, ikey) +} + // lowerBoundNext implements tableIndexTxn. func (r *partIndexTxn) lowerBoundNext(key index.Key) (func() ([]byte, object, bool), <-chan struct{}) { if !r.unique { @@ -284,6 +372,18 @@ func (r *partIndexTxn) lowerBoundNext(key index.Key) (func() ([]byte, object, bo return newNonUniqueLowerBoundPartIterator(iter, key).Next, r.rootWatch() } +func (r *partIndexTxn) lowerBoundNextNoWatch(key index.Key) func() ([]byte, object, bool) { + if !r.unique { + key = encodeNonUniqueBytes(key) + } + snapshot := r.tx.Clone() + iter := snapshot.LowerBound(key) + if r.unique { + return iter.Next + } + return newNonUniqueLowerBoundPartIterator(iter, key).Next +} + // rootWatch implements tableIndexTxn. func (r *partIndexTxn) rootWatch() <-chan struct{} { return r.tx.RootWatch() @@ -310,11 +410,19 @@ func (r *partIndexTxn) get(key index.Key) (iobj object, watch <-chan struct{}, o return partGet(r.unique, r.tx, key) } +func (r *partIndexTxn) getNoWatch(key index.Key) (iobj object, ok bool) { + return partGetNoWatch(r.unique, r.tx, key) +} + // insert implements tableIndexTxn. func (r *partIndexTxn) insert(key index.Key, obj object) (old object, hadOld bool, watch <-chan struct{}) { return r.tx.InsertWatch(key, obj) } +func (r *partIndexTxn) insertNoWatch(key index.Key, obj object) (old object, hadOld bool) { + return r.tx.Insert(key, obj) +} + // len implements tableIndexTxn. func (r *partIndexTxn) len() int { return r.tx.Len() @@ -325,6 +433,10 @@ func (r *partIndexTxn) modify(key index.Key, obj object, mod func(old, new objec return r.tx.ModifyWatch(key, obj, mod) } +func (r *partIndexTxn) modifyNoWatch(key index.Key, obj object, mod func(old, new object) object) (old object, newObj object, hadOld bool) { + return r.tx.Modify(key, obj, mod) +} + // notify implements tableIndexTxn. func (r *partIndexTxn) notify() { if r.tx != nil { @@ -339,6 +451,11 @@ func (r *partIndexTxn) prefix(ikey index.Key) (tableIndexIterator, <-chan struct return partPrefix(r.unique, &snapshot, ikey) } +func (r *partIndexTxn) prefixNoWatch(ikey index.Key) tableIndexIterator { + snapshot := r.tx.Clone() + return partPrefixNoWatch(r.unique, &snapshot, ikey) +} + func (r *partIndexTxn) objectToKey(obj object) index.Key { return r.objectToKeys(obj).First() } diff --git a/read_txn.go b/read_txn.go index 4a098b5a..e35912c7 100644 --- a/read_txn.go +++ b/read_txn.go @@ -88,7 +88,7 @@ func marshalJSON(data any) (out []byte) { func writeTableAsJSON(buf *bufio.Writer, txn ReadTxn, table *tableEntry) (err error) { indexTxn := txn.mustIndexReadTxn(table.meta, PrimaryIndexPos) - iter, _ := indexTxn.all() + iter := indexTxn.allNoWatch() writeString := func(s string) { if err != nil { diff --git a/reconciler/multi_test.go b/reconciler/multi_test.go index 8274dda7..863bd43c 100644 --- a/reconciler/multi_test.go +++ b/reconciler/multi_test.go @@ -460,20 +460,28 @@ func TestMultipleReconcilersPerModuleMetrics(t *testing.T) { <-watch } - require.NotNil(t, metrics.ReconciliationCountVar.Get("test/left")) - require.NotNil(t, metrics.ReconciliationCountVar.Get("test/right")) - require.NotNil(t, metrics.ReconciliationCurrentErrorsVar.Get("test/left")) - require.NotNil(t, metrics.ReconciliationCurrentErrorsVar.Get("test/right")) - assert.NotEqual(t, "0", metrics.ReconciliationCountVar.Get("test/left").String()) - assert.NotEqual(t, "0", metrics.ReconciliationCountVar.Get("test/right").String()) - assert.Equal(t, "0", metrics.ReconciliationCurrentErrorsVar.Get("test/left").String()) - assert.Equal(t, "0", metrics.ReconciliationCurrentErrorsVar.Get("test/right").String()) + // Status updates are committed before the reconciliation round publishes + // its metrics, so wait for the expected values to appear. + require.Eventually(t, func() bool { + leftCount := metrics.ReconciliationCountVar.Get("test/left") + rightCount := metrics.ReconciliationCountVar.Get("test/right") + leftErrors := metrics.ReconciliationCurrentErrorsVar.Get("test/left") + rightErrors := metrics.ReconciliationCurrentErrorsVar.Get("test/right") + return leftCount != nil && leftCount.String() != "0" && + rightCount != nil && rightCount.String() != "0" && + leftErrors != nil && leftErrors.String() == "0" && + rightErrors != nil && rightErrors.String() == "0" + }, 5*time.Second, time.Millisecond) assert.Nil(t, metrics.ReconciliationCountVar.Get("test")) - assert.NotNil(t, health.GetChild("job-reconcile-left")) - assert.NotNil(t, health.GetChild("job-refresh-left")) - assert.NotNil(t, health.GetChild("job-reconcile-right")) - assert.NotNil(t, health.GetChild("job-refresh-right")) + // The refresh jobs start independently from reconciliation, so reaching the + // done statuses does not guarantee that their health nodes exist yet. + require.Eventually(t, func() bool { + return health.GetChild("job-reconcile-left") != nil && + health.GetChild("job-refresh-left") != nil && + health.GetChild("job-reconcile-right") != nil && + health.GetChild("job-refresh-right") != nil + }, 5*time.Second, time.Millisecond) require.NoError(t, hive.Stop(log, context.TODO()), "Stop") } diff --git a/table.go b/table.go index ea97a17f..dab0290f 100644 --- a/table.go +++ b/table.go @@ -422,7 +422,12 @@ func (t *genTable[Obj]) numDeletedObjects(txn ReadTxn) int { } func (t *genTable[Obj]) Get(txn ReadTxn, q Query[Obj]) (obj Obj, revision uint64, ok bool) { - obj, revision, _, ok = t.GetWatch(txn, q) + index := txn.root()[t.pos].indexes[t.indexPos(q.index)] + iobj, ok := index.getNoWatch(q.key) + if ok { + obj = iobj.data.(Obj) + revision = iobj.revision + } return } @@ -437,8 +442,8 @@ func (t *genTable[Obj]) GetWatch(txn ReadTxn, q Query[Obj]) (obj Obj, revision u } func (t *genTable[Obj]) LowerBound(txn ReadTxn, q Query[Obj]) iter.Seq2[Obj, Revision] { - iter, _ := t.LowerBoundWatch(txn, q) - return iter + indexTxn := txn.mustIndexReadTxn(t, t.indexPos(q.index)) + return objSeq[Obj](indexTxn.lowerBoundNoWatch(q.key)) } func (t *genTable[Obj]) LowerBoundWatch(txn ReadTxn, q Query[Obj]) (iter.Seq2[Obj, Revision], <-chan struct{}) { @@ -448,8 +453,8 @@ func (t *genTable[Obj]) LowerBoundWatch(txn ReadTxn, q Query[Obj]) (iter.Seq2[Ob } func (t *genTable[Obj]) Prefix(txn ReadTxn, q Query[Obj]) iter.Seq2[Obj, Revision] { - iter, _ := t.PrefixWatch(txn, q) - return iter + indexTxn := txn.mustIndexReadTxn(t, t.indexPos(q.index)) + return objSeq[Obj](indexTxn.prefixNoWatch(q.key)) } func (t *genTable[Obj]) PrefixWatch(txn ReadTxn, q Query[Obj]) (iter.Seq2[Obj, Revision], <-chan struct{}) { @@ -459,8 +464,13 @@ func (t *genTable[Obj]) PrefixWatch(txn ReadTxn, q Query[Obj]) (iter.Seq2[Obj, R } func (t *genTable[Obj]) All(txn ReadTxn) iter.Seq2[Obj, Revision] { - iter, _ := t.AllWatch(txn) - return iter + indexTxn := txn.mustIndexReadTxn(t, PrimaryIndexPos) + iter := indexTxn.allNoWatch() + return func(yield func(Obj, Revision) bool) { + iter.All(func(_ []byte, obj object) bool { + return yield(obj.data.(Obj), obj.revision) + }) + } } func (t *genTable[Obj]) AllWatch(txn ReadTxn) (iter.Seq2[Obj, Revision], <-chan struct{}) { @@ -474,8 +484,8 @@ func (t *genTable[Obj]) AllWatch(txn ReadTxn) (iter.Seq2[Obj, Revision], <-chan } func (t *genTable[Obj]) List(txn ReadTxn, q Query[Obj]) iter.Seq2[Obj, Revision] { - iter, _ := t.ListWatch(txn, q) - return iter + indexTxn := txn.mustIndexReadTxn(t, t.indexPos(q.index)) + return objSeq[Obj](indexTxn.listNoWatch(q.key)) } func (t *genTable[Obj]) ListWatch(txn ReadTxn, q Query[Obj]) (iter.Seq2[Obj, Revision], <-chan struct{}) { @@ -485,13 +495,17 @@ func (t *genTable[Obj]) ListWatch(txn ReadTxn, q Query[Obj]) (iter.Seq2[Obj, Rev } func (t *genTable[Obj]) Insert(txn WriteTxn, obj Obj) (oldObj Obj, hadOld bool, err error) { - oldObj, hadOld, _, err = t.InsertWatch(txn, obj) + var old object + old, hadOld, _, err = txn.unwrap().insert(t, Revision(0), obj, false) + if hadOld { + oldObj = old.data.(Obj) + } return } func (t *genTable[Obj]) InsertWatch(txn WriteTxn, obj Obj) (oldObj Obj, hadOld bool, watch <-chan struct{}, err error) { var old object - old, hadOld, watch, err = txn.unwrap().insert(t, Revision(0), obj) + old, hadOld, watch, err = txn.unwrap().insert(t, Revision(0), obj, true) if hadOld { oldObj = old.data.(Obj) } @@ -504,7 +518,7 @@ func (t *genTable[Obj]) Modify(txn WriteTxn, obj Obj, merge func(old, new Obj) O return new } var old object - old, hadOld, _, err = txn.unwrap().modify(t, Revision(0), obj, mergeObjects) + old, hadOld, _, err = txn.unwrap().modify(t, Revision(0), obj, mergeObjects, false) if hadOld { oldObj = old.data.(Obj) } @@ -513,7 +527,7 @@ func (t *genTable[Obj]) Modify(txn WriteTxn, obj Obj, merge func(old, new Obj) O func (t *genTable[Obj]) CompareAndSwap(txn WriteTxn, rev Revision, obj Obj) (oldObj Obj, hadOld bool, err error) { var old object - old, hadOld, _, err = txn.unwrap().insert(t, rev, obj) + old, hadOld, _, err = txn.unwrap().insert(t, rev, obj, false) if hadOld { oldObj = old.data.(Obj) } diff --git a/types.go b/types.go index 0ec8414b..466c0767 100644 --- a/types.go +++ b/types.go @@ -410,11 +410,17 @@ type tableIndexIterator interface { type tableIndexReader interface { len() int get(key index.Key) (object, <-chan struct{}, bool) + getNoWatch(key index.Key) (object, bool) prefix(key index.Key) (tableIndexIterator, <-chan struct{}) + prefixNoWatch(key index.Key) tableIndexIterator lowerBound(key index.Key) (tableIndexIterator, <-chan struct{}) + lowerBoundNoWatch(key index.Key) tableIndexIterator lowerBoundNext(key index.Key) (func() ([]byte, object, bool), <-chan struct{}) + lowerBoundNextNoWatch(key index.Key) func() ([]byte, object, bool) list(key index.Key) (tableIndexIterator, <-chan struct{}) + listNoWatch(key index.Key) tableIndexIterator all() (tableIndexIterator, <-chan struct{}) + allNoWatch() tableIndexIterator rootWatch() <-chan struct{} objectToKey(obj object) index.Key } @@ -429,7 +435,9 @@ type tableIndexTxn interface { tableIndex insert(key index.Key, obj object) (old object, hadOld bool, watch <-chan struct{}) + insertNoWatch(key index.Key, obj object) (old object, hadOld bool) modify(key index.Key, obj object, mod func(old, new object) object) (old object, new object, hadOld bool, watch <-chan struct{}) + modifyNoWatch(key index.Key, obj object, mod func(old, new object) object) (old object, new object, hadOld bool) delete(key index.Key) (old object, hadOld bool) reindex(primaryKey index.Key, old object, new object) } diff --git a/write_txn.go b/write_txn.go index ea4beaca..5eee7193 100644 --- a/write_txn.go +++ b/write_txn.go @@ -112,11 +112,11 @@ func (txn *writeTxnState) mustIndexWriteTxn(meta TableMeta, indexPos int) tableI return indexTxn } -func (txn *writeTxnState) insert(meta TableMeta, guardRevision Revision, data any) (object, bool, <-chan struct{}, error) { - return txn.modify(meta, guardRevision, data, nil) +func (txn *writeTxnState) insert(meta TableMeta, guardRevision Revision, data any, withWatch bool) (object, bool, <-chan struct{}, error) { + return txn.modify(meta, guardRevision, data, nil, withWatch) } -func (txn *writeTxnState) modify(meta TableMeta, guardRevision Revision, newData any, merge func(old, new object) object) (object, bool, <-chan struct{}, error) { +func (txn *writeTxnState) modify(meta TableMeta, guardRevision Revision, newData any, merge func(old, new object) object, withWatch bool) (object, bool, <-chan struct{}, error) { if txn == nil { return object{}, false, nil, ErrTransactionClosed } @@ -142,11 +142,19 @@ func (txn *writeTxnState) modify(meta TableMeta, guardRevision Revision, newData watch <-chan struct{} ) if merge == nil { - oldObj, oldExists, watch = idIndexTxn.insert(idKey, obj) + if withWatch { + oldObj, oldExists, watch = idIndexTxn.insert(idKey, obj) + } else { + oldObj, oldExists = idIndexTxn.insertNoWatch(idKey, obj) + } } else { // Insert the object into the primary index. This returns the merged new // object which we'll then insert into the secondary indexes. - oldObj, obj, oldExists, watch = idIndexTxn.modify(idKey, obj, merge) + if withWatch { + oldObj, obj, oldExists, watch = idIndexTxn.modify(idKey, obj, merge) + } else { + oldObj, obj, oldExists = idIndexTxn.modifyNoWatch(idKey, obj, merge) + } } // Sanity check: is the same object being inserted back and thus the @@ -176,7 +184,7 @@ func (txn *writeTxnState) modify(meta TableMeta, guardRevision Revision, newData // Revert the change. We're assuming here that it's rarer for CompareAndSwap() to // fail and thus we're optimizing to have only one lookup in the common case // (versus doing a Get() and then Insert()). - idIndexTxn.insert(idKey, oldObj) + idIndexTxn.insertNoWatch(idKey, oldObj) table.revision = oldRevision return oldObj, true, watch, ErrRevisionNotEqual } @@ -188,12 +196,12 @@ func (txn *writeTxnState) modify(meta TableMeta, guardRevision Revision, newData binary.BigEndian.PutUint64(txn.revKey[:], oldObj.revision) revIndexTxn.delete(txn.revKey[:]) } - revIndexTxn.insert(index.Uint64(obj.revision), obj) + revIndexTxn.insertNoWatch(index.Uint64(obj.revision), obj) // If it's new, possibly remove an older deleted object with the same // primary key from the graveyard. if !oldExists { - if old, _, existed := txn.mustIndexReadTxn(meta, GraveyardIndexPos).get(idKey); existed { + if old, existed := txn.mustIndexReadTxn(meta, GraveyardIndexPos).getNoWatch(idKey); existed { txn.mustIndexWriteTxn(meta, GraveyardIndexPos).delete(idKey) binary.BigEndian.PutUint64(txn.revKey[:], old.revision) txn.mustIndexWriteTxn(meta, GraveyardRevisionIndexPos).delete(txn.revKey[:]) @@ -256,7 +264,7 @@ func (txn *writeTxnState) delete(meta TableMeta, guardRevision Revision, data an // revert the change. if guardRevision > 0 { if obj.revision != guardRevision { - idIndex.insert(idKey, obj) + idIndex.insertNoWatch(idKey, obj) return obj, true, ErrRevisionNotEqual } } @@ -277,10 +285,10 @@ func (txn *writeTxnState) delete(meta TableMeta, guardRevision Revision, data an if txn.hasDeleteTrackers(meta) { graveyardIndex := txn.mustIndexWriteTxn(meta, GraveyardIndexPos) obj.revision = revision - if _, existed, _ := graveyardIndex.insert(idKey, obj); existed { + if _, existed := graveyardIndex.insertNoWatch(idKey, obj); existed { panic("BUG: Double deletion! Deleted object already existed in graveyard") } - txn.mustIndexWriteTxn(meta, GraveyardRevisionIndexPos).insert(index.Uint64(revision), obj) + txn.mustIndexWriteTxn(meta, GraveyardRevisionIndexPos).insertNoWatch(index.Uint64(revision), obj) } return obj, true, nil