part: Store header pointer instead of uintptr in nodeMutated

If we store uintptr then there's a chance that a pointer there
is garbage collected and reallocated in which case the assumption
that a node present in nodeMutated has had its watch channel
replaced and marked for closing no longer holds. This lead
to the validation failing with a watch channel retained on a node
which was mutated (e.g. leaf or children changed).

Signed-off-by: Jussi Maki <jussi.maki@isovalent.com>

Author: joamaki

part/cache.go

@@ -9,36 +9,56 @@ import (
 
 const nodeMutatedSize = 256 // must be power-of-two
 
-type nodeMutated struct {
-	ptrs [nodeMutatedSize]uintptr
+// nodeMutated is a probabilistic check for seeing if a node has
+// been cloned within a transaction and thus can be modified in-place
+// since it has not been seen outside. This significantly speeds up
+// writes within a single write transaction as inner nodes no longer
+// need to be cloned on every change, effectively making the immutable
+// radix tree perform as if it's a mutable one.
+//
+// Earlier versions of StateDB just used a map[*header[T]]struct{}, but
+// that was fairly costly and experiments showed that it's enough to most
+// of the time avoid the clone to perform well.
+//
+// The value for [nodeMutatedSize] is a balance between making Txn()
+// not too costly (due to e.g. clear()) and between giving a high likelyhood
+// that we mutate nodes in-place.
+type nodeMutated[T any] struct {
+	ptrs [nodeMutatedSize]*header[T]
 	used bool
 }
 
-func nodeMutatedSet[T any](nm *nodeMutated, ptr *header[T]) {
+func (nm *nodeMutated[T]) set(n *header[T]) {
 	if nm == nil {
 		return
 	}
-	ptrInt := uintptr(unsafe.Pointer(ptr))
-	nm.ptrs[slot(ptrInt)] = ptrInt
+	ptrInt := uintptr(unsafe.Pointer(n))
+	nm.ptrs[slot(ptrInt)] = n
 	nm.used = true
 }
 
-func nodeMutatedExists[T any](nm *nodeMutated, ptr *header[T]) bool {
+func (nm *nodeMutated[T]) exists(n *header[T]) bool {
 	if nm == nil {
 		return false
 	}
-	ptrInt := uintptr(unsafe.Pointer(ptr))
-	return nm.ptrs[slot(ptrInt)] == ptrInt
+	ptrInt := uintptr(unsafe.Pointer(n))
+	return nm.ptrs[slot(ptrInt)] == n
 }
 
+// slot returns the index in the [ptrs] array for a given pointer.
+// The Go spec allows objects to be moved so it may be that the same
+// instance of an object is assigned to a different memory location in
+// which case we'd no longer report that node as being in the cache.
+// This is fine though as we do compare the actual *header[T] pointers
+// and this is probabilistic anyway as this is a fixed size cache.
 func slot(p uintptr) int {
 	p >>= 4 // ignore low order bits
 	// use some relevant bits from the pointer
 	slot := uint8(p) ^ uint8(p>>8) ^ uint8(p>>16)
 	return int(slot & (nodeMutatedSize - 1))
 }
 
-func (nm *nodeMutated) clear() {
+func (nm *nodeMutated[T]) clear() {
 	if nm == nil {
 		return
 	}

part/tree.go

@@ -51,7 +51,7 @@ func newTxn[T any](o options) *Txn[T] {
 		watches: make(map[chan struct{}]struct{}),
 	}
 	if !o.noCache() {
-		txn.mutated = &nodeMutated{}
+		txn.mutated = &nodeMutated[T]{}
 		txn.deleteParentsCache = make([]deleteParent[T], 0, 32)
 	}
 	txn.opts = o

part/txn.go

@@ -26,7 +26,7 @@ type Txn[T any] struct {
 	// that we can keep mutating without cloning them again.
 	// It is cleared if the transaction is cloned or iterated
 	// upon.
-	mutated *nodeMutated
+	mutated *nodeMutated[T]
 
 	// watches contains the channels of cloned nodes that should be closed
 	// when transaction is committed.
@@ -213,14 +213,14 @@ func (txn *Txn[T]) PrintTree() {
 }
 
 func (txn *Txn[T]) cloneNode(n *header[T]) *header[T] {
-	if nodeMutatedExists(txn.mutated, n) {
+	if txn.mutated.exists(n) {
 		return n
 	}
 	if n.watch != nil {
 		txn.watches[n.watch] = struct{}{}
 	}
 	n = n.clone(!txn.opts.rootOnlyWatch())
-	nodeMutatedSet(txn.mutated, n)
+	txn.mutated.set(n)
 	return n
 }
 
@@ -265,7 +265,7 @@ func (txn *Txn[T]) modify(root *header[T], key []byte, mod func(T) T) (oldValue
 					txn.watches[this.watch] = struct{}{}
 				}
 				this = this.promote()
-				nodeMutatedSet(txn.mutated, this)
+				txn.mutated.set(this)
 			} else {
 				// Node is big enough, clone it so we can mutate it
 				this = txn.cloneNode(this)
@@ -371,7 +371,7 @@ func (txn *Txn[T]) modify(root *header[T], key []byte, mod func(T) T) (oldValue
 		newNode.setSize(2)
 	}
 	*thisp = newNode.self()
-	nodeMutatedSet(txn.mutated, newNode.self())
+	txn.mutated.set(newNode.self())
 
 	return
 }
@@ -602,7 +602,7 @@ func (txn *Txn[T]) removeChild(parent *header[T], index int) (newParent *header[
 	if parent.watch != nil {
 		txn.watches[parent.watch] = struct{}{}
 	}
-	nodeMutatedSet(txn.mutated, newParent)
+	txn.mutated.set(newParent)
 	return newParent
 }