Iteration across a node's children was slower than expected and creat…

…ing memory for the GC.

The solution was to have a node return its children as a []node. Since node256 is sparse the upper layers need to check for nil, but this improved the performance.

Signed-off-by: Derek Collison <derek@nats.io>

server/stree/leaf.go

@@ -37,10 +37,11 @@ func (n *leaf[T]) setSuffix(suffix []byte)                    { n.suffix = copyB
 func (n *leaf[T]) isFull() bool                               { return true }
 func (n *leaf[T]) matchParts(parts [][]byte) ([][]byte, bool) { return matchParts(parts, n.suffix) }
 func (n *leaf[T]) iter(f func(node) bool)                     {}
+func (n *leaf[T]) children() []node                           { return nil }
 func (n *leaf[T]) numChildren() uint16                        { return 0 }
 func (n *leaf[T]) path() []byte                               { return n.suffix }
 
-// Not applicable to leafs.
+// Not applicable to leafs and should not be called, so panic if we do.
 func (n *leaf[T]) setPrefix(pre []byte)    { panic("setPrefix called on leaf") }
 func (n *leaf[T]) addChild(_ byte, _ node) { panic("addChild called on leaf") }
 func (n *leaf[T]) findChild(_ byte) *node  { panic("findChild called on leaf") }

server/stree/node.go

@@ -27,6 +27,7 @@ type node interface {
 	matchParts(parts [][]byte) ([][]byte, bool)
 	kind() string
 	iter(f func(node) bool)
+	children() []node
 	numChildren() uint16
 	path() []byte
 }

server/stree/node16.go

@@ -16,8 +16,8 @@ package stree
 // Node with 16 children
 type node16 struct {
 	meta
-	children [16]node
-	keys     [16]byte
+	child [16]node
+	key   [16]byte
 }
 
 func newNode16(prefix []byte) *node16 {
@@ -42,8 +42,8 @@ func (n *node16) addChild(c byte, nn node) {
 	if n.size >= 16 {
 		panic("node16 full!")
 	}
-	n.keys[n.size] = c
-	n.children[n.size] = nn
+	n.key[n.size] = c
+	n.child[n.size] = nn
 	n.size++
 }
 
@@ -52,8 +52,8 @@ func (n *node16) path() []byte        { return n.prefix[:n.prefixLen] }
 
 func (n *node16) findChild(c byte) *node {
 	for i := uint16(0); i < n.size; i++ {
-		if n.keys[i] == c {
-			return &n.children[i]
+		if n.key[i] == c {
+			return &n.child[i]
 		}
 	}
 	return nil
@@ -64,24 +64,24 @@ func (n *node16) isFull() bool { return n.size >= 16 }
 func (n *node16) grow() node {
 	nn := newNode256(n.prefix[:n.prefixLen])
 	for i := 0; i < 16; i++ {
-		nn.addChild(n.keys[i], n.children[i])
+		nn.addChild(n.key[i], n.child[i])
 	}
 	return nn
 }
 
 // Deletes a child from the node.
 func (n *node16) deleteChild(c byte) {
 	for i, last := uint16(0), n.size-1; i < n.size; i++ {
-		if n.keys[i] == c {
+		if n.key[i] == c {
 			// Unsorted so just swap in last one here, else nil if last.
 			if i < last {
-				n.keys[i] = n.keys[last]
-				n.children[i] = n.children[last]
-				n.keys[last] = 0
-				n.children[last] = nil
+				n.key[i] = n.key[last]
+				n.child[i] = n.child[last]
+				n.key[last] = 0
+				n.child[last] = nil
 			} else {
-				n.keys[i] = 0
-				n.children[i] = nil
+				n.key[i] = 0
+				n.child[i] = nil
 			}
 			n.size--
 			return
@@ -96,7 +96,7 @@ func (n *node16) shrink() node {
 	}
 	nn := newNode4(nil)
 	for i := uint16(0); i < n.size; i++ {
-		nn.addChild(n.keys[i], n.children[i])
+		nn.addChild(n.key[i], n.child[i])
 	}
 	return nn
 }
@@ -109,8 +109,13 @@ func (n *node16) matchParts(parts [][]byte) ([][]byte, bool) {
 // Iterate over all children calling func f.
 func (n *node16) iter(f func(node) bool) {
 	for i := uint16(0); i < n.size; i++ {
-		if !f(n.children[i]) {
+		if !f(n.child[i]) {
 			return
 		}
 	}
 }
+
+// Return our children as a slice.
+func (n *node16) children() []node {
+	return n.child[:n.size]
+}

server/stree/node256.go

@@ -16,7 +16,7 @@ package stree
 // Node with 256 children
 type node256 struct {
 	meta
-	children [256]node
+	child [256]node
 }
 
 func newNode256(prefix []byte) *node256 {
@@ -36,16 +36,16 @@ func (n *node256) setPrefix(pre []byte) {
 }
 
 func (n *node256) addChild(c byte, nn node) {
-	n.children[c] = nn
+	n.child[c] = nn
 	n.size++
 }
 
 func (n *node256) numChildren() uint16 { return n.size }
 func (n *node256) path() []byte        { return n.prefix[:n.prefixLen] }
 
 func (n *node256) findChild(c byte) *node {
-	if n.children[c] != nil {
-		return &n.children[c]
+	if n.child[c] != nil {
+		return &n.child[c]
 	}
 	return nil
 }
@@ -55,8 +55,8 @@ func (n *node256) grow() node   { panic("grow can not be called on node256") }
 
 // Deletes a child from the node.
 func (n *node256) deleteChild(c byte) {
-	if n.children[c] != nil {
-		n.children[c] = nil
+	if n.child[c] != nil {
+		n.child[c] = nil
 		n.size--
 	}
 }
@@ -67,9 +67,9 @@ func (n *node256) shrink() node {
 		return nil
 	}
 	nn := newNode16(nil)
-	for c, child := range n.children {
+	for c, child := range n.child {
 		if child != nil {
-			nn.addChild(byte(c), n.children[c])
+			nn.addChild(byte(c), n.child[c])
 		}
 	}
 	return nn
@@ -83,10 +83,15 @@ func (n *node256) matchParts(parts [][]byte) ([][]byte, bool) {
 // Iterate over all children calling func f.
 func (n *node256) iter(f func(node) bool) {
 	for i := 0; i < 256; i++ {
-		if n.children[i] != nil {
-			if !f(n.children[i]) {
+		if n.child[i] != nil {
+			if !f(n.child[i]) {
 				return
 			}
 		}
 	}
 }
+
+// Return our children as a slice.
+func (n *node256) children() []node {
+	return n.child[:256]
+}

server/stree/node4.go

@@ -16,8 +16,8 @@ package stree
 // Node with 4 children
 type node4 struct {
 	meta
-	children [4]node
-	keys     [4]byte
+	child [4]node
+	key   [4]byte
 }
 
 func newNode4(prefix []byte) *node4 {
@@ -41,8 +41,8 @@ func (n *node4) addChild(c byte, nn node) {
 	if n.size >= 4 {
 		panic("node4 full!")
 	}
-	n.keys[n.size] = c
-	n.children[n.size] = nn
+	n.key[n.size] = c
+	n.child[n.size] = nn
 	n.size++
 }
 
@@ -51,8 +51,8 @@ func (n *node4) path() []byte        { return n.prefix[:n.prefixLen] }
 
 func (n *node4) findChild(c byte) *node {
 	for i := uint16(0); i < n.size; i++ {
-		if n.keys[i] == c {
-			return &n.children[i]
+		if n.key[i] == c {
+			return &n.child[i]
 		}
 	}
 	return nil
@@ -63,24 +63,24 @@ func (n *node4) isFull() bool { return n.size >= 4 }
 func (n *node4) grow() node {
 	nn := newNode16(n.prefix[:n.prefixLen])
 	for i := 0; i < 4; i++ {
-		nn.addChild(n.keys[i], n.children[i])
+		nn.addChild(n.key[i], n.child[i])
 	}
 	return nn
 }
 
 // Deletes a child from the node.
 func (n *node4) deleteChild(c byte) {
 	for i, last := uint16(0), n.size-1; i < n.size; i++ {
-		if n.keys[i] == c {
+		if n.key[i] == c {
 			// Unsorted so just swap in last one here, else nil if last.
 			if i < last {
-				n.keys[i] = n.keys[last]
-				n.children[i] = n.children[last]
-				n.keys[last] = 0
-				n.children[last] = nil
+				n.key[i] = n.key[last]
+				n.child[i] = n.child[last]
+				n.key[last] = 0
+				n.child[last] = nil
 			} else {
-				n.keys[i] = 0
-				n.children[i] = nil
+				n.key[i] = 0
+				n.child[i] = nil
 			}
 			n.size--
 			return
@@ -91,7 +91,7 @@ func (n *node4) deleteChild(c byte) {
 // Shrink if needed and return new node, otherwise return nil.
 func (n *node4) shrink() node {
 	if n.size == 1 {
-		return n.children[0]
+		return n.child[0]
 	}
 	return nil
 }
@@ -104,8 +104,13 @@ func (n *node4) matchParts(parts [][]byte) ([][]byte, bool) {
 // Iterate over all children calling func f.
 func (n *node4) iter(f func(node) bool) {
 	for i := uint16(0); i < n.size; i++ {
-		if !f(n.children[i]) {
+		if !f(n.child[i]) {
 			return
 		}
 	}
 }
+
+// Return our children as a slice.
+func (n *node4) children() []node {
+	return n.child[:n.size]
+}

server/stree/stree.go

@@ -303,7 +303,10 @@ func (t *SubjectTree[T]) match(n node, parts [][]byte, pre []byte, cb func(subje
 				nparts = parts[len(parts)-1:]
 				hasTermPWC = true
 			}
-			n.iter(func(cn node) bool {
+			for _, cn := range n.children() {
+				if cn == nil {
+					continue
+				}
 				if cn.isLeaf() {
 					ln := cn.(*leaf[T])
 					if len(ln.suffix) == 0 {
@@ -315,8 +318,7 @@ func (t *SubjectTree[T]) match(n node, parts [][]byte, pre []byte, cb func(subje
 					// We have terminal pwc so call into match again with the child node.
 					t.match(cn, nparts, pre, cb)
 				}
-				return true
-			})
+			}
 			// Return regardless.
 			return
 		}
@@ -330,19 +332,13 @@ func (t *SubjectTree[T]) match(n node, parts [][]byte, pre []byte, cb func(subje
 		fp := nparts[0]
 		p := pivot(fp, 0)
 		// Check if we have a pwc/fwc part here. This will cause us to iterate.
-		if len(fp) == 1 {
-			if p == pwc {
-				// We need to iterate over all children here for the current node
-				// to see if we match further down.
-				n.iter(func(cn node) bool {
-					t.match(cn, nparts, pre, cb)
-					return true
-				})
-			} else if p == fwc {
-				n.iter(func(cn node) bool {
+		if len(fp) == 1 && (p == pwc || p == fwc) {
+			// We need to iterate over all children here for the current node
+			// to see if we match further down.
+			for _, cn := range n.children() {
+				if cn != nil {
 					t.match(cn, nparts, pre, cb)
-					return true
-				})
+				}
 			}
 		}
 		// Here we have normal traversal, so find the next child.
@@ -367,10 +363,11 @@ func (t *SubjectTree[T]) iter(n node, pre []byte, cb func(subject []byte, val *T
 	// Collect nodes since unsorted.
 	var _nodes [256]node
 	nodes := _nodes[:0]
-	n.iter(func(cn node) bool {
-		nodes = append(nodes, cn)
-		return true
-	})
+	for _, cn := range n.children() {
+		if cn != nil {
+			nodes = append(nodes, cn)
+		}
+	}
 	// Now sort.
 	sort.SliceStable(nodes, func(i, j int) bool { return bytes.Compare(nodes[i].path(), nodes[j].path()) < 0 })
 	// Now walk the nodes in order and call into next iter.