Add node10 to stree (#6106)

Although we probably don't want too many different node sizes here, the
`node10` case is particularly interesting because it perfectly fits the
full 0-9 numeric range without wasting bytes. In fact it saves 96 bytes
(208 bytes instead of 304) compared to using `node16` for the same
purpose.

This means memory savings for tracking subjects which are either mostly
numerical throughout, or have tokens that are primarily numerical.

For a subject space that is mostly numerical, this can be several GBs
less at the half-billion subjects mark and the saving can grow above
that.

Signed-off-by: Neil Twigg <neil@nats.io>

server/stree/dump.go

@@ -50,6 +50,7 @@ func (t *SubjectTree[T]) dump(w io.Writer, n node, depth int) {
 // For individual node/leaf dumps.
 func (n *leaf[T]) kind() string { return "LEAF" }
 func (n *node4) kind() string   { return "NODE4" }
+func (n *node10) kind() string  { return "NODE10" }
 func (n *node16) kind() string  { return "NODE16" }
 func (n *node48) kind() string  { return "NODE48" }
 func (n *node256) kind() string { return "NODE256" }

server/stree/node10.go

@@ -0,0 +1,106 @@
+// Copyright 2023-2024 The NATS Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package stree
+
+// Node with 10 children
+// This node size is for the particular case that a part of the subject is numeric
+// in nature, i.e. it only needs to satisfy the range 0-9 without wasting bytes
+// Order of struct fields for best memory alignment (as per govet/fieldalignment)
+type node10 struct {
+	child [10]node
+	meta
+	key [10]byte
+}
+
+func newNode10(prefix []byte) *node10 {
+	nn := &node10{}
+	nn.setPrefix(prefix)
+	return nn
+}
+
+// Currently we do not keep node10 sorted or use bitfields for traversal so just add to the end.
+// TODO(dlc) - We should revisit here with more detailed benchmarks.
+func (n *node10) addChild(c byte, nn node) {
+	if n.size >= 10 {
+		panic("node10 full!")
+	}
+	n.key[n.size] = c
+	n.child[n.size] = nn
+	n.size++
+}
+
+func (n *node10) findChild(c byte) *node {
+	for i := uint16(0); i < n.size; i++ {
+		if n.key[i] == c {
+			return &n.child[i]
+		}
+	}
+	return nil
+}
+
+func (n *node10) isFull() bool { return n.size >= 10 }
+
+func (n *node10) grow() node {
+	nn := newNode16(n.prefix)
+	for i := 0; i < 10; i++ {
+		nn.addChild(n.key[i], n.child[i])
+	}
+	return nn
+}
+
+// Deletes a child from the node.
+func (n *node10) deleteChild(c byte) {
+	for i, last := uint16(0), n.size-1; i < n.size; i++ {
+		if n.key[i] == c {
+			// Unsorted so just swap in last one here, else nil if last.
+			if i < last {
+				n.key[i] = n.key[last]
+				n.child[i] = n.child[last]
+				n.key[last] = 0
+				n.child[last] = nil
+			} else {
+				n.key[i] = 0
+				n.child[i] = nil
+			}
+			n.size--
+			return
+		}
+	}
+}
+
+// Shrink if needed and return new node, otherwise return nil.
+func (n *node10) shrink() node {
+	if n.size > 4 {
+		return nil
+	}
+	nn := newNode4(nil)
+	for i := uint16(0); i < n.size; i++ {
+		nn.addChild(n.key[i], n.child[i])
+	}
+	return nn
+}
+
+// Iterate over all children calling func f.
+func (n *node10) iter(f func(node) bool) {
+	for i := uint16(0); i < n.size; i++ {
+		if !f(n.child[i]) {
+			return
+		}
+	}
+}
+
+// Return our children as a slice.
+func (n *node10) children() []node {
+	return n.child[:n.size]
+}

server/stree/node16.go

@@ -79,10 +79,10 @@ func (n *node16) deleteChild(c byte) {
 
 // Shrink if needed and return new node, otherwise return nil.
 func (n *node16) shrink() node {
-	if n.size > 4 {
+	if n.size > 10 {
 		return nil
 	}
-	nn := newNode4(nil)
+	nn := newNode10(nil)
 	for i := uint16(0); i < n.size; i++ {
 		nn.addChild(n.key[i], n.child[i])
 	}

server/stree/node4.go

@@ -49,7 +49,7 @@ func (n *node4) findChild(c byte) *node {
 func (n *node4) isFull() bool { return n.size >= 4 }
 
 func (n *node4) grow() node {
-	nn := newNode16(n.prefix)
+	nn := newNode10(n.prefix)
 	for i := 0; i < 4; i++ {
 		nn.addChild(n.key[i], n.child[i])
 	}

server/stree/stree_test.go

@@ -79,9 +79,22 @@ func TestSubjectTreeNodeGrow(t *testing.T) {
 	old, updated := st.Insert(b("foo.bar.E"), 22)
 	require_True(t, old == nil)
 	require_False(t, updated)
+	_, ok = st.root.(*node10)
+	require_True(t, ok)
+	for i := 5; i < 10; i++ {
+		subj := b(fmt.Sprintf("foo.bar.%c", 'A'+i))
+		old, updated := st.Insert(subj, 22)
+		require_True(t, old == nil)
+		require_False(t, updated)
+	}
+	// This one will trigger us to grow.
+	old, updated = st.Insert(b("foo.bar.K"), 22)
+	require_True(t, old == nil)
+	require_False(t, updated)
+	// We have filled a node10.
 	_, ok = st.root.(*node16)
 	require_True(t, ok)
-	for i := 5; i < 16; i++ {
+	for i := 11; i < 16; i++ {
 		subj := b(fmt.Sprintf("foo.bar.%c", 'A'+i))
 		old, updated := st.Insert(subj, 22)
 		require_True(t, old == nil)
@@ -164,18 +177,33 @@ func TestSubjectTreeNodeDelete(t *testing.T) {
 	require_True(t, found)
 	require_Equal(t, *v, 11)
 	require_Equal(t, st.root, nil)
-	// Now pop up to a node16 and make sure we can shrink back down.
+	// Now pop up to a node10 and make sure we can shrink back down.
 	for i := 0; i < 5; i++ {
 		subj := fmt.Sprintf("foo.bar.%c", 'A'+i)
 		st.Insert(b(subj), 22)
 	}
-	_, ok := st.root.(*node16)
+	_, ok := st.root.(*node10)
 	require_True(t, ok)
 	v, found = st.Delete(b("foo.bar.A"))
 	require_True(t, found)
 	require_Equal(t, *v, 22)
 	_, ok = st.root.(*node4)
 	require_True(t, ok)
+	// Now pop up to node16
+	for i := 0; i < 11; i++ {
+		subj := fmt.Sprintf("foo.bar.%c", 'A'+i)
+		st.Insert(b(subj), 22)
+	}
+	_, ok = st.root.(*node16)
+	require_True(t, ok)
+	v, found = st.Delete(b("foo.bar.A"))
+	require_True(t, found)
+	require_Equal(t, *v, 22)
+	_, ok = st.root.(*node10)
+	require_True(t, ok)
+	v, found = st.Find(b("foo.bar.B"))
+	require_True(t, found)
+	require_Equal(t, *v, 22)
 	// Now pop up to node48
 	st = NewSubjectTree[int]()
 	for i := 0; i < 17; i++ {