part: Store header instead of uintptr in node mutated

There's no need to store an uintptr. In StateDB v0.5.x this
actually caused issues due to the compression change in which
watch channels from old tree were retained. The extended validation
shows no such for v0.4 since old watches are never retained when
a node is cloned, but let's be on safe side and change this here as well.

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

@@ -44,7 +44,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

@@ -20,7 +20,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.
@@ -188,14 +188,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() || n == txn.root)
-	nodeMutatedSet(txn.mutated, n)
+	txn.mutated.set(n)
 	return n
 }
 
@@ -232,7 +232,7 @@ func (txn *Txn[T]) modify(root *header[T], key []byte, mod func(T) T) (oldValue
 					txn.watches[this.watch] = struct{}{}
 				}
 				this = this.promote(!txn.opts.rootOnlyWatch() || this == root)
-				nodeMutatedSet(txn.mutated, this)
+				txn.mutated.set(this)
 			} else {
 				// Node is big enough, clone it so we can mutate it
 				this = txn.cloneNode(this)
@@ -490,41 +490,89 @@ func validateRemovedWatches[T any](oldRoot *header[T], newRoot *header[T]) {
 	if !runValidation {
 		return
 	}
-	var collectWatches func(depth int, watches map[<-chan struct{}]int, node *header[T])
-	collectWatches = func(depth int, watches map[<-chan struct{}]int, node *header[T]) {
+
+	// nodeStruct memoizes the structure for a node.
+	nodeStruct := map[*header[T]]string{}
+
+	// summarizeNodeStructure "lazily" summarizes the internal structure of a node,
+	// e.g. it's watch channel, size, leaf and children. The lazy construction speeds
+	// things up a lot as we only look at the structure in certain specific cases.
+	var summarizeNodeStructure func(node *header[T]) func() string
+	summarizeNodeStructure = func(node *header[T]) func() string {
+		if node == nil {
+			return func() string { return "" }
+		}
+		if s, found := nodeStruct[node]; found {
+			return func() string { return s }
+		}
+		return func() string {
+			var childS string
+			for _, child := range node.children() {
+				if child != nil {
+					childS += summarizeNodeStructure(child)()
+				}
+			}
+			var leafS string
+			if leaf := node.getLeaf(); leaf != nil && !node.isLeaf() {
+				leafS = summarizeNodeStructure(leaf.self())()
+			}
+			s := fmt.Sprintf("K:%d S:%d W:%p L:[%s] C:[%s]", node.kind(), node.size(), node.watch, leafS, childS)
+			nodeStruct[node] = s
+			return s
+		}
+	}
+
+	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]) {
 		if node == nil {
 			return
 		}
 		if node.watch == nil {
 			panic("nil watch channel")
 		}
-		watches[node.watch] = depth
+		watches[node.watch] = summarizeNodeStructure(node)
 		if leaf := node.getLeaf(); leaf != nil && !node.isLeaf() {
-			watches[leaf.watch] = depth
+			watches[leaf.watch] = summarizeNodeStructure(leaf.self())
 		}
 		for _, child := range node.children() {
 			if child != nil {
 				collectWatches(depth+1, watches, child)
 			}
 		}
 	}
-	oldWatches := map[<-chan struct{}]int{}
+
+	oldWatches := map[<-chan struct{}]func() string{}
 	collectWatches(0, oldWatches, oldRoot)
-	newWatches := map[<-chan struct{}]int{}
+	newWatches := map[<-chan struct{}]func() string{}
 	collectWatches(0, newWatches, newRoot)
 
+	// Check that any nodes that kept the old watch channel have exactly
+	// the same leaf and children structure.
+	for watch, oldDescFn := range oldWatches {
+		newDescFn, found := newWatches[watch]
+		if found {
+			oldDesc := oldDescFn()
+			newDesc := newDescFn()
+			if oldDesc != newDesc {
+				panic(fmt.Sprintf("node with retained watch channel has different structure:\nexpected: %s\n     got: %s", oldDesc, newDesc))
+			}
+		}
+
+	}
+
 	// Any nodes that are not part of the new tree must have their watch channels closed.
 	for watch := range newWatches {
 		delete(oldWatches, watch)
 	}
-	for watch, depth := range oldWatches {
+
+	for watch, desc := range oldWatches {
 		select {
 		case <-watch:
 		default:
 			oldRoot.printTree(0)
 			fmt.Println("---")
 			newRoot.printTree(0)
-			panic(fmt.Sprintf("dropped watch channel %p at depth %d not closed", watch, depth))
+			panic(fmt.Sprintf("dropped watch channel %p not closed %s", watch, desc()))
 		}
 	}
 }