wehub-resource-sync
355 lines
13 KiB
Go
355 lines
13 KiB
Go
package radixtree_test
|
|
|
|
import (
|
|
"time"
|
|
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
|
|
"github.com/mudler/LocalAI/pkg/radixtree"
|
|
)
|
|
|
|
var t0 = time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC)
|
|
|
|
var _ = Describe("Tree construction", func() {
|
|
It("returns an empty tree that matches nothing", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Minute})
|
|
_, depth, ok := tr.LongestMatch([]uint64{1, 2, 3}, t0)
|
|
Expect(ok).To(BeFalse())
|
|
Expect(depth).To(Equal(0))
|
|
})
|
|
})
|
|
|
|
var _ = Describe("Insert and LongestMatch", func() {
|
|
It("returns the deepest matching prefix value", func() {
|
|
// Non-overlapping chains keep the longest-prefix intent clean: every
|
|
// node on the value's own chain records that value, and no other Insert
|
|
// overwrites a shared prefix node. A query that runs off the end of a
|
|
// chain stops matching at the deepest stored element it reached.
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
tr.Insert([]uint64{1, 2, 3, 4}, "nodeB", t0)
|
|
tr.Insert([]uint64{7, 8}, "nodeA", t0)
|
|
|
|
v, depth, ok := tr.LongestMatch([]uint64{1, 2, 3, 4, 5}, t0)
|
|
Expect(ok).To(BeTrue())
|
|
Expect(v).To(Equal("nodeB"))
|
|
Expect(depth).To(Equal(4))
|
|
|
|
v, depth, ok = tr.LongestMatch([]uint64{7, 8, 9}, t0)
|
|
Expect(ok).To(BeTrue())
|
|
Expect(v).To(Equal("nodeA"))
|
|
Expect(depth).To(Equal(2))
|
|
})
|
|
|
|
It("lets the last writer own a shared prefix node", func() {
|
|
// When two chains share a leading block, value-at-every-node means the
|
|
// later Insert overwrites the shared prefix node. Inserting nodeA on
|
|
// [1,2] then nodeB on [1,2,3,4] makes nodeB own [1] and [1,2], so a
|
|
// query that diverges within the shared block resolves to nodeB. This
|
|
// is the intended recency heuristic: the most recent chain through that
|
|
// block is the one most likely still warm.
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
tr.Insert([]uint64{1, 2}, "nodeA", t0)
|
|
tr.Insert([]uint64{1, 2, 3, 4}, "nodeB", t0)
|
|
|
|
v, depth, ok := tr.LongestMatch([]uint64{1, 2, 3, 4, 5}, t0)
|
|
Expect(ok).To(BeTrue())
|
|
Expect(v).To(Equal("nodeB"))
|
|
Expect(depth).To(Equal(4))
|
|
|
|
// The shared prefix [1,2] is now owned by nodeB (last writer wins).
|
|
v, depth, ok = tr.LongestMatch([]uint64{1, 2, 9}, t0)
|
|
Expect(ok).To(BeTrue())
|
|
Expect(v).To(Equal("nodeB"))
|
|
Expect(depth).To(Equal(2))
|
|
})
|
|
|
|
It("returns ok=false when no prefix is stored", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
tr.Insert([]uint64{7, 8}, "nodeA", t0)
|
|
_, _, ok := tr.LongestMatch([]uint64{1, 2}, t0)
|
|
Expect(ok).To(BeFalse())
|
|
})
|
|
|
|
It("matches a shared prefix when the query tail diverges", func() {
|
|
// SGLang/vLLM-style prefix matching: a single Insert of a full chain
|
|
// must let any query that shares a leading block match at the depth of
|
|
// the deepest shared element, even though the tails differ. This is the
|
|
// core use case (shared system prompt / multi-turn extension / volatile
|
|
// tail), not exact-repeat.
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
tr.Insert([]uint64{1, 2, 3, 4, 5}, "nodeA", t0)
|
|
v, depth, ok := tr.LongestMatch([]uint64{1, 2, 3, 9, 9}, t0)
|
|
Expect(ok).To(BeTrue())
|
|
Expect(depth).To(Equal(3)) // shared prefix [1,2,3]
|
|
Expect(v).To(Equal("nodeA"))
|
|
})
|
|
})
|
|
|
|
var _ = Describe("TTL expiry", func() {
|
|
It("does not match an entry past its TTL", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Minute})
|
|
tr.Insert([]uint64{1, 2}, "nodeA", t0)
|
|
_, _, ok := tr.LongestMatch([]uint64{1, 2}, t0.Add(2*time.Minute))
|
|
Expect(ok).To(BeFalse())
|
|
})
|
|
|
|
It("refreshes lastSeen on re-insert so a live path survives", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Minute})
|
|
tr.Insert([]uint64{1, 2}, "nodeA", t0)
|
|
tr.Insert([]uint64{1, 2}, "nodeA", t0.Add(50*time.Second))
|
|
_, _, ok := tr.LongestMatch([]uint64{1, 2}, t0.Add(90*time.Second))
|
|
Expect(ok).To(BeTrue())
|
|
})
|
|
|
|
It("Evict reclaims expired nodes", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Minute})
|
|
// Value-at-every-node: Insert of a 2-element chain records nodeA at both
|
|
// {1} and {1,2}, so Len is 2 (one valued node per distinct prefix).
|
|
tr.Insert([]uint64{1, 2}, "nodeA", t0)
|
|
Expect(tr.Len()).To(Equal(2))
|
|
tr.Evict(t0.Add(2 * time.Minute))
|
|
Expect(tr.Len()).To(Equal(0))
|
|
})
|
|
})
|
|
|
|
var _ = Describe("Weight", func() {
|
|
It("counts live entries for a value with no decay", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour}) // HalfLife=0
|
|
tr.Insert([]uint64{1}, "A", t0)
|
|
tr.Insert([]uint64{1, 2}, "A", t0)
|
|
tr.Insert([]uint64{9}, "B", t0)
|
|
Expect(tr.Weight("A", t0)).To(BeNumerically("==", 2))
|
|
Expect(tr.Weight("B", t0)).To(BeNumerically("==", 1))
|
|
Expect(tr.Weight("C", t0)).To(BeNumerically("==", 0))
|
|
})
|
|
|
|
It("decays older entries by half-life", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour, HalfLife: time.Minute})
|
|
tr.Insert([]uint64{1}, "A", t0)
|
|
// one half-life later, the entry weighs 0.5
|
|
Expect(tr.Weight("A", t0.Add(time.Minute))).To(BeNumerically("~", 0.5, 0.001))
|
|
})
|
|
|
|
It("ignores expired entries", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Minute})
|
|
tr.Insert([]uint64{1}, "A", t0)
|
|
Expect(tr.Weight("A", t0.Add(2*time.Minute))).To(BeNumerically("==", 0))
|
|
})
|
|
})
|
|
|
|
var _ = Describe("WeightsFor", func() {
|
|
It("matches per-value Weight with no decay", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour}) // HalfLife=0
|
|
tr.Insert([]uint64{1}, "A", t0)
|
|
tr.Insert([]uint64{1, 2}, "A", t0)
|
|
tr.Insert([]uint64{9}, "B", t0)
|
|
|
|
got := tr.WeightsFor([]string{"A", "B", "C"}, t0)
|
|
Expect(got).To(HaveLen(3))
|
|
Expect(got["A"]).To(BeNumerically("==", 2))
|
|
Expect(got["B"]).To(BeNumerically("==", 1))
|
|
Expect(got["C"]).To(BeNumerically("==", 0))
|
|
})
|
|
|
|
It("matches per-value Weight under decay", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour, HalfLife: time.Minute})
|
|
tr.Insert([]uint64{1}, "A", t0)
|
|
tr.Insert([]uint64{1, 2}, "A", t0.Add(30*time.Second))
|
|
tr.Insert([]uint64{9}, "B", t0)
|
|
|
|
now := t0.Add(time.Minute)
|
|
got := tr.WeightsFor([]string{"A", "B", "C"}, now)
|
|
Expect(got["A"]).To(BeNumerically("~", tr.Weight("A", now), 1e-12))
|
|
Expect(got["B"]).To(BeNumerically("~", tr.Weight("B", now), 1e-12))
|
|
Expect(got["C"]).To(BeNumerically("==", 0))
|
|
})
|
|
|
|
It("respects TTL expiry and matches Weight at a non-zero age under decay", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Minute, HalfLife: 30 * time.Second})
|
|
tr.Insert([]uint64{1}, "A", t0) // will be expired at now
|
|
tr.Insert([]uint64{2}, "A", t0.Add(90*time.Second)) // live, aged 30s at now
|
|
tr.Insert([]uint64{9}, "B", t0) // expired at now
|
|
|
|
now := t0.Add(2 * time.Minute)
|
|
got := tr.WeightsFor([]string{"A", "B"}, now)
|
|
Expect(got["A"]).To(BeNumerically("~", tr.Weight("A", now), 1e-12))
|
|
Expect(got["A"]).To(BeNumerically("~", 0.5, 0.001)) // single live entry aged one half-life
|
|
Expect(got["B"]).To(BeNumerically("==", 0))
|
|
})
|
|
|
|
It("returns an empty map for an empty values slice", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
tr.Insert([]uint64{1}, "A", t0)
|
|
Expect(tr.WeightsFor(nil, t0)).To(BeEmpty())
|
|
Expect(tr.WeightsFor([]string{}, t0)).To(BeEmpty())
|
|
})
|
|
|
|
It("maps a value not present in the tree to 0", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
tr.Insert([]uint64{1}, "A", t0)
|
|
got := tr.WeightsFor([]string{"Z"}, t0)
|
|
Expect(got).To(HaveLen(1))
|
|
Expect(got["Z"]).To(BeNumerically("==", 0))
|
|
})
|
|
})
|
|
|
|
var _ = Describe("Remove", func() {
|
|
It("drops every entry anchored to a value and prunes", func() {
|
|
// Non-overlapping chains so Remove("A") and the survival of B are both
|
|
// meaningful: with value-at-every-node, overlapping chains would let the
|
|
// later writer own the shared prefix nodes, so A could own nothing and
|
|
// the test would be vacuous.
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
tr.Insert([]uint64{1, 2}, "A", t0)
|
|
tr.Insert([]uint64{7, 8, 9}, "B", t0)
|
|
tr.Remove("A")
|
|
_, _, ok := tr.LongestMatch([]uint64{1, 2}, t0)
|
|
Expect(ok).To(BeFalse()) // A gone; its branch is fully reclaimed
|
|
v, _, ok := tr.LongestMatch([]uint64{7, 8, 9}, t0)
|
|
Expect(ok).To(BeTrue())
|
|
Expect(v).To(Equal("B")) // B survives
|
|
Expect(tr.Weight("A", t0)).To(BeNumerically("==", 0))
|
|
})
|
|
})
|
|
|
|
var _ = Describe("RemoveFunc", func() {
|
|
It("drops every entry matching the predicate, prunes, and keeps the rest", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
tr.Insert([]uint64{1, 2}, "drop-a", t0)
|
|
tr.Insert([]uint64{3, 4}, "drop-b", t0)
|
|
tr.Insert([]uint64{7, 8, 9}, "keep", t0)
|
|
// Drop everything whose value starts with "drop".
|
|
tr.RemoveFunc(func(v string) bool { return len(v) >= 4 && v[:4] == "drop" })
|
|
_, _, ok := tr.LongestMatch([]uint64{1, 2}, t0)
|
|
Expect(ok).To(BeFalse())
|
|
_, _, ok = tr.LongestMatch([]uint64{3, 4}, t0)
|
|
Expect(ok).To(BeFalse())
|
|
v, _, ok := tr.LongestMatch([]uint64{7, 8, 9}, t0)
|
|
Expect(ok).To(BeTrue())
|
|
Expect(v).To(Equal("keep"))
|
|
Expect(tr.Len()).To(Equal(3)) // only the 3-node "keep" chain remains
|
|
})
|
|
|
|
It("makes Remove a special case of RemoveFunc", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
tr.Insert([]uint64{1, 2}, "A", t0)
|
|
tr.Insert([]uint64{7, 8, 9}, "B", t0)
|
|
tr.RemoveFunc(func(v string) bool { return v == "A" })
|
|
_, _, ok := tr.LongestMatch([]uint64{1, 2}, t0)
|
|
Expect(ok).To(BeFalse())
|
|
v, _, ok := tr.LongestMatch([]uint64{7, 8, 9}, t0)
|
|
Expect(ok).To(BeTrue())
|
|
Expect(v).To(Equal("B"))
|
|
})
|
|
})
|
|
|
|
var _ = Describe("TTL boundary", func() {
|
|
It("treats age exactly equal to TTL as still live, and one tick past as expired", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Minute})
|
|
tr.Insert([]uint64{1, 2}, "A", t0)
|
|
|
|
// age == TTL: strict greater-than means this is still live.
|
|
_, _, ok := tr.LongestMatch([]uint64{1, 2}, t0.Add(time.Minute))
|
|
Expect(ok).To(BeTrue())
|
|
|
|
// one nanosecond past TTL: expired.
|
|
_, _, ok = tr.LongestMatch([]uint64{1, 2}, t0.Add(time.Minute+time.Nanosecond))
|
|
Expect(ok).To(BeFalse())
|
|
})
|
|
})
|
|
|
|
var _ = Describe("MaxEntries eviction", func() {
|
|
It("drops the least-recently-seen entry when the cap is exceeded", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour, MaxEntries: 2})
|
|
tr.Insert([]uint64{1}, "A", t0)
|
|
tr.Insert([]uint64{2}, "B", t0.Add(time.Second))
|
|
tr.Insert([]uint64{3}, "C", t0.Add(2*time.Second))
|
|
|
|
Expect(tr.Len()).To(Equal(2))
|
|
|
|
// A was the least-recently-seen, so it is the one dropped.
|
|
_, _, ok := tr.LongestMatch([]uint64{1}, t0.Add(2*time.Second))
|
|
Expect(ok).To(BeFalse())
|
|
|
|
// B and C survive.
|
|
_, _, ok = tr.LongestMatch([]uint64{2}, t0.Add(2*time.Second))
|
|
Expect(ok).To(BeTrue())
|
|
_, _, ok = tr.LongestMatch([]uint64{3}, t0.Add(2*time.Second))
|
|
Expect(ok).To(BeTrue())
|
|
})
|
|
|
|
It("prunes value-less ancestors left behind by an eviction", func() {
|
|
// Value-at-every-node: Inserting the deep chain B = [1,2,3] records B at
|
|
// {1}, {1,2}, and {1,2,3} (three valued nodes). With the cap at 2, the
|
|
// least-recently-seen valued nodes are evicted one per subsequent Insert.
|
|
// The two fresh single-element keys (C, D) are newer, so eviction keeps
|
|
// peeling B's nodes off until B's entire branch is reclaimed - none of
|
|
// its internal nodes may linger and inflate Len past the cap.
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour, MaxEntries: 2})
|
|
tr.Insert([]uint64{1, 2, 3}, "B", t0)
|
|
tr.Insert([]uint64{5}, "C", t0.Add(time.Second))
|
|
tr.Insert([]uint64{6}, "D", t0.Add(2*time.Second))
|
|
|
|
Expect(tr.Len()).To(Equal(2))
|
|
// B (oldest) evicted; its deep branch reclaimed.
|
|
_, _, ok := tr.LongestMatch([]uint64{1, 2, 3}, t0.Add(2*time.Second))
|
|
Expect(ok).To(BeFalse())
|
|
_, _, ok = tr.LongestMatch([]uint64{1, 2}, t0.Add(2*time.Second))
|
|
Expect(ok).To(BeFalse())
|
|
Expect(tr.Weight("B", t0.Add(2*time.Second))).To(BeNumerically("==", 0))
|
|
})
|
|
|
|
It("reclaims structure so the tree never grows past the cap under churn", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour, MaxEntries: 2})
|
|
tr.Insert([]uint64{1}, "A", t0)
|
|
tr.Insert([]uint64{2}, "B", t0.Add(time.Second))
|
|
Expect(tr.Len()).To(Equal(2))
|
|
|
|
for i := range 10 {
|
|
tr.Insert([]uint64{uint64(100 + i)}, "X", t0.Add(time.Duration(i+2)*time.Second))
|
|
Expect(tr.Len()).To(Equal(2))
|
|
}
|
|
})
|
|
})
|
|
|
|
var _ = Describe("Empty key", func() {
|
|
It("LongestMatch on an empty key returns ok=false", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
tr.Insert([]uint64{1, 2}, "A", t0)
|
|
_, depth, ok := tr.LongestMatch([]uint64{}, t0)
|
|
Expect(ok).To(BeFalse())
|
|
Expect(depth).To(Equal(0))
|
|
})
|
|
|
|
It("Insert with an empty key is a no-op that creates no root value", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
Expect(func() { tr.Insert([]uint64{}, "A", t0) }).NotTo(Panic())
|
|
Expect(tr.Len()).To(Equal(0))
|
|
_, _, ok := tr.LongestMatch([]uint64{}, t0)
|
|
Expect(ok).To(BeFalse())
|
|
Expect(tr.Weight("A", t0)).To(BeNumerically("==", 0))
|
|
})
|
|
})
|
|
|
|
var _ = Describe("Concurrent access", func() {
|
|
It("is race-free under parallel insert/match/weight", func() {
|
|
tr := radixtree.New[string](radixtree.Options{TTL: time.Hour})
|
|
done := make(chan struct{})
|
|
for g := range 8 {
|
|
go func(g int) {
|
|
defer GinkgoRecover()
|
|
for i := range 1000 {
|
|
tr.Insert([]uint64{uint64(g), uint64(i % 10)}, "n", t0)
|
|
tr.LongestMatch([]uint64{uint64(g), 1}, t0)
|
|
tr.Weight("n", t0)
|
|
}
|
|
done <- struct{}{}
|
|
}(g)
|
|
}
|
|
for range 8 {
|
|
<-done
|
|
}
|
|
})
|
|
})
|