mirror of
https://github.com/ollama/ollama.git
synced 2026-07-31 01:48:44 -04:00
556245843a
A draft cache pairs each slot with the token that follows it, so the deepest stored pair always names one token past what a prefix match can verify - at generation end, the sampled-but-never-committed final token. Restoring at the match point reuses that pair blind: a stop token stripped from the next prompt, or any divergence at the boundary, leaves it stale, and pairing never rewrites below the resume position, quietly lowering draft acceptance. Key the trie by token pairs instead: the key for offset i packs (token i, token i+1), so matching k keys verifies k+1 tokens and every match is a valid restore point. A pair is reused only if the token it names matched, and prefill re-evaluates the boundary token, rebuilding its pair with the token that actually follows. A token gets a key only once its successor is recorded, so endings record the final sampled token - never forwarded - and the trie stays level with the caches. Without a look-ahead the keys are the tokens and behavior is unchanged. The recorded tokens' slice bounds used to reject state past them for free; close now checks the invariant against the stored keys directly. The test harness rests requests the way the pipeline does - the deepest recorded token never enters the caches.
1191 lines
38 KiB
Go
1191 lines
38 KiB
Go
package mlxrunner
|
|
|
|
import (
|
|
"slices"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/ollama/ollama/x/mlxrunner/cache"
|
|
"github.com/ollama/ollama/x/mlxrunner/mlx"
|
|
)
|
|
|
|
// snapshotTracker records every fakeSnapshot created and every Close() call
|
|
// so tests can detect leaked (created but never closed) or double-closed snapshots.
|
|
type snapshotTracker struct {
|
|
all []*fakeSnapshot
|
|
}
|
|
|
|
func (tr *snapshotTracker) track(s *fakeSnapshot) {
|
|
if s == nil {
|
|
return
|
|
}
|
|
s.tracker = tr
|
|
tr.all = append(tr.all, s)
|
|
}
|
|
|
|
// Fake caches that store actual token sequences so tests can verify the right
|
|
// data was restored, not just the right offset.
|
|
|
|
// fakePending mirrors the production pendingSnapshots capture machinery for the
|
|
// fakes: it schedules offsets, captures as feed crosses each one (edge-local,
|
|
// via a running base cursor), and returns the captures in scheduled order.
|
|
// capture is the owning fake's snapshot-at-an-offset function.
|
|
type fakePending struct {
|
|
offsets []int
|
|
captured []cache.Snapshot
|
|
base int
|
|
}
|
|
|
|
func (p *fakePending) prepare(currentOffset int, offsets []int) {
|
|
p.offsets = slices.Clone(offsets)
|
|
p.captured = make([]cache.Snapshot, len(offsets))
|
|
p.base = currentOffset
|
|
}
|
|
|
|
func (p *fakePending) take() []cache.Snapshot {
|
|
out := p.captured
|
|
p.offsets, p.captured = nil, nil
|
|
return out
|
|
}
|
|
|
|
// feedCapturing advances the cache from start over the fed tokens, capturing at
|
|
// every scheduled offset crossed (including start and end). capture(from,
|
|
// reached) produces the snapshot for the edge ending at reached;
|
|
// advance(tokens) appends a token segment to the live state. Segment boundaries
|
|
// fall at scheduled offsets.
|
|
func (p *fakePending) feedCapturing(start int, tokens []int32, capture func(from, reached int) cache.Snapshot, advance func([]int32)) {
|
|
end := start + len(tokens)
|
|
captureAt := func(reached int) {
|
|
for i, o := range p.offsets {
|
|
if p.captured[i] == nil && o == reached {
|
|
p.captured[i] = capture(p.base, reached)
|
|
}
|
|
}
|
|
p.base = reached
|
|
}
|
|
|
|
if len(p.offsets) == 0 {
|
|
advance(tokens)
|
|
return
|
|
}
|
|
|
|
captureAt(start)
|
|
prev := 0
|
|
for cut := start + 1; cut < end; cut++ {
|
|
if !slices.Contains(p.offsets, cut) {
|
|
continue
|
|
}
|
|
advance(tokens[prev : cut-start])
|
|
captureAt(cut)
|
|
prev = cut - start
|
|
}
|
|
advance(tokens[prev:])
|
|
captureAt(end)
|
|
}
|
|
|
|
// fakeSnapshot stores a copy of the token sub-sequence it covers.
|
|
type fakeSnapshot struct {
|
|
tokens []int32
|
|
from, to int
|
|
byteSize int // configurable for eviction tests
|
|
|
|
tracker *snapshotTracker
|
|
closeCount int
|
|
|
|
onMaterialize func(delta int)
|
|
}
|
|
|
|
func (s *fakeSnapshot) Size() int { return s.byteSize }
|
|
func (s *fakeSnapshot) SetMaterializeHook(fn func(delta int)) { s.onMaterialize = fn }
|
|
func (s *fakeSnapshot) Close() { s.closeCount++ }
|
|
|
|
// materialize simulates a lazy snapshot copying out: grow byteSize by delta
|
|
// and fire the trie's hook if one is attached. Tests use this to verify the
|
|
// trie's counter responds to materialization events.
|
|
func (s *fakeSnapshot) materialize(delta int) {
|
|
s.byteSize += delta
|
|
if s.onMaterialize != nil {
|
|
s.onMaterialize(delta)
|
|
s.onMaterialize = nil
|
|
}
|
|
}
|
|
|
|
// fakeRewindableCache tracks the full token sequence and supports
|
|
// arbitrary rewind via Restore(nil, target).
|
|
type fakeRewindableCache struct {
|
|
tokens []int32
|
|
tracker *snapshotTracker
|
|
pending fakePending
|
|
}
|
|
|
|
func (c *fakeRewindableCache) feed(tokens []int32) {
|
|
c.pending.feedCapturing(len(c.tokens), tokens,
|
|
func(from, reached int) cache.Snapshot { return c.Snapshot(from) },
|
|
func(seg []int32) { c.tokens = append(c.tokens, seg...) })
|
|
}
|
|
|
|
func (c *fakeRewindableCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
|
|
return nil, nil
|
|
}
|
|
func (c *fakeRewindableCache) State() []*mlx.Array { return nil }
|
|
func (c *fakeRewindableCache) Offset() int { return len(c.tokens) }
|
|
|
|
func (c *fakeRewindableCache) Free() {
|
|
c.tokens = nil
|
|
}
|
|
|
|
func (c *fakeRewindableCache) Snapshot(fromOffset int) cache.Snapshot {
|
|
if fromOffset >= len(c.tokens) {
|
|
return nil
|
|
}
|
|
from := fromOffset
|
|
if from < 0 {
|
|
from = 0
|
|
}
|
|
s := &fakeSnapshot{
|
|
tokens: slices.Clone(c.tokens[from:]),
|
|
from: from,
|
|
to: len(c.tokens),
|
|
}
|
|
c.tracker.track(s)
|
|
return s
|
|
}
|
|
|
|
func (c *fakeRewindableCache) Restore(snapshot cache.Snapshot, target int) bool {
|
|
if target < 0 {
|
|
return false
|
|
}
|
|
|
|
if snapshot == nil {
|
|
if target > len(c.tokens) {
|
|
return false
|
|
}
|
|
c.tokens = c.tokens[:target]
|
|
return true
|
|
}
|
|
s := snapshot.(*fakeSnapshot)
|
|
if target > s.to || len(c.tokens) < s.from {
|
|
return false
|
|
}
|
|
c.tokens = append(c.tokens[:s.from], s.tokens...)
|
|
if target < len(c.tokens) {
|
|
c.tokens = c.tokens[:target]
|
|
}
|
|
return true
|
|
}
|
|
|
|
func (c *fakeRewindableCache) Merge(parent, child cache.Snapshot) cache.Snapshot {
|
|
if parent == nil || child == nil {
|
|
if parent != nil {
|
|
parent.Close()
|
|
}
|
|
if child != nil {
|
|
child.Close()
|
|
}
|
|
return nil
|
|
}
|
|
p := parent.(*fakeSnapshot)
|
|
ch := child.(*fakeSnapshot)
|
|
merged := make([]int32, len(p.tokens)+len(ch.tokens))
|
|
copy(merged, p.tokens)
|
|
copy(merged[len(p.tokens):], ch.tokens)
|
|
s := &fakeSnapshot{
|
|
tokens: merged,
|
|
from: p.from,
|
|
to: ch.to,
|
|
byteSize: p.byteSize + ch.byteSize,
|
|
}
|
|
c.tracker.track(s)
|
|
p.Close()
|
|
ch.Close()
|
|
return s
|
|
}
|
|
|
|
func (c *fakeRewindableCache) Split(snapshot cache.Snapshot, at int) (cache.Snapshot, cache.Snapshot) {
|
|
if snapshot == nil {
|
|
return nil, nil
|
|
}
|
|
s := snapshot.(*fakeSnapshot)
|
|
relAt := at - s.from
|
|
if relAt <= 0 {
|
|
return nil, snapshot
|
|
}
|
|
if relAt >= len(s.tokens) {
|
|
return snapshot, nil
|
|
}
|
|
p := &fakeSnapshot{
|
|
tokens: slices.Clone(s.tokens[:relAt]),
|
|
from: s.from,
|
|
to: at,
|
|
byteSize: s.byteSize,
|
|
}
|
|
ch := &fakeSnapshot{
|
|
tokens: slices.Clone(s.tokens[relAt:]),
|
|
from: at,
|
|
to: s.to,
|
|
byteSize: s.byteSize,
|
|
}
|
|
c.tracker.track(p)
|
|
c.tracker.track(ch)
|
|
s.Close()
|
|
return p, ch
|
|
}
|
|
|
|
func (c *fakeRewindableCache) PrepareSnapshots(offsets []int) {
|
|
c.pending.prepare(len(c.tokens), offsets)
|
|
}
|
|
func (c *fakeRewindableCache) TakeSnapshots() []cache.Snapshot { return c.pending.take() }
|
|
|
|
// fakeSlidingWindowCache models RotatingKVCache semantics: stores the full
|
|
// token sequence but only the trailing maxSize tokens are "live" in the window.
|
|
// Once the window fills, live rewind is impossible without a snapshot.
|
|
type fakeSlidingWindowCache struct {
|
|
tokens []int32
|
|
maxSize int
|
|
tracker *snapshotTracker
|
|
pending fakePending
|
|
}
|
|
|
|
func (c *fakeSlidingWindowCache) feed(tokens []int32) {
|
|
c.pending.feedCapturing(len(c.tokens), tokens,
|
|
func(from, reached int) cache.Snapshot { return c.Snapshot(0) },
|
|
func(seg []int32) { c.tokens = append(c.tokens, seg...) })
|
|
}
|
|
|
|
func (c *fakeSlidingWindowCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
|
|
return nil, nil
|
|
}
|
|
func (c *fakeSlidingWindowCache) State() []*mlx.Array { return nil }
|
|
func (c *fakeSlidingWindowCache) Offset() int { return len(c.tokens) }
|
|
|
|
func (c *fakeSlidingWindowCache) Free() {
|
|
c.tokens = nil
|
|
}
|
|
|
|
func (c *fakeSlidingWindowCache) Snapshot(fromOffset int) cache.Snapshot {
|
|
if len(c.tokens) == 0 || len(c.tokens) <= fromOffset {
|
|
return nil
|
|
}
|
|
// Snapshot captures the full window state (like RotatingKVCache.Snapshot).
|
|
s := &fakeSnapshot{
|
|
tokens: slices.Clone(c.tokens),
|
|
from: 0,
|
|
to: len(c.tokens),
|
|
}
|
|
c.tracker.track(s)
|
|
return s
|
|
}
|
|
|
|
func (c *fakeSlidingWindowCache) Restore(snapshot cache.Snapshot, target int) bool {
|
|
if target < 0 {
|
|
return false
|
|
}
|
|
|
|
if snapshot == nil {
|
|
if target >= len(c.tokens) {
|
|
return target == len(c.tokens)
|
|
}
|
|
// Live rewind only works when buffer hasn't filled (offset <= maxSize).
|
|
if len(c.tokens) > c.maxSize {
|
|
return false
|
|
}
|
|
c.tokens = c.tokens[:target]
|
|
return true
|
|
}
|
|
s := snapshot.(*fakeSnapshot)
|
|
if target > s.to {
|
|
return false
|
|
}
|
|
// Reject if clamping would leave an incomplete window
|
|
// (matches RotatingKVCache behavior).
|
|
if target < s.to && s.to > c.maxSize {
|
|
return false
|
|
}
|
|
c.tokens = slices.Clone(s.tokens)
|
|
if target < len(c.tokens) {
|
|
c.tokens = c.tokens[:target]
|
|
}
|
|
return true
|
|
}
|
|
|
|
func (c *fakeSlidingWindowCache) Merge(parent, child cache.Snapshot) cache.Snapshot {
|
|
// Child supersedes parent for sliding window (full window state).
|
|
if parent != nil {
|
|
parent.Close()
|
|
}
|
|
return child
|
|
}
|
|
|
|
func (c *fakeSlidingWindowCache) Split(snapshot cache.Snapshot, at int) (cache.Snapshot, cache.Snapshot) {
|
|
// Can't split a ring buffer at an arbitrary point.
|
|
return nil, snapshot
|
|
}
|
|
|
|
func (c *fakeSlidingWindowCache) PrepareSnapshots(offsets []int) {
|
|
c.pending.prepare(len(c.tokens), offsets)
|
|
}
|
|
func (c *fakeSlidingWindowCache) TakeSnapshots() []cache.Snapshot { return c.pending.take() }
|
|
|
|
// fakeRecurrentCache models RecurrentCache semantics: stores tokens
|
|
// but cannot rewind without a snapshot.
|
|
type fakeRecurrentCache struct {
|
|
tokens []int32
|
|
tracker *snapshotTracker
|
|
pending fakePending
|
|
}
|
|
|
|
func (c *fakeRecurrentCache) feed(tokens []int32) {
|
|
c.pending.feedCapturing(len(c.tokens), tokens,
|
|
func(from, reached int) cache.Snapshot { return c.Snapshot(0) },
|
|
func(seg []int32) { c.tokens = append(c.tokens, seg...) })
|
|
}
|
|
|
|
func (c *fakeRecurrentCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
|
|
return nil, nil
|
|
}
|
|
func (c *fakeRecurrentCache) State() []*mlx.Array { return nil }
|
|
func (c *fakeRecurrentCache) Offset() int { return len(c.tokens) }
|
|
|
|
func (c *fakeRecurrentCache) Free() {
|
|
c.tokens = nil
|
|
}
|
|
|
|
func (c *fakeRecurrentCache) Snapshot(fromOffset int) cache.Snapshot {
|
|
// Recurrent state is cumulative; snapshot captures the full state.
|
|
if len(c.tokens) == 0 {
|
|
return nil
|
|
}
|
|
s := &fakeSnapshot{
|
|
tokens: slices.Clone(c.tokens),
|
|
from: 0,
|
|
to: len(c.tokens),
|
|
}
|
|
c.tracker.track(s)
|
|
return s
|
|
}
|
|
|
|
func (c *fakeRecurrentCache) Restore(snapshot cache.Snapshot, target int) bool {
|
|
if snapshot == nil {
|
|
return target == len(c.tokens) // can only no-op
|
|
}
|
|
s := snapshot.(*fakeSnapshot)
|
|
if target != s.to {
|
|
return false // cumulative state requires exact match
|
|
}
|
|
c.tokens = slices.Clone(s.tokens)
|
|
return true
|
|
}
|
|
|
|
func (c *fakeRecurrentCache) Merge(parent, child cache.Snapshot) cache.Snapshot {
|
|
// Child supersedes parent for cumulative state.
|
|
if parent != nil {
|
|
parent.Close()
|
|
}
|
|
return child
|
|
}
|
|
|
|
func (c *fakeRecurrentCache) Split(snapshot cache.Snapshot, at int) (cache.Snapshot, cache.Snapshot) {
|
|
return nil, snapshot // can't split cumulative state
|
|
}
|
|
|
|
func (c *fakeRecurrentCache) PrepareSnapshots(offsets []int) {
|
|
c.pending.prepare(len(c.tokens), offsets)
|
|
}
|
|
func (c *fakeRecurrentCache) TakeSnapshots() []cache.Snapshot { return c.pending.take() }
|
|
|
|
type feedableCache interface {
|
|
cache.Cache
|
|
feed(tokens []int32)
|
|
}
|
|
|
|
// testEnv encapsulates a prefixCache and its fake caches for a test scenario.
|
|
type testEnv struct {
|
|
pc *prefixCache
|
|
caches []cache.Cache // typed references for assertions
|
|
tracker *snapshotTracker
|
|
rewindable bool // true when all caches support arbitrary Restore(nil, target)
|
|
}
|
|
|
|
// newTransformerEnv creates a test environment with a single rewindable cache
|
|
// (pure transformer model).
|
|
func newTransformerEnv() *testEnv {
|
|
tracker := &snapshotTracker{}
|
|
caches := []cache.Cache{&fakeRewindableCache{tracker: tracker}}
|
|
return &testEnv{
|
|
pc: &prefixCache{caches: caches},
|
|
caches: caches,
|
|
tracker: tracker,
|
|
rewindable: true,
|
|
}
|
|
}
|
|
|
|
// newSlidingWindowEnv creates a test environment with one rewindable cache and
|
|
// one sliding window cache (Mistral-style architecture). The sliding window
|
|
// maxSize is set small enough that test sequences fill it, making
|
|
// Restore(nil, target) fail — the same behavior as production models where
|
|
// the window fills after a few turns.
|
|
func newSlidingWindowEnv() *testEnv {
|
|
tr := &snapshotTracker{}
|
|
rc := &fakeRewindableCache{tracker: tr}
|
|
sw := &fakeSlidingWindowCache{maxSize: 4, tracker: tr}
|
|
caches := []cache.Cache{rc, sw}
|
|
return &testEnv{
|
|
pc: &prefixCache{caches: caches},
|
|
caches: caches,
|
|
tracker: tr,
|
|
rewindable: false,
|
|
}
|
|
}
|
|
|
|
// newRecurrentEnv creates a test environment with one rewindable cache and one
|
|
// non-rewindable cache (Jamba-style architecture).
|
|
func newRecurrentEnv() *testEnv {
|
|
tr := &snapshotTracker{}
|
|
rc := &fakeRewindableCache{tracker: tr}
|
|
nrc := &fakeRecurrentCache{tracker: tr}
|
|
caches := []cache.Cache{rc, nrc}
|
|
return &testEnv{
|
|
pc: &prefixCache{caches: caches},
|
|
caches: caches,
|
|
tracker: tr,
|
|
rewindable: false,
|
|
}
|
|
}
|
|
|
|
// assertAllTokens checks that every cache in the environment contains exactly
|
|
// the expected token sequence.
|
|
func (e *testEnv) assertAllTokens(t *testing.T, label string, expected []int32) {
|
|
t.Helper()
|
|
for i, c := range e.caches {
|
|
assertTokens(t, label, c, expected)
|
|
// Verify all caches report the same offset.
|
|
if i > 0 && c.Offset() != e.caches[0].Offset() {
|
|
t.Errorf("%s: cache %d offset=%d != cache 0 offset=%d",
|
|
label, i, c.Offset(), e.caches[0].Offset())
|
|
}
|
|
}
|
|
}
|
|
|
|
// simulateRequest mirrors the production pipeline lifecycle:
|
|
// begin -> schedule snapshots -> prefill in one pass -> attach snapshots -> generate -> close
|
|
|
|
type requestResult struct {
|
|
remaining []int32
|
|
pendingSnapshots int
|
|
}
|
|
|
|
// simulateRequest runs a request through the harness. If userSnapshotAt > 0,
|
|
// a user snapshot is requested at that offset during prefill.
|
|
func simulateRequest(t *testing.T, pc *prefixCache, inputs, generated []int32, userSnapshotAt ...int) requestResult {
|
|
t.Helper()
|
|
|
|
session := pc.begin(inputs)
|
|
var snapshotOffsets []int
|
|
for _, at := range userSnapshotAt {
|
|
if at > 0 {
|
|
snapshotOffsets = append(snapshotOffsets, at)
|
|
}
|
|
}
|
|
|
|
result := requestResult{
|
|
remaining: slices.Clone(session.remaining),
|
|
pendingSnapshots: len(session.pendingSnapshots),
|
|
}
|
|
|
|
assertCacheOffsetAlignment(t, pc, "after begin")
|
|
|
|
baseOffset := pc.minCacheOffset()
|
|
seed := len(inputs) - 1
|
|
|
|
// Prefill: schedule the pending snapshots, feed the prompt up to the
|
|
// seed token in one pass (the caches self-segment at the scheduled
|
|
// offsets), then attach the captures to the trie.
|
|
session.schedulePrefillSnapshots(snapshotOffsets)
|
|
if baseOffset < seed {
|
|
feedAll(pc.caches, inputs[baseOffset:seed])
|
|
}
|
|
session.attachPrefillSnapshots()
|
|
|
|
assertCacheOffsetAlignment(t, pc, "after prefill")
|
|
|
|
// Decode: feed the seed and all but the last generated token.
|
|
if len(generated) > 0 {
|
|
session.outputs = generated
|
|
feedAll(pc.caches, inputs[seed:])
|
|
feedAll(pc.caches, generated[:len(generated)-1])
|
|
}
|
|
|
|
assertCacheOffsetAlignment(t, pc, "before close")
|
|
session.close()
|
|
return result
|
|
}
|
|
|
|
func feedAll(caches []cache.Cache, tokens []int32) {
|
|
for _, c := range caches {
|
|
if fc, ok := c.(feedableCache); ok {
|
|
fc.feed(tokens)
|
|
}
|
|
}
|
|
}
|
|
|
|
// assertCacheOffsetAlignment verifies all caches report the same offset.
|
|
func assertCacheOffsetAlignment(t *testing.T, pc *prefixCache, label string) {
|
|
t.Helper()
|
|
if len(pc.caches) < 2 {
|
|
return
|
|
}
|
|
expected := pc.caches[0].Offset()
|
|
for i := 1; i < len(pc.caches); i++ {
|
|
if got := pc.caches[i].Offset(); got != expected {
|
|
t.Errorf("%s: cache %d offset=%d != cache 0 offset=%d", label, i, got, expected)
|
|
}
|
|
}
|
|
}
|
|
|
|
// assertTokens checks that a feedable cache contains the expected token sequence.
|
|
// For sliding window caches, only the trailing maxSize tokens are checked.
|
|
func assertTokens(t *testing.T, label string, c cache.Cache, expected []int32) {
|
|
t.Helper()
|
|
switch fc := c.(type) {
|
|
case *fakeRewindableCache:
|
|
if !slices.Equal(fc.tokens, expected) {
|
|
t.Errorf("%s: rewindable tokens = %v, want %v", label, fc.tokens, expected)
|
|
}
|
|
case *fakeSlidingWindowCache:
|
|
// Sliding window stores full history but only trailing maxSize are live.
|
|
// Verify the full token sequence matches (the window semantics are
|
|
// enforced by Snapshot/Restore, not by the token log).
|
|
if !slices.Equal(fc.tokens, expected) {
|
|
t.Errorf("%s: sliding window tokens = %v, want %v", label, fc.tokens, expected)
|
|
}
|
|
case *fakeRecurrentCache:
|
|
if !slices.Equal(fc.tokens, expected) {
|
|
t.Errorf("%s: non-rewindable tokens = %v, want %v", label, fc.tokens, expected)
|
|
}
|
|
default:
|
|
t.Fatalf("%s: unknown cache type %T", label, c)
|
|
}
|
|
}
|
|
|
|
// checkTrieInvariants walks the trie and checks structural invariants.
|
|
func checkTrieInvariants(t *testing.T, root *trieNode) {
|
|
t.Helper()
|
|
walkNodes(root, func(n *trieNode) bool {
|
|
if n.parent != nil {
|
|
if n.startOffset() != n.parent.endOffset {
|
|
t.Errorf("node [%d,%d): startOffset %d != parent endOffset %d",
|
|
n.startOffset(), n.endOffset, n.startOffset(), n.parent.endOffset)
|
|
}
|
|
}
|
|
if len(n.tokens) != n.endOffset-n.startOffset() {
|
|
t.Errorf("node [%d,%d): token count %d != offset span %d",
|
|
n.startOffset(), n.endOffset, len(n.tokens), n.endOffset-n.startOffset())
|
|
}
|
|
for _, c := range n.children {
|
|
if c.parent != n {
|
|
t.Errorf("child [%d,%d) parent mismatch", c.startOffset(), c.endOffset)
|
|
}
|
|
}
|
|
// No two siblings should start with the same token.
|
|
seen := make(map[trieKey]bool)
|
|
for _, c := range n.children {
|
|
if len(c.tokens) > 0 {
|
|
first := c.tokens[0]
|
|
if seen[first] {
|
|
t.Errorf("node [%d,%d): duplicate sibling first token %d",
|
|
n.startOffset(), n.endOffset, first)
|
|
}
|
|
seen[first] = true
|
|
}
|
|
}
|
|
return true
|
|
})
|
|
}
|
|
|
|
// checkSnapshotLeaks verifies that every tracked snapshot is either still live
|
|
// in the trie (closeCount == 0) or has been closed exactly once. It reports
|
|
// leaked snapshots (not in trie, never closed) and double-closes.
|
|
func checkSnapshotLeaks(t *testing.T, tracker *snapshotTracker, root *trieNode) {
|
|
t.Helper()
|
|
if tracker == nil {
|
|
return
|
|
}
|
|
|
|
// Collect all live snapshots still referenced by trie nodes.
|
|
live := make(map[*fakeSnapshot]bool)
|
|
walkNodes(root, func(n *trieNode) bool {
|
|
for _, s := range n.snapshots {
|
|
if s != nil {
|
|
if fs, ok := s.(*fakeSnapshot); ok {
|
|
live[fs] = true
|
|
}
|
|
}
|
|
}
|
|
return true
|
|
})
|
|
|
|
for i, s := range tracker.all {
|
|
if live[s] {
|
|
if s.closeCount != 0 {
|
|
t.Errorf("snapshot #%d [%d,%d) is still in trie but was closed %d time(s)",
|
|
i, s.from, s.to, s.closeCount)
|
|
}
|
|
} else {
|
|
if s.closeCount == 0 {
|
|
t.Errorf("snapshot #%d [%d,%d) leaked: created but never closed and not in trie",
|
|
i, s.from, s.to)
|
|
} else if s.closeCount > 1 {
|
|
t.Errorf("snapshot #%d [%d,%d) double-closed: closed %d times",
|
|
i, s.from, s.to, s.closeCount)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// forEachEnv runs fn as subtests for three realistic model configurations —
|
|
// pure transformer, transformer + sliding window (Mistral-style), and
|
|
// transformer + recurrent (Jamba-style) — each with and without a draft
|
|
// look-ahead. Leak checking runs automatically at the end of each subtest.
|
|
func forEachEnv(t *testing.T, fn func(t *testing.T, env *testEnv)) {
|
|
t.Helper()
|
|
envs := []struct {
|
|
name string
|
|
make func() *testEnv
|
|
}{
|
|
{"Transformer", newTransformerEnv},
|
|
{"SlidingWindow", newSlidingWindowEnv},
|
|
{"Recurrent", newRecurrentEnv},
|
|
}
|
|
for _, e := range envs {
|
|
for _, lookahead := range []int{0, 1} {
|
|
name := e.name
|
|
if lookahead > 0 {
|
|
name += "+Lookahead"
|
|
}
|
|
t.Run(name, func(t *testing.T) {
|
|
env := e.make()
|
|
env.pc.draftLookahead = lookahead
|
|
t.Cleanup(func() {
|
|
checkSnapshotLeaks(t, env.tracker, env.pc.root)
|
|
})
|
|
fn(t, env)
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestBranchCreationAndReuse exercises the core multi-conversation lifecycle:
|
|
// two conversations share a prefix and diverge, creating a branch point.
|
|
// A third conversation extends the first. Verifies trie structure, cache
|
|
// hit lengths, and that semantic caches contain the correct token sequences.
|
|
func TestBranchCreationAndReuse(t *testing.T) {
|
|
forEachEnv(t, func(t *testing.T, env *testEnv) {
|
|
pc := env.pc
|
|
|
|
// Request A: [1,2,3,4,5,6,7,8] + generate [20,21] — full miss.
|
|
resA := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 6, 7, 8}, []int32{20, 21})
|
|
if len(resA.remaining) != 8 {
|
|
t.Fatalf("A: remaining = %d, want 8 (full miss)", len(resA.remaining))
|
|
}
|
|
env.assertAllTokens(t, "after A", []int32{1, 2, 3, 4, 5, 6, 7, 8, 20})
|
|
|
|
// Verify trie was populated by close(): everything in the caches
|
|
// is findable; the last generated token is not.
|
|
seqA := []int32{1, 2, 3, 4, 5, 6, 7, 8, 20, 21}
|
|
_, mA := findBestMatch(pc.root, pc.key(seqA))
|
|
if want := len(seqA) - 1; mA != want {
|
|
t.Fatalf("A findable: expected %d matched, got %d", want, mA)
|
|
}
|
|
|
|
// Request B: [1,2,3,4,5,10,11,12] — shares 5-token prefix with A.
|
|
// For rewindable caches, switchToPath rewinds to the match point
|
|
// (one below the shared tokens with a look-ahead) so only the suffix
|
|
// needs evaluation. For non-rewindable caches (RecurrentCache), the
|
|
// rewind fails and freeAll fires.
|
|
resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12}, []int32{30, 31})
|
|
if env.rewindable {
|
|
if resB.pendingSnapshots != 0 {
|
|
t.Fatalf("B: pendingSnapshots = %d, want 0 (rewind succeeded)", resB.pendingSnapshots)
|
|
}
|
|
if want := 3 + pc.draftLookahead; len(resB.remaining) != want {
|
|
t.Fatalf("B: remaining = %d, want %d (rewind to match point)", len(resB.remaining), want)
|
|
}
|
|
} else {
|
|
if resB.pendingSnapshots != 1 {
|
|
t.Fatalf("B: pendingSnapshots = %d, want 1", resB.pendingSnapshots)
|
|
}
|
|
if len(resB.remaining) != 8 {
|
|
t.Fatalf("B: remaining = %d, want 8 (freeAll fallback)", len(resB.remaining))
|
|
}
|
|
}
|
|
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 10, 11, 12, 30})
|
|
|
|
// Both A and B should be findable in the trie.
|
|
_, mA2 := findBestMatch(pc.root, pc.key(seqA))
|
|
if mA2 < 5 {
|
|
t.Fatalf("A still findable: expected >= 5 matched, got %d", mA2)
|
|
}
|
|
_, mB := findBestMatch(pc.root, pc.key([]int32{1, 2, 3, 4, 5, 10, 11, 12, 30, 31}))
|
|
if mB < 5 {
|
|
t.Fatalf("B findable: expected >= 5 matched, got %d", mB)
|
|
}
|
|
|
|
// Request C: [1,2,3,4,5,6,7,8,40,41] — extends A's prefix.
|
|
// Should get a cache hit for the shared prefix.
|
|
resC := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 6, 7, 8, 40, 41}, nil)
|
|
if len(resC.remaining) >= 10 {
|
|
t.Fatalf("C: remaining = %d, want < 10 (should get cache hit)", len(resC.remaining))
|
|
}
|
|
env.assertAllTokens(t, "after C", []int32{1, 2, 3, 4, 5, 6, 7, 8, 40})
|
|
|
|
checkTrieInvariants(t, pc.root)
|
|
})
|
|
}
|
|
|
|
// TestExactMatchSeedBehavior verifies the holdback mechanism: when the exact
|
|
// same prompt is requested twice, the cache does not overclaim cached work.
|
|
// The last token must be re-evaluated to seed generation.
|
|
func TestExactMatchSeedBehavior(t *testing.T) {
|
|
forEachEnv(t, func(t *testing.T, env *testEnv) {
|
|
pc := env.pc
|
|
|
|
// Request A: first time.
|
|
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11})
|
|
|
|
// Request B: identical prompt. Holdback means matched=4, partial in
|
|
// the 5-token edge. For rewindable caches, switchToPath rewinds to
|
|
// offset 4, so only the held-back token needs re-evaluation. For
|
|
// non-rewindable caches, the rewind fails and freeAll fires.
|
|
resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{20, 21})
|
|
if env.rewindable {
|
|
if len(resB.remaining) != 1 {
|
|
t.Fatalf("B: remaining = %d, want 1 (rewind to holdback point)", len(resB.remaining))
|
|
}
|
|
if resB.pendingSnapshots != 0 {
|
|
t.Fatalf("B: pendingSnapshots = %d, want 0 (rewind succeeded)", resB.pendingSnapshots)
|
|
}
|
|
} else {
|
|
if len(resB.remaining) != 5 {
|
|
t.Fatalf("B: remaining = %d, want 5 (freeAll fallback)", len(resB.remaining))
|
|
}
|
|
if resB.pendingSnapshots != 1 {
|
|
t.Fatalf("B: pendingSnapshots = %d, want 1", resB.pendingSnapshots)
|
|
}
|
|
}
|
|
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 20})
|
|
|
|
checkTrieInvariants(t, pc.root)
|
|
})
|
|
}
|
|
|
|
// TestConversationResumption tests the most common pattern: user sends a message,
|
|
// gets a response, then sends a follow-up. The follow-up should reuse the cached
|
|
// prefix (system prompt + first turn + assistant response).
|
|
func TestConversationResumption(t *testing.T) {
|
|
forEachEnv(t, func(t *testing.T, env *testEnv) {
|
|
pc := env.pc
|
|
|
|
// Turn 1: system prompt + user message, assistant generates response.
|
|
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11, 12})
|
|
env.assertAllTokens(t, "turn 1", []int32{1, 2, 3, 4, 5, 10, 11})
|
|
|
|
// Turn 2: full history + new user message; the match lands exactly
|
|
// at the caches' offset, so even exact-offset state is reused.
|
|
resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21}, []int32{30})
|
|
if len(resB.remaining) > 5 {
|
|
t.Fatalf("turn 2: remaining = %d, want <= 5 (should reuse most of history)", len(resB.remaining))
|
|
}
|
|
env.assertAllTokens(t, "turn 2", []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21})
|
|
|
|
// Turn 3: even longer history.
|
|
resC := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21, 30, 40, 41}, nil)
|
|
if len(resC.remaining) > 5 {
|
|
t.Fatalf("turn 3: remaining = %d, want <= 5", len(resC.remaining))
|
|
}
|
|
env.assertAllTokens(t, "turn 3", []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21, 30, 40})
|
|
|
|
checkTrieInvariants(t, pc.root)
|
|
})
|
|
}
|
|
|
|
// TestEvictionPreservesActiveConversations creates multiple conversations sharing
|
|
// a system prompt, triggers eviction via large snapshot sizes, and verifies the
|
|
// active path and shared prefix survive while memory stays bounded.
|
|
func TestEvictionPreservesActiveConversations(t *testing.T) {
|
|
forEachEnv(t, func(t *testing.T, env *testEnv) {
|
|
pc := env.pc
|
|
systemPrompt := []int32{1, 2, 3, 4, 5}
|
|
|
|
// Create 5 conversations with unique suffixes.
|
|
for i := range 5 {
|
|
suffix := []int32{int32(100 + i*10), int32(101 + i*10), int32(102 + i*10)}
|
|
inputs := append(slices.Clone(systemPrompt), suffix...)
|
|
simulateRequest(t, pc, inputs, []int32{int32(200 + i)})
|
|
}
|
|
|
|
// Inflate snapshot sizes to trigger eviction.
|
|
walkNodes(pc.root, func(n *trieNode) bool {
|
|
if !n.hasSnapshots() {
|
|
return true
|
|
}
|
|
snaps := make([]cache.Snapshot, len(n.snapshots))
|
|
for i, s := range n.snapshots {
|
|
if s != nil {
|
|
snaps[i] = &fakeSnapshot{byteSize: 2 * 1024 * 1024 * 1024} // 2 GiB per snapshot
|
|
}
|
|
}
|
|
n.setSnapshots(snaps, &pc.pagedOutBytes)
|
|
return true
|
|
})
|
|
|
|
// Run eviction.
|
|
pc.enforceEvictionPolicy()
|
|
|
|
// Memory should be within limits.
|
|
if pc.pagedOutBytes > maxPagedOutBytes {
|
|
t.Fatalf("pagedOutBytes = %d, want <= %d", pc.pagedOutBytes, maxPagedOutBytes)
|
|
}
|
|
|
|
// Active path should be untouched.
|
|
if len(pc.activePath) < 2 {
|
|
t.Fatalf("activePath should have >= 2 nodes, got %d", len(pc.activePath))
|
|
}
|
|
|
|
// System prompt prefix should still be findable (multi-child
|
|
// branch points are protected from eviction entirely).
|
|
_, matched := findBestMatch(pc.root, pc.key(systemPrompt))
|
|
if want := len(pc.key(systemPrompt)); matched < want {
|
|
t.Fatalf("system prompt match = %d, want %d", matched, want)
|
|
}
|
|
|
|
checkTrieInvariants(t, pc.root)
|
|
})
|
|
}
|
|
|
|
// TestUserSnapshotPreservesRestorePoint verifies that user-created snapshots
|
|
// (snapshot(true)) are exact restore points that resist structural changes:
|
|
// - A user node forces new tokens into a child instead of extending in-place
|
|
// - A prompt diverging at the snapshot resumes there, even for caches that
|
|
// cannot rewind
|
|
// - The snapshot remains restorable after other branches are added
|
|
func TestUserSnapshotPreservesRestorePoint(t *testing.T) {
|
|
forEachEnv(t, func(t *testing.T, env *testEnv) {
|
|
pc := env.pc
|
|
inputs := []int32{1, 2, 3, 4, 5}
|
|
|
|
// Request A: user snapshot at offset 4, then generate.
|
|
simulateRequest(t, pc, inputs, []int32{10, 11}, 4)
|
|
|
|
assertUserNodeExists(t, pc, "after A")
|
|
|
|
// Request B: extends A's prefix. The user node should force tokens
|
|
// into a child rather than extending in-place.
|
|
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 20, 21}, nil)
|
|
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 10, 11, 20})
|
|
assertUserNodeExists(t, pc, "after B")
|
|
|
|
// Request C: diverge at the snapshot — prefill resumes at its capture
|
|
// point, one token lower with a look-ahead (the boundary token
|
|
// re-evaluates to rebuild its draft pair).
|
|
divergeC := append(slices.Clone(inputs[:4]), 30, 31)
|
|
resC := simulateRequest(t, pc, divergeC, []int32{40})
|
|
if want := divergeC[4-pc.draftLookahead:]; !slices.Equal(resC.remaining, want) {
|
|
t.Fatalf("C: remaining = %v, want %v", resC.remaining, want)
|
|
}
|
|
|
|
// Request D: switch back to A's branch — user snapshot still restorable.
|
|
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 20, 21, 50}, nil)
|
|
env.assertAllTokens(t, "back to A", []int32{1, 2, 3, 4, 5, 10, 11, 20, 21})
|
|
|
|
checkTrieInvariants(t, pc.root)
|
|
})
|
|
}
|
|
|
|
// TestUserSnapshotResistsAutoMerge verifies that when a sibling leaf is evicted,
|
|
// a user-marked parent node is not auto-merged with its remaining single child.
|
|
func TestUserSnapshotResistsAutoMerge(t *testing.T) {
|
|
forEachEnv(t, func(t *testing.T, env *testEnv) {
|
|
pc := env.pc
|
|
inputs := []int32{1, 2, 3, 4, 5}
|
|
|
|
// Request A: user snapshot at offset 3, then continue to offset 5.
|
|
simulateRequest(t, pc, inputs, []int32{10}, 3)
|
|
|
|
// Request B: diverges at the user node, creating a second child.
|
|
simulateRequest(t, pc, []int32{1, 2, 3, 6, 7}, []int32{20})
|
|
|
|
userNode := findUserNode(t, pc)
|
|
if len(userNode.children) != 2 {
|
|
t.Fatalf("user node children = %d, want 2", len(userNode.children))
|
|
}
|
|
|
|
// Inflate snapshot sizes and evict. The non-active branch should be
|
|
// evicted, leaving the user node with one child.
|
|
walkNodes(pc.root, func(n *trieNode) bool {
|
|
if !n.hasSnapshots() {
|
|
return true
|
|
}
|
|
snaps := make([]cache.Snapshot, len(n.snapshots))
|
|
for i, s := range n.snapshots {
|
|
if s != nil {
|
|
snaps[i] = &fakeSnapshot{byteSize: 5 * 1024 * 1024 * 1024}
|
|
}
|
|
}
|
|
n.setSnapshots(snaps, &pc.pagedOutBytes)
|
|
return true
|
|
})
|
|
pc.enforceEvictionPolicy()
|
|
|
|
// The user node should still exist (not auto-merged) even with one child.
|
|
assertUserNodeExists(t, pc, "after eviction")
|
|
|
|
checkTrieInvariants(t, pc.root)
|
|
})
|
|
}
|
|
|
|
// TestSnapshotBeyondPrefillSkipped verifies that a snapshot scheduled at an
|
|
// offset the prefill never reaches (prefill leaves one token for decode
|
|
// seeding, so the last token is never written during prefill) is dropped rather
|
|
// than materialized as a trie node claiming tokens the cache never wrote.
|
|
func TestSnapshotBeyondPrefillSkipped(t *testing.T) {
|
|
forEachEnv(t, func(t *testing.T, env *testEnv) {
|
|
pc := env.pc
|
|
inputs := []int32{1, 2, 3, 4, 5}
|
|
|
|
session := pc.begin(inputs)
|
|
// Request a snapshot at 3 and one at len(inputs); captures land at
|
|
// the requested prefix minus the look-ahead.
|
|
session.schedulePrefillSnapshots([]int{3, len(inputs)})
|
|
// Prefill writes all but the final token (mirrors total-processed > 1).
|
|
feedAll(pc.caches, inputs[pc.minCacheOffset():len(inputs)-1])
|
|
session.attachPrefillSnapshots()
|
|
|
|
// The first request became a node at its capture point; nothing may
|
|
// claim offsets the prefill never wrote.
|
|
if at := 3 - pc.draftLookahead; !nodeExistsAtOffset(pc.root, at) {
|
|
t.Errorf("no trie node at capture point %d", at)
|
|
}
|
|
reached := pc.minCacheOffset()
|
|
walkNodes(pc.root, func(n *trieNode) bool {
|
|
if n.endOffset > reached {
|
|
t.Errorf("trie node materialized at unwritten offset %d", n.endOffset)
|
|
}
|
|
return true
|
|
})
|
|
|
|
checkTrieInvariants(t, pc.root)
|
|
})
|
|
}
|
|
|
|
// TestPrefillSnapshotsDiscardedOnCancel mirrors a prefill canceled after the
|
|
// caches captured interior snapshots but before attachPrefillSnapshots ran. The
|
|
// abandoned captures must be released when the session closes; otherwise the
|
|
// next request's PrepareSnapshots overwrites the schedule without closing them,
|
|
// leaking the snapshots (caught by checkSnapshotLeaks in the env cleanup).
|
|
func TestPrefillSnapshotsDiscardedOnCancel(t *testing.T) {
|
|
forEachEnv(t, func(t *testing.T, env *testEnv) {
|
|
pc := env.pc
|
|
inputs := []int32{1, 2, 3, 4, 5}
|
|
|
|
session := pc.begin(inputs)
|
|
session.schedulePrefillSnapshots([]int{3})
|
|
// Cross offset 3 so the caches capture it, then close the session as a
|
|
// canceled prefill would, before the captures are attached to the trie.
|
|
feedAll(pc.caches, inputs[pc.minCacheOffset():3])
|
|
session.close()
|
|
|
|
// close advances the trie over the committed tokens, but the abandoned
|
|
// captures must not be attached as snapshots to any node.
|
|
walkNodes(pc.root, func(n *trieNode) bool {
|
|
if n != pc.root && n.hasSnapshots() {
|
|
t.Errorf("abandoned capture attached as snapshot at offset %d", n.endOffset)
|
|
}
|
|
return true
|
|
})
|
|
|
|
// A second request re-prepares snapshots on the same caches: if the
|
|
// discarded ones were not closed, prepare() orphans them here.
|
|
simulateRequest(t, pc, inputs, nil, 5)
|
|
|
|
checkTrieInvariants(t, pc.root)
|
|
})
|
|
}
|
|
|
|
func nodeExistsAtOffset(root *trieNode, offset int) bool {
|
|
var found bool
|
|
walkNodes(root, func(n *trieNode) bool {
|
|
if n.endOffset == offset && n != root {
|
|
found = true
|
|
}
|
|
return true
|
|
})
|
|
return found
|
|
}
|
|
|
|
func findUserNode(t *testing.T, pc *prefixCache) *trieNode {
|
|
t.Helper()
|
|
var found *trieNode
|
|
walkNodes(pc.root, func(n *trieNode) bool {
|
|
if n.user {
|
|
found = n
|
|
}
|
|
return true
|
|
})
|
|
if found == nil {
|
|
t.Fatal("no user-marked node found")
|
|
}
|
|
return found
|
|
}
|
|
|
|
func assertUserNodeExists(t *testing.T, pc *prefixCache, label string) {
|
|
t.Helper()
|
|
var exists bool
|
|
walkNodes(pc.root, func(n *trieNode) bool {
|
|
if n.user {
|
|
exists = true
|
|
}
|
|
return true
|
|
})
|
|
if !exists {
|
|
t.Fatalf("%s: no user-marked node found", label)
|
|
}
|
|
}
|
|
|
|
// TestBranchSwitchRestoresCorrectState exercises switching back to an older
|
|
// branch after working on a different one, verifying that the restored cache
|
|
// state contains the correct token sequence for both rewindable and
|
|
// non-rewindable caches.
|
|
func TestBranchSwitchRestoresCorrectState(t *testing.T) {
|
|
forEachEnv(t, func(t *testing.T, env *testEnv) {
|
|
pc := env.pc
|
|
|
|
// Request A: [1,2,3,4,5] + generate [10,11]
|
|
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11})
|
|
env.assertAllTokens(t, "after A", []int32{1, 2, 3, 4, 5, 10})
|
|
|
|
// Request B: [1,2,3,6,7] — diverges at token 4
|
|
simulateRequest(t, pc, []int32{1, 2, 3, 6, 7}, []int32{12, 13})
|
|
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 6, 7, 12})
|
|
|
|
// Request C: switch back to A's branch [1,2,3,4,5,10,11,20]
|
|
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 20}, nil)
|
|
env.assertAllTokens(t, "after C (back to A)", []int32{1, 2, 3, 4, 5, 10, 11})
|
|
|
|
checkTrieInvariants(t, pc.root)
|
|
})
|
|
}
|
|
|
|
// TestLRUOnlyUpdatesUsedNodes verifies that intermediate nodes on the active
|
|
// path whose snapshots were not actually restored don't get their lastUsed
|
|
// refreshed, allowing them to age out and collapse.
|
|
func TestLRUOnlyUpdatesUsedNodes(t *testing.T) {
|
|
forEachEnv(t, func(t *testing.T, env *testEnv) {
|
|
pc := env.pc
|
|
|
|
// Request A: creates path [1,2,3,4,5] + generate [10,11]
|
|
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11})
|
|
|
|
// Request B: diverges at token 4, creating a branch point at offset 3
|
|
// with a split snapshot.
|
|
simulateRequest(t, pc, []int32{1, 2, 3, 6, 7}, []int32{20, 21})
|
|
|
|
// Set all lastUsed to a known old time.
|
|
oldTime := time.Now().Add(-1 * time.Hour)
|
|
walkNodes(pc.root, func(n *trieNode) bool {
|
|
n.lastUsed = oldTime
|
|
return true
|
|
})
|
|
|
|
// Request C: continue on B's branch. This will match B's path
|
|
// and extend it. The branch point's snapshot may be paged in
|
|
// for some cache types but not others.
|
|
beforeRequest := time.Now()
|
|
inputsC := []int32{1, 2, 3, 6, 7, 20, 21, 30}
|
|
resC := simulateRequest(t, pc, inputsC, nil)
|
|
landing := len(inputsC) - len(resC.remaining)
|
|
|
|
// The path must have enough depth to exercise intermediate nodes.
|
|
if len(pc.activePath) < 3 {
|
|
t.Fatalf("activePath too short to test intermediate nodes: got %d nodes", len(pc.activePath))
|
|
}
|
|
|
|
// The frontier (deepest node on the active path) must be updated.
|
|
frontier := pc.activePath[len(pc.activePath)-1]
|
|
if frontier.lastUsed.Before(beforeRequest) {
|
|
t.Errorf("frontier lastUsed was not updated: got %v, want >= %v",
|
|
frontier.lastUsed, beforeRequest)
|
|
}
|
|
|
|
// Only used nodes refresh — the frontier and the restore landing;
|
|
// merely traversed nodes keep their age so they can still evict.
|
|
for i, node := range pc.activePath[:len(pc.activePath)-1] {
|
|
if node.endOffset == landing {
|
|
continue
|
|
}
|
|
if !node.lastUsed.Before(beforeRequest) {
|
|
t.Errorf("activePath[%d] (endOffset=%d) lastUsed was refreshed: got %v, want < %v",
|
|
i, node.endOffset, node.lastUsed, beforeRequest)
|
|
}
|
|
}
|
|
|
|
checkTrieInvariants(t, pc.root)
|
|
})
|
|
}
|
|
|
|
// TestPagedOutBytesUpdatesOnMaterialize verifies that when a snapshot owned
|
|
// by a trie node materializes (allocates owned bytes from a previously lazy
|
|
// state), the trie's pagedOutBytes counter picks up the delta via the
|
|
// installed materialize hook.
|
|
func TestPagedOutBytesUpdatesOnMaterialize(t *testing.T) {
|
|
pc := &prefixCache{}
|
|
pc.ensureRoot()
|
|
|
|
node := &trieNode{parent: pc.root, tokens: []trieKey{1, 2, 3}, endOffset: 3}
|
|
pc.root.children = append(pc.root.children, node)
|
|
|
|
snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
|
|
node.setSnapshots([]cache.Snapshot{snap}, &pc.pagedOutBytes)
|
|
|
|
if pc.pagedOutBytes != 0 {
|
|
t.Fatalf("pagedOutBytes after install = %d, want 0 (lazy snapshot)", pc.pagedOutBytes)
|
|
}
|
|
|
|
const materialized = 1 << 20
|
|
snap.materialize(materialized)
|
|
|
|
if pc.pagedOutBytes != materialized {
|
|
t.Fatalf("pagedOutBytes after materialize = %d, want %d", pc.pagedOutBytes, materialized)
|
|
}
|
|
}
|
|
|
|
// TestSwapSnapshotsDetachesHook verifies that snapshots removed from a trie
|
|
// node via swapSnapshots no longer feed the trie's counter when they later
|
|
// materialize. Without detach, a Split/Merge that folds an old snapshot
|
|
// elsewhere would double-count its bytes if it copied out afterward.
|
|
func TestSwapSnapshotsDetachesHook(t *testing.T) {
|
|
pc := &prefixCache{}
|
|
pc.ensureRoot()
|
|
|
|
node := &trieNode{parent: pc.root, tokens: []trieKey{1, 2, 3}, endOffset: 3}
|
|
pc.root.children = append(pc.root.children, node)
|
|
|
|
snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
|
|
node.setSnapshots([]cache.Snapshot{snap}, &pc.pagedOutBytes)
|
|
|
|
replacement := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
|
|
old := node.swapSnapshots([]cache.Snapshot{replacement}, &pc.pagedOutBytes)
|
|
if len(old) != 1 || old[0] != snap {
|
|
t.Fatalf("swapSnapshots returned %v, want the original snap", old)
|
|
}
|
|
|
|
snap.materialize(1 << 20)
|
|
if pc.pagedOutBytes != 0 {
|
|
t.Fatalf("pagedOutBytes after detached materialize = %d, want 0", pc.pagedOutBytes)
|
|
}
|
|
|
|
replacement.materialize(2 << 20)
|
|
if pc.pagedOutBytes != 2<<20 {
|
|
t.Fatalf("pagedOutBytes after replacement materialize = %d, want %d", pc.pagedOutBytes, 2<<20)
|
|
}
|
|
}
|