llmgw/internal/router/prefix/tree_test.go

package prefix

import (
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"
)

func tree(t *testing.T, cfg Config) *Tree {
	t.Helper()
	return New(cfg)
}

func TestTheLongestMatchWins(t *testing.T) {
	// Two replicas have seen overlapping prefixes; the one that has seen more
	// of this prompt is the one that has more of it cached.
	tr := tree(t, Config{MinMatch: 4})
	tr.Observe("You are a helpful assistant. ", "a")
	tr.Observe("You are a helpful assistant. Answer in French. ", "b")

	replica, n := tr.Match("You are a helpful assistant. Answer in French. What is 2+2?")
	if replica != "b" {
		t.Errorf("matched %q with %d bytes; the longer prefix belongs to b", replica, n)
	}
	if n < len("You are a helpful assistant. Answer in French. ") {
		t.Errorf("matched only %d bytes", n)
	}
}

func TestAShortOverlapIsNotAffinity(t *testing.T) {
	// Two English prompts share "The". Treating coincidence as affinity is how
	// prefix routing drowns one replica.
	tr := tree(t, Config{MinMatch: 64})
	tr.Observe("The quick brown fox", "a")

	if replica, n := tr.Match("The slow green turtle"); replica != "" {
		t.Errorf("a %d-byte coincidence was treated as a match to %q", n, replica)
	}
}

func TestSplittingAnExistingBranch(t *testing.T) {
	// The case a radix tree exists for: a new prompt diverges partway through
	// a label that was one node.
	tr := tree(t, Config{MinMatch: 4})
	tr.Observe("abcdefgh", "a")
	tr.Observe("abcdXYZW", "b")

	if r, _ := tr.Match("abcdefgh and more"); r != "a" {
		t.Errorf("the original branch matched %q", r)
	}
	if r, _ := tr.Match("abcdXYZW and more"); r != "b" {
		t.Errorf("the new branch matched %q", r)
	}
	// The shared head belongs to both, because both were sent something
	// beginning with it — and a prefix cache holds leading bytes, not whole
	// prompts. An earlier version treated this as no match, which made the
	// whole tree inert in practice.
	if r, n := tr.Match("abcd"); r == "" || n != 4 {
		t.Errorf("the shared head matched %q with %d bytes; both replicas hold it", r, n)
	}
}

func TestAPartialLabelIsStillAMatch(t *testing.T) {
	// The correction that made this feature work at all.
	//
	// The first request for a prompt becomes one long label. The second differs
	// in its last few bytes — a new question after the same system prompt,
	// which is the single most common shape there is. Requiring the whole label
	// to match meant every lookup after the first returned nothing: measured as
	// 144 misses and zero hits.
	//
	// What matters is the shared head, because that is what the engine cached.
	tr := tree(t, Config{MinMatch: 2})
	tr.Observe("abcdefghij", "a")

	r, n := tr.Match("abcdeZZZZ")
	if r != "a" {
		t.Errorf("a partial match returned %q", r)
	}
	if n != 5 {
		t.Errorf("matched %d bytes, want the 5 they share", n)
	}
}

func TestEntriesExpire(t *testing.T) {
	// An engine evicts on its own schedule, so an old record describes a cache
	// that may be long gone. Believing it forever sends traffic towards
	// something that is not there.
	now := time.Now()
	tr := New(Config{MinMatch: 4, TTL: time.Minute})
	tr.now = func() time.Time { return now }

	tr.Observe("shared prefix here", "a")
	if r, _ := tr.Match("shared prefix here and more"); r != "a" {
		t.Fatalf("fresh entry matched %q", r)
	}

	now = now.Add(2 * time.Minute)
	if r, _ := tr.Match("shared prefix here and more"); r != "" {
		t.Errorf("a stale entry still matched %q", r)
	}
}

func TestForgettingAReplica(t *testing.T) {
	// A replica that has gone away must stop attracting traffic. Nothing reuses
	// a replica id, so the entries would otherwise only ever expire.
	tr := tree(t, Config{MinMatch: 4})
	tr.Observe("a long enough shared prefix", "a")
	tr.Observe("a long enough shared prefix", "b")

	tr.Forget("a")
	r, _ := tr.Match("a long enough shared prefix and more")
	if r == "a" {
		t.Error("a forgotten replica still matches")
	}
	if r != "b" {
		t.Errorf("forgetting one replica lost the other: %q", r)
	}
}

func TestTheTreeIsBounded(t *testing.T) {
	// A router that runs for months would otherwise accumulate every prompt it
	// has ever seen.
	tr := tree(t, Config{MinMatch: 4, MaxNodes: 200})
	for i := range 2000 {
		tr.Observe(fmt.Sprintf("prompt number %d with some trailing text", i), "a")
	}
	if got := tr.Size(); got > 400 {
		t.Errorf("the tree grew to %d nodes against a limit of 200", got)
	}
	// And it still works afterwards.
	tr.Observe("a fresh prompt after all that eviction", "b")
	if r, _ := tr.Match("a fresh prompt after all that eviction, continued"); r != "b" {
		t.Errorf("the tree stopped working after eviction: %q", r)
	}
}

func TestConcurrentUse(t *testing.T) {
	tr := tree(t, Config{MinMatch: 4, MaxNodes: 500})
	var wg sync.WaitGroup
	for i := range 32 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			p := strings.Repeat("x", 40) + fmt.Sprint(i%4)
			for range 50 {
				tr.Observe(p, fmt.Sprintf("r%d", i%4))
				tr.Match(p)
				tr.Forget("r0")
			}
		}()
	}
	wg.Wait()
}

func TestEmptyInputsAreHarmless(t *testing.T) {
	tr := tree(t, Config{})
	tr.Observe("", "a")
	tr.Observe("something", "")
	if r, n := tr.Match(""); r != "" || n != 0 {
		t.Errorf("an empty prompt matched %q/%d", r, n)
	}
}