Files
mesh-controller/vendor/github.com/nats-io/nats-server/v2/server/stree/stree.go
T
jochen be92762969 Give every test a bus of its own, at the release the mesh runs (hq ADR 0237)
The live tests reached one shared bus and assert, read and remove the mesh's own objects by
their fixed names, so packages run in parallel deleted what each other read and the suite
passed only one package at a time; a red suite read as noise. internal/testbus starts a server
per test, linked in at the nats-server release go.mod pins, and a test holds that pin to the
catalogue's bus image and to the facts snapshot's bus when there is one, so the tests never run
a bus the mesh does not. The waiter test read a timing (the most connections held at one look)
and now reads the state it means (the fewest held across the wait). make check runs the packages
in parallel under the race detector, with a timeout.
2026-10-06 21:17:02 +02:00

511 lines
15 KiB
Go

// Copyright 2023-2025 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
import (
"bytes"
"slices"
"unsafe"
"github.com/nats-io/nats-server/v2/server/gsl"
)
// SubjectTree is an adaptive radix trie (ART) for storing subject information on literal subjects.
// Will use dynamic nodes, path compression and lazy expansion.
// The reason this exists is to not only save some memory in our filestore but to greatly optimize matching
// a wildcard subject to certain members, e.g. consumer NumPending calculations.
type SubjectTree[T any] struct {
root node
size int
}
// NewSubjectTree creates a new SubjectTree with values T.
func NewSubjectTree[T any]() *SubjectTree[T] {
return &SubjectTree[T]{}
}
// Size returns the number of elements stored.
func (t *SubjectTree[T]) Size() int {
if t == nil {
return 0
}
return t.size
}
// Will empty out the tree, or if tree is nil create a new one.
func (t *SubjectTree[T]) Empty() *SubjectTree[T] {
if t == nil {
return NewSubjectTree[T]()
}
t.root, t.size = nil, 0
return t
}
// Insert a value into the tree. Will return if the value was updated and if so the old value.
func (t *SubjectTree[T]) Insert(subject []byte, value T) (*T, bool) {
if t == nil {
return nil, false
}
// Make sure we never insert anything with a noPivot byte.
if bytes.IndexByte(subject, noPivot) >= 0 {
return nil, false
}
old, updated := t.insert(&t.root, subject, value, 0)
if !updated {
t.size++
}
return old, updated
}
// Find will find the value and return it or false if it was not found.
func (t *SubjectTree[T]) Find(subject []byte) (*T, bool) {
if t == nil {
return nil, false
}
var si int
for n := t.root; n != nil; {
if n.isLeaf() {
if ln := n.(*leaf[T]); ln.match(subject[si:]) {
return &ln.value, true
}
return nil, false
}
// We are a node type here, grab meta portion.
if bn := n.base(); len(bn.prefix) > 0 {
end := min(si+len(bn.prefix), len(subject))
if !bytes.Equal(subject[si:end], bn.prefix) {
return nil, false
}
// Increment our subject index.
si += len(bn.prefix)
}
if an := n.findChild(pivot(subject, si)); an != nil {
n = *an
} else {
return nil, false
}
}
return nil, false
}
// Delete will delete the item and return its value, or not found if it did not exist.
func (t *SubjectTree[T]) Delete(subject []byte) (*T, bool) {
if t == nil {
return nil, false
}
val, deleted := t.delete(&t.root, subject, 0)
if deleted {
t.size--
}
return val, deleted
}
// Match will match against a subject that can have wildcards and invoke the callback func for each matched value.
func (t *SubjectTree[T]) Match(filter []byte, cb func(subject []byte, val *T)) {
if t == nil || t.root == nil || len(filter) == 0 || cb == nil {
return
}
// We need to break this up into chunks based on wildcards, either pwc '*' or fwc '>'.
var raw [16][]byte
parts := genParts(filter, raw[:0])
var _pre [256]byte
t.match(t.root, parts, _pre[:0], cb)
}
// IterOrdered will walk all entries in the SubjectTree lexicographically. The callback can return false to terminate the walk.
func (t *SubjectTree[T]) IterOrdered(cb func(subject []byte, val *T) bool) {
if t == nil || t.root == nil {
return
}
var _pre [256]byte
t.iter(t.root, _pre[:0], true, cb)
}
// IterFast will walk all entries in the SubjectTree with no guarantees of ordering. The callback can return false to terminate the walk.
func (t *SubjectTree[T]) IterFast(cb func(subject []byte, val *T) bool) {
if t == nil || t.root == nil {
return
}
var _pre [256]byte
t.iter(t.root, _pre[:0], false, cb)
}
// Internal methods
// Internal call to insert that can be recursive.
func (t *SubjectTree[T]) insert(np *node, subject []byte, value T, si int) (*T, bool) {
n := *np
if n == nil {
*np = newLeaf(subject, value)
return nil, false
}
if n.isLeaf() {
ln := n.(*leaf[T])
if ln.match(subject[si:]) {
// Replace with new value.
old := ln.value
ln.value = value
return &old, true
}
// Here we need to split this leaf.
cpi := commonPrefixLen(ln.suffix, subject[si:])
nn := newNode4(subject[si : si+cpi])
ln.setSuffix(ln.suffix[cpi:])
si += cpi
// Make sure we have different pivot, normally this will be the case unless we have overflowing prefixes.
if p := pivot(ln.suffix, 0); cpi > 0 && si < len(subject) && p == subject[si] {
// We need to split the original leaf. Recursively call into insert.
t.insert(np, subject, value, si)
// Now add the update version of *np as a child to the new node4.
nn.addChild(p, *np)
} else {
// Can just add this new leaf as a sibling.
nl := newLeaf(subject[si:], value)
nn.addChild(pivot(nl.suffix, 0), nl)
// Add back original.
nn.addChild(pivot(ln.suffix, 0), ln)
}
*np = nn
return nil, false
}
// Non-leaf nodes.
bn := n.base()
if len(bn.prefix) > 0 {
cpi := commonPrefixLen(bn.prefix, subject[si:])
if pli := len(bn.prefix); cpi >= pli {
// Move past this node. We look for an existing child node to recurse into.
// If one does not exist we can create a new leaf node.
si += pli
if nn := n.findChild(pivot(subject, si)); nn != nil {
return t.insert(nn, subject, value, si)
}
if n.isFull() {
n = n.grow()
*np = n
}
n.addChild(pivot(subject, si), newLeaf(subject[si:], value))
return nil, false
} else {
// We did not match the prefix completely here.
// Calculate new prefix for this node.
prefix := subject[si : si+cpi]
si += len(prefix)
// We will insert a new node4 and attach our current node below after adjusting prefix.
nn := newNode4(prefix)
// Shift the prefix for our original node.
n.setPrefix(bn.prefix[cpi:])
nn.addChild(pivot(bn.prefix[:], 0), n)
// Add in our new leaf.
nn.addChild(pivot(subject[si:], 0), newLeaf(subject[si:], value))
// Update our node reference.
*np = nn
}
} else {
if nn := n.findChild(pivot(subject, si)); nn != nil {
return t.insert(nn, subject, value, si)
}
// No prefix and no matched child, so add in new leafnode as needed.
if n.isFull() {
n = n.grow()
*np = n
}
n.addChild(pivot(subject, si), newLeaf(subject[si:], value))
}
return nil, false
}
// internal function to recursively find the leaf to delete. Will do compaction if the item is found and removed.
func (t *SubjectTree[T]) delete(np *node, subject []byte, si int) (*T, bool) {
if t == nil || np == nil || *np == nil || len(subject) == 0 {
return nil, false
}
n := *np
if n.isLeaf() {
ln := n.(*leaf[T])
if ln.match(subject[si:]) {
*np = nil
return &ln.value, true
}
return nil, false
}
// Not a leaf node.
if bn := n.base(); len(bn.prefix) > 0 {
// subject could be shorter and would panic on bad index into subject slice.
if len(subject) < si+len(bn.prefix) {
return nil, false
}
if !bytes.Equal(subject[si:si+len(bn.prefix)], bn.prefix) {
return nil, false
}
// Increment our subject index.
si += len(bn.prefix)
}
p := pivot(subject, si)
nna := n.findChild(p)
if nna == nil {
return nil, false
}
nn := *nna
if nn.isLeaf() {
ln := nn.(*leaf[T])
if ln.match(subject[si:]) {
n.deleteChild(p)
if sn := n.shrink(); sn != nil {
bn := n.base()
// Make sure to set cap so we force an append to copy below.
pre := bn.prefix[:len(bn.prefix):len(bn.prefix)]
// Need to fix up prefixes/suffixes.
if sn.isLeaf() {
ln := sn.(*leaf[T])
// Make sure to set cap so we force an append to copy.
ln.suffix = append(pre, ln.suffix...)
} else {
// We are a node here, we need to add in the old prefix.
if len(pre) > 0 {
bsn := sn.base()
sn.setPrefix(append(pre, bsn.prefix...))
}
}
*np = sn
}
return &ln.value, true
}
return nil, false
}
return t.delete(nna, subject, si)
}
// Internal function which can be called recursively to match all leaf nodes to a given filter subject which
// once here has been decomposed to parts. These parts only care about wildcards, both pwc and fwc.
func (t *SubjectTree[T]) match(n node, parts [][]byte, pre []byte, cb func(subject []byte, val *T)) {
// Capture if we are sitting on a terminal fwc.
var hasFWC bool
if lp := len(parts); lp > 0 && len(parts[lp-1]) > 0 && parts[lp-1][0] == fwc {
hasFWC = true
}
for n != nil {
nparts, matched := n.matchParts(parts)
// Check if we did not match.
if !matched {
return
}
// We have matched here. If we are a leaf and have exhausted all parts or he have a FWC fire callback.
if n.isLeaf() {
if len(nparts) == 0 || (hasFWC && len(nparts) == 1) {
ln := n.(*leaf[T])
cb(append(pre, ln.suffix...), &ln.value)
}
return
}
// We have normal nodes here.
// We need to append our prefix
bn := n.base()
if len(bn.prefix) > 0 {
// Note that this append may reallocate, but it doesn't modify "pre" at the "match" callsite.
pre = append(pre, bn.prefix...)
}
// Check our remaining parts.
if len(nparts) == 0 && !hasFWC {
// We are a node with no parts left and we are not looking at a fwc.
// We could have a leafnode with no suffix which would be a match.
// We could also have a terminal pwc. Check for those here.
var hasTermPWC bool
if lp := len(parts); lp > 0 && len(parts[lp-1]) == 1 && parts[lp-1][0] == pwc {
// If we are sitting on a terminal pwc, put the pwc back and continue.
nparts = parts[len(parts)-1:]
hasTermPWC = true
}
for _, cn := range n.children() {
if cn == nil {
continue
}
if cn.isLeaf() {
ln := cn.(*leaf[T])
if len(ln.suffix) == 0 {
cb(append(pre, ln.suffix...), &ln.value)
} else if hasTermPWC && bytes.IndexByte(ln.suffix, tsep) < 0 {
cb(append(pre, ln.suffix...), &ln.value)
}
} else if hasTermPWC {
// We have terminal pwc so call into match again with the child node.
t.match(cn, nparts, pre, cb)
}
}
// Return regardless.
return
}
// If we are sitting on a terminal fwc, put back and continue.
if hasFWC && len(nparts) == 0 {
nparts = parts[len(parts)-1:]
}
// Here we are a node type with a partial match.
// Check if the first part is a wildcard.
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 && (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
}
// Here we have normal traversal, so find the next child.
nn := n.findChild(p)
if nn == nil {
return
}
n, parts = *nn, nparts
}
}
// Internal iter function to walk nodes in lexicographical order.
func (t *SubjectTree[T]) iter(n node, pre []byte, ordered bool, cb func(subject []byte, val *T) bool) bool {
if n.isLeaf() {
ln := n.(*leaf[T])
return cb(append(pre, ln.suffix...), &ln.value)
}
// We are normal node here.
bn := n.base()
// Note that this append may reallocate, but it doesn't modify "pre" at the "iter" callsite.
pre = append(pre, bn.prefix...)
// Not everything requires lexicographical sorting, so support a fast path for iterating in
// whatever order the stree has things stored instead.
if !ordered {
for _, cn := range n.children() {
if cn == nil {
continue
}
if !t.iter(cn, pre, false, cb) {
return false
}
}
return true
}
// Collect nodes since unsorted.
var _nodes [256]node
nodes := _nodes[:0]
for _, cn := range n.children() {
if cn != nil {
nodes = append(nodes, cn)
}
}
// Now sort.
slices.SortStableFunc(nodes, func(a, b node) int { return bytes.Compare(a.path(), b.path()) })
// Now walk the nodes in order and call into next iter.
for i := range nodes {
if !t.iter(nodes[i], pre, true, cb) {
return false
}
}
return true
}
// LazyIntersect iterates the smaller of the two provided subject trees and
// looks for matching entries in the other. It is lazy in that it does not
// aggressively optimize against repeated walks, but is considerably faster
// in most cases than intersecting against a potentially large sublist.
func LazyIntersect[TL, TR any](tl *SubjectTree[TL], tr *SubjectTree[TR], cb func([]byte, *TL, *TR)) {
if tl == nil || tr == nil || tl.root == nil || tr.root == nil {
return
}
// Iterate over the smaller tree to reduce the number of rounds.
if tl.Size() <= tr.Size() {
tl.IterFast(func(key []byte, v1 *TL) bool {
if v2, ok := tr.Find(key); ok {
cb(key, v1, v2)
}
return true
})
} else {
tr.IterFast(func(key []byte, v2 *TR) bool {
if v1, ok := tl.Find(key); ok {
cb(key, v1, v2)
}
return true
})
}
}
// IntersectGSL will match all items in the given subject tree that
// have interest expressed in the given sublist. The callback will only be called
// once for each subject, regardless of overlapping subscriptions in the sublist.
func IntersectGSL[T any, SL comparable](t *SubjectTree[T], sl *gsl.GenericSublist[SL], cb func(subject []byte, val *T)) {
if t == nil || t.root == nil || sl == nil {
return
}
var _pre [256]byte
_intersectGSL(t.root, _pre[:0], sl, cb)
}
func _intersectGSL[T any, SL comparable](n node, pre []byte, sl *gsl.GenericSublist[SL], cb func(subject []byte, val *T)) {
if n.isLeaf() {
ln := n.(*leaf[T])
subj := append(pre, ln.suffix...)
if sl.HasInterest(bytesToString(subj)) {
cb(subj, &ln.value)
}
return
}
bn := n.base()
pre = append(pre, bn.prefix...)
for _, cn := range n.children() {
if cn == nil {
continue
}
subj := append(pre, cn.path()...)
if !hasInterestForTokens(sl, subj, len(pre)) {
continue
}
_intersectGSL(cn, pre, sl, cb)
}
}
// The subject tree can return partial tokens so we need to check starting interest
// only from whole tokens when we encounter a tsep.
func hasInterestForTokens[SL comparable](sl *gsl.GenericSublist[SL], subj []byte, since int) bool {
for i := since; i < len(subj); i++ {
if subj[i] == tsep {
if !sl.HasInterestStartingIn(bytesToString(subj[:i])) {
return false
}
}
}
return true
}
// Note this will avoid a copy of the data used for the string, but it will also reference the existing slice's data pointer.
// So this should be used sparingly when we know the encompassing byte slice's lifetime is the same.
func bytesToString(b []byte) string {
if len(b) == 0 {
return ""
}
p := unsafe.SliceData(b)
return unsafe.String(p, len(b))
}