part: Add missing closing of root watch channel

The commit f8e6b9e changed the implementation not to always have
a zero-prefix node4 at the root. This change unfortunately missed marking
the replaced root's watch channel for closing. This fixes the issue and
adds a validation to make sure any watch channel that no longer is referenced
in the new tree is closed.

=== RUN   Test_replaceRoot
node4[61]: 61 -> 1 (L:0x4000041b00 W:0x4000014a80 true) (N:0x400012c500, W:0x4000014bd0 false)
 leaf[62]: 6162 -> 3 (L:0x4000041b30 W:0x4000014b60 false) (N:0x4000041b30, W:0x4000014b60 false)
---
leaf[6162]: 6162 -> 3 (L:0x4000041bf0 W:0x4000014b60 false) (N:0x4000041bf0, W:0x4000014b60 false)
--- FAIL: Test_replaceRoot (0.00s)
panic: dropped watch channel 0x4000014bd0 at depth 0 not closed [recovered]
        panic: dropped watch channel 0x4000014bd0 at depth 0 not closed

Fixes: f8e6b9e ("part: Remove the unnecessary empty root node")
Signed-off-by: Jussi Maki <jussi@isovalent.com>

Author: joamaki

part/txn.go

@@ -194,6 +194,9 @@ func (txn *Txn[T]) Notify() {
 		close(txn.rootWatch)
 		txn.rootWatch = nil
 	}
+	if !txn.opts.rootOnlyWatch() {
+		validateRemovedWatches(txn.oldRoot, txn.root)
+	}
 }
 
 // PrintTree to the standard output. For debugging.
@@ -426,9 +429,15 @@ func (txn *Txn[T]) delete(root *header[T], key []byte) (oldValue T, hadOld bool,
 	if target == root {
 		switch {
 		case root.isLeaf() || root.size() == 0:
+			if root.watch != nil {
+				txn.watches[root.watch] = struct{}{}
+			}
 			// Root is a leaf or node without children
 			newRoot = nil
 		case root.size() == 1:
+			if root.watch != nil {
+				txn.watches[root.watch] = struct{}{}
+			}
 			// Root is a non-leaf node with single child. We can replace
 			// the root with the child.
 			child := root.children()[0]
@@ -471,6 +480,9 @@ 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{}{}
+		}
 		parent.node = txn.removeChild(parent.node, this.index)
 	}
 	index--
@@ -675,3 +687,47 @@ func validateTree[T any](node *header[T], parents []*header[T], watches map[chan
 		}
 	}
 }
+
+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]) {
+		if node == nil {
+			return
+		}
+		if node.watch == nil {
+			panic("nil watch channel")
+		}
+		watches[node.watch] = depth
+		if leaf := node.getLeaf(); leaf != nil && !node.isLeaf() {
+			watches[leaf.watch] = depth
+		}
+		for _, child := range node.children() {
+			if child != nil {
+				collectWatches(depth+1, watches, child)
+			}
+		}
+	}
+	oldWatches := map[<-chan struct{}]int{}
+	collectWatches(0, oldWatches, oldRoot)
+	newWatches := map[<-chan struct{}]int{}
+	collectWatches(0, newWatches, newRoot)
+
+	// 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 {
+		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))
+		}
+	}
+}