part: Fix returned watch channel in search()

If there was a leaf node with a key that was shorter than
our search key we incorrectly returned that leaf node's watch
channel instead of the channel of the parent non-leaf node.

This caused the watch channel returned by GetWatch() to not
close when expected.

Signed-off-by: Jussi Maki <[email protected]>

Author: joamaki

part/node.go

@@ -474,13 +474,13 @@ func newNode4[T any]() *header[T] {
 
 func search[T any](root *header[T], key []byte) (value T, watch <-chan struct{}, ok bool) {
 	this := root
+	watch = root.watch
 	for {
-		watch = this.watch
-
-		// Consume the prefix
 		if !bytes.HasPrefix(key, this.prefix) {
 			return
 		}
+
+		// Consume the prefix
 		key = key[len(this.prefix):]
 
 		if len(key) == 0 {
@@ -492,6 +492,12 @@ func search[T any](root *header[T], key []byte) (value T, watch <-chan struct{},
 			return
 		}
 
+		// Prefix matched. Remember this as the closest watch channel as we traverse
+		// further.
+		if !this.isLeaf() && this.watch != nil {
+			watch = this.watch
+		}
+
 		this = this.find(key[0])
 		if this == nil {
 			return