wehub-resource-sync
119 lines
4.8 KiB
Go
119 lines
4.8 KiB
Go
package prefixcache_test
|
|
|
|
import (
|
|
"time"
|
|
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
|
|
"github.com/mudler/LocalAI/core/services/messaging"
|
|
"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
|
|
)
|
|
|
|
type fakePub struct{ published []any }
|
|
|
|
func (f *fakePub) Publish(subject string, v any) error {
|
|
f.published = append(f.published, v)
|
|
return nil
|
|
}
|
|
|
|
// Sync must satisfy the Provider seam so SmartRouter can hold a single
|
|
// prefixcache.Provider that broadcasts via NATS.
|
|
var _ prefixcache.Provider = (*prefixcache.Sync)(nil)
|
|
|
|
var _ = Describe("Sync", func() {
|
|
It("delegates Evict to the wrapped index", func() {
|
|
cfg := prefixcache.DefaultConfig()
|
|
cfg.TTL = time.Minute
|
|
idx := prefixcache.NewIndex(cfg)
|
|
s := prefixcache.NewSync(idx, &fakePub{})
|
|
s.Observe("m", []uint64{1, 2}, rk("A", 0), t0)
|
|
// Before TTL: still hot.
|
|
Expect(idx.Decide("m", []uint64{1, 2}, []prefixcache.ReplicaKey{rk("A", 0)}, t0).HasHot).To(BeTrue())
|
|
// After TTL via Sync.Evict: entry is swept.
|
|
s.Evict(t0.Add(2 * time.Minute))
|
|
Expect(idx.Decide("m", []uint64{1, 2}, []prefixcache.ReplicaKey{rk("A", 0)}, t0.Add(2*time.Minute)).HasHot).To(BeFalse())
|
|
})
|
|
|
|
It("publishes an observe event with the replica when Observe is new", func() {
|
|
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
|
pub := &fakePub{}
|
|
s := prefixcache.NewSync(idx, pub)
|
|
s.Observe("m", []uint64{1, 2}, rk("A", 1), t0) // first time -> publish
|
|
Expect(pub.published).To(HaveLen(1))
|
|
ev := pub.published[0].(messaging.PrefixCacheObserveEvent)
|
|
Expect(ev.NodeID).To(Equal("A"))
|
|
Expect(ev.Replica).To(Equal(1))
|
|
s.Observe("m", []uint64{1, 2}, rk("A", 1), t0) // same -> no publish
|
|
Expect(pub.published).To(HaveLen(1))
|
|
})
|
|
|
|
It("broadcasts an invalidate even for a model with no local tree, without interning one", func() {
|
|
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
|
pub := &fakePub{}
|
|
s := prefixcache.NewSync(idx, pub)
|
|
// A peer frontend may hold a stale entry for this model even though THIS
|
|
// frontend never cached it, so the invalidate MUST be broadcast for
|
|
// cross-frontend coherence. The local drop must still not intern a tree.
|
|
s.Invalidate("never-cached", rk("A", 0))
|
|
Expect(pub.published).To(HaveLen(1))
|
|
ev := pub.published[0].(messaging.PrefixCacheInvalidateEvent)
|
|
Expect(ev.NodeID).To(Equal("A"))
|
|
Expect(ev.Replica).To(Equal(0))
|
|
Expect(idx.TreeCountForTest()).To(Equal(0))
|
|
})
|
|
|
|
It("broadcasts an invalidate for a cached replica too", func() {
|
|
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
|
pub := &fakePub{}
|
|
s := prefixcache.NewSync(idx, pub)
|
|
s.Observe("m", []uint64{1, 2}, rk("A", 0), t0) // creates the tree (also publishes observe)
|
|
pub.published = nil
|
|
s.Invalidate("m", rk("A", 0))
|
|
Expect(pub.published).To(HaveLen(1))
|
|
Expect(pub.published[0]).To(BeAssignableToTypeOf(messaging.PrefixCacheInvalidateEvent{}))
|
|
})
|
|
|
|
It("broadcasts a node-wide invalidate with a negative replica", func() {
|
|
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
|
pub := &fakePub{}
|
|
s := prefixcache.NewSync(idx, pub)
|
|
s.InvalidateNode("m", "A")
|
|
Expect(pub.published).To(HaveLen(1))
|
|
ev := pub.published[0].(messaging.PrefixCacheInvalidateEvent)
|
|
Expect(ev.NodeID).To(Equal("A"))
|
|
Expect(ev.Replica).To(BeNumerically("<", 0))
|
|
})
|
|
|
|
It("applies a peer observe event into the local index with the replica", func() {
|
|
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
|
s := prefixcache.NewSync(idx, &fakePub{})
|
|
s.ApplyObserve(messaging.PrefixCacheObserveEvent{Model: "m", Chain: []uint64{1, 2}, NodeID: "A", Replica: 2}, t0)
|
|
d := idx.Decide("m", []uint64{1, 2}, []prefixcache.ReplicaKey{rk("A", 2)}, t0)
|
|
Expect(d.HasHot).To(BeTrue())
|
|
Expect(d.Hot).To(Equal(rk("A", 2)))
|
|
})
|
|
|
|
It("applies a peer single-replica invalidate", func() {
|
|
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
|
s := prefixcache.NewSync(idx, &fakePub{})
|
|
s.Observe("m", []uint64{1, 2}, rk("A", 0), t0)
|
|
s.Observe("m", []uint64{3, 4}, rk("A", 1), t0)
|
|
s.ApplyInvalidate(messaging.PrefixCacheInvalidateEvent{Model: "m", NodeID: "A", Replica: 0})
|
|
cands := []prefixcache.ReplicaKey{rk("A", 0), rk("A", 1)}
|
|
Expect(idx.Decide("m", []uint64{1, 2}, cands, t0).HasHot).To(BeFalse())
|
|
Expect(idx.Decide("m", []uint64{3, 4}, cands, t0).HasHot).To(BeTrue())
|
|
})
|
|
|
|
It("applies a peer node-wide invalidate when replica is negative", func() {
|
|
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
|
s := prefixcache.NewSync(idx, &fakePub{})
|
|
s.Observe("m", []uint64{1, 2}, rk("A", 0), t0)
|
|
s.Observe("m", []uint64{3, 4}, rk("A", 1), t0)
|
|
s.ApplyInvalidate(messaging.PrefixCacheInvalidateEvent{Model: "m", NodeID: "A", Replica: -1})
|
|
cands := []prefixcache.ReplicaKey{rk("A", 0), rk("A", 1)}
|
|
Expect(idx.Decide("m", []uint64{1, 2}, cands, t0).HasHot).To(BeFalse())
|
|
Expect(idx.Decide("m", []uint64{3, 4}, cands, t0).HasHot).To(BeFalse())
|
|
})
|
|
})
|