511 lines
15 KiB
Go
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))
|
|
}
|