diff --git a/internal/lua/vm.go b/internal/lua/vm.go index 37ca840..207b494 100644 --- a/internal/lua/vm.go +++ b/internal/lua/vm.go @@ -4,9 +4,15 @@ // // The runtime uses LuaJIT (github.com/aarzilli/golua) so every adapter runs on // an independent interpreter state. Because a lua.State is NOT goroutine-safe, -// each adapter owns a pool of workers (states); the pool size equals the sum -// of max_concurrent over all sources that use the adapter, and a worker is -// checked out for the duration of a single hook call. +// each adapter owns a pool of workers (states) and a worker is checked out for +// the duration of a single hook call. +// +// Pools are ELASTIC: the sum of max_concurrent over the sources using an adapter +// is a ceiling, not a preallocation. States are booted on demand in batches +// sized by that ceiling, and reclaimed by a single background janitor at a rate +// scaled inversely to the adapter's live connection count. An idle gateway +// therefore holds close to zero Lua states regardless of how much aggregate +// concurrency is configured. See adapterPool for the sizing rules. package lua import ( @@ -23,6 +29,7 @@ import ( "strconv" "strings" "sync" + "time" // golua exposes the LuaJIT C API. Build with -tags luajit so the cgo // LDFLAGS resolve to -lluajit-5.1 (see golua's lua.go). @@ -64,7 +71,20 @@ type staticInfo struct { } // adapterPool owns the worker states for one adapter. Idle workers are held in -// a slice guarded by a cond; up to target workers are created lazily. +// a slice guarded by a cond; workers are created lazily on demand and reclaimed +// by the VM janitor once demand drops. +// +// Sizing is elastic rather than fixed (plan 阶段 2). Three inputs drive it: +// +// - maxW — the adapter's maximum concurrency (Σ max_concurrent of every source +// using it). It is the hard ceiling AND sets the growth step, so a hot, +// high-concurrency adapter ramps in batches while a low-concurrency one +// creeps up one state at a time. +// - created / len(idle) — how many states already exist, i.e. how much room is +// left to grow and how much is available to reclaim. +// - inUse — how many states are checked out right now, i.e. the live +// connection count for this adapter. It sets the shrink step: the fewer +// connections, the more aggressively idle states are released. type adapterPool struct { name string script string @@ -73,24 +93,160 @@ type adapterPool struct { mu sync.Mutex cond *sync.Cond idle []*worker - created int - target int + created int // live states: len(idle) + inUse (+ states being booted) + inUse int // states checked out right now == live connections + waiting int // goroutines blocked waiting for a free state + maxW int // hard ceiling = Σ max_concurrent of the sources using this adapter + used bool closed bool + + lastGrow time.Time + idleRounds int // consecutive janitor rounds that saw reclaimable slack + peakInUse int // high-water mark of inUse, for observability +} + +const ( + // growDivisor turns the adapter's max concurrency into a growth step: + // step = ceil(maxW / growDivisor), clamped to [1, growStepCap]. A source set + // at max_concurrent=8 grows one state at a time; a 64-wide adapter warms 8 + // at once instead of paying 8 sequential boots on a traffic spike. + growDivisor = 8 + growStepCap = 8 + // growCooldown keeps a burst of misses from batching repeatedly while the + // previous batch is still booting. + growCooldown = time.Second + // residentWorkers is how many states an adapter keeps warm once it has + // served at least one request. Booting is milliseconds, but keeping one warm + // removes that from the critical path of the next request. Adapters that + // were never used keep nothing. + residentWorkers = 1 + // idleHeadroom is the spare idle state kept above the live connection count, + // so the next concurrent request does not have to boot. + idleHeadroom = 1 + // shrinkInterval is how often the VM janitor reclaims idle states. + shrinkInterval = 30 * time.Second + // shrinkGraceRounds is how many consecutive janitor rounds must see slack + // before anything is released, so a gap between requests is not mistaken + // for the end of a load period. + shrinkGraceRounds = 2 +) + +// PoolStats is a snapshot of one adapter pool's elastic sizing, surfaced by the +// status API so the algorithm is observable instead of a black box. +type PoolStats struct { + Name string `json:"name"` + Created int `json:"created"` + Idle int `json:"idle"` + InUse int `json:"in_use"` + Waiting int `json:"waiting"` + Max int `json:"max"` + Resident int `json:"resident"` + GrowStep int `json:"grow_step"` + ShrinkStep int `json:"shrink_step"` + PeakInUse int `json:"peak_in_use"` } func newAdapterPool(name, script string, static staticInfo) *adapterPool { - p := &adapterPool{name: name, script: script, static: static, target: 1} + p := &adapterPool{name: name, script: script, static: static, maxW: 1} p.cond = sync.NewCond(&p.mu) return p } -func (p *adapterPool) setTarget(n int) { +// setMax updates the pool's hard ceiling (Σ max_concurrent of its sources). +// Lowering it does not kill live states: the janitor reclaims the excess. +func (p *adapterPool) setMax(n int) { + if n < 1 { + n = 1 + } p.mu.Lock() - p.target = n + p.maxW = n p.cond.Broadcast() p.mu.Unlock() } +func ceilDiv(a, b int) int { + if b <= 0 { + return a + } + return (a + b - 1) / b +} + +// growStepLocked is how many states to warm on a miss, derived from the +// adapter's maximum concurrency. Caller holds p.mu. +func (p *adapterPool) growStepLocked() int { + step := ceilDiv(p.maxW, growDivisor) + if step < 1 { + return 1 + } + if step > growStepCap { + return growStepCap + } + return step +} + +// residentLocked is the floor the janitor will not shrink below. Caller holds p.mu. +func (p *adapterPool) residentLocked() int { + if !p.used { + return 0 + } + if residentWorkers > p.maxW { + return p.maxW + } + return residentWorkers +} + +// shrinkStepLocked is how many idle states to release this round. The step is +// inversely proportional to the live connection count: with no connections the +// slack collapses in one round, while a busy adapter gives up one state at a +// time so the hot path keeps its warm states. Caller holds p.mu. +func (p *adapterPool) shrinkStepLocked() int { + excess := p.created - p.keepLocked() + if excess <= 0 { + return 0 + } + step := ceilDiv(excess, 1+p.inUse) + if step > excess { + step = excess + } + if step > len(p.idle) { + step = len(p.idle) + } + if step < 0 { + return 0 + } + return step +} + +// keepLocked is the number of live states this pool should retain right now. +// Caller holds p.mu. +func (p *adapterPool) keepLocked() int { + keep := p.residentLocked() + if h := p.inUse + idleHeadroom; h > keep { + keep = h + } + if keep > p.maxW { + keep = p.maxW + } + return keep +} + +func (p *adapterPool) stats() PoolStats { + p.mu.Lock() + defer p.mu.Unlock() + return PoolStats{ + Name: p.name, + Created: p.created, + Idle: len(p.idle), + InUse: p.inUse, + Waiting: p.waiting, + Max: p.maxW, + Resident: p.residentLocked(), + GrowStep: p.growStepLocked(), + ShrinkStep: p.shrinkStepLocked(), + PeakInUse: p.peakInUse, + } +} + // boot creates a fresh LuaJIT state, registers the shared globals and executes // the adapter script, storing the returned adapter table in the state. func (p *adapterPool) boot() (*worker, error) { @@ -110,6 +266,81 @@ func (p *adapterPool) boot() (*worker, error) { return &worker{L: L}, nil } +// growPlanLocked reserves capacity for a batch prewarm and returns how many +// EXTRA states to boot in the background (the caller already reserved one for +// itself). contended says the miss happened while every existing state was +// already checked out, which is the real ramp signal: sequential traffic keeps +// reusing one warm state and must never trigger a batch, while genuinely +// concurrent traffic warms a whole step at once instead of paying one boot per +// request all the way up the ramp. Caller holds p.mu. +func (p *adapterPool) growPlanLocked(contended bool) int { + if !contended { + return 0 + } + if time.Since(p.lastGrow) < growCooldown { + return 0 + } + extra := p.growStepLocked() - 1 + if extra <= 0 { + return 0 + } + if room := p.maxW - p.created; extra > room { + extra = room + } + if extra <= 0 { + return 0 + } + // Reserve immediately so concurrent acquires respect the ceiling while the + // batch is still booting. + p.created += extra + p.lastGrow = time.Now() + return extra +} + +// prewarm boots n reserved states and publishes them as idle. A boot failure +// releases its reservation; it is not fatal, since the synchronous path +// surfaces adapter errors already. +func (p *adapterPool) prewarm(n int) { + for i := 0; i < n; i++ { + p.mu.Lock() + closed := p.closed + p.mu.Unlock() + if closed { + p.releaseReservation(n - i) + return + } + w, err := p.boot() + if err != nil { + p.releaseReservation(n - i) + return + } + p.mu.Lock() + if p.closed { + p.created-- + p.mu.Unlock() + w.close() + p.releaseReservation(n - i - 1) + return + } + p.idle = append(p.idle, w) + p.cond.Signal() + p.mu.Unlock() + } +} + +func (p *adapterPool) releaseReservation(n int) { + if n <= 0 { + return + } + p.mu.Lock() + p.created -= n + if p.created < 0 { + p.created = 0 + } + p.cond.Broadcast() + p.mu.Unlock() +} + func (p *adapterPool) acquire() (*worker, error) { p.mu.Lock() for { @@ -120,32 +351,67 @@ func (p *adapterPool) acquire() (*worker, error) { if n := len(p.idle); n > 0 { w := p.idle[n-1] p.idle = p.idle[:n-1] + p.checkoutLocked() p.mu.Unlock() return w, nil } - if p.created < p.target { + if p.created < p.maxW { + // Every existing state busy on a miss = real concurrency, not a + // sequential caller reusing one warm state. + contended := p.created > 0 && p.inUse >= p.created p.created++ - break + extra := p.growPlanLocked(contended) + p.mu.Unlock() + if extra > 0 { + go p.prewarm(extra) + } + w, err := p.boot() + if err != nil { + p.mu.Lock() + p.created-- + p.cond.Signal() + p.mu.Unlock() + return nil, err + } + p.mu.Lock() + p.checkoutLocked() + p.mu.Unlock() + return w, nil } + p.waiting++ p.cond.Wait() + p.waiting-- } - p.mu.Unlock() +} - w, err := p.boot() - if err != nil { - p.mu.Lock() - p.created-- - p.cond.Signal() - p.mu.Unlock() - return nil, err +// checkoutLocked accounts one state as handed out. Caller holds p.mu. +func (p *adapterPool) checkoutLocked() { + p.inUse++ + p.used = true + if p.inUse > p.peakInUse { + p.peakInUse = p.inUse } - return w, nil + // Live demand invalidates any pending shrink decision. + p.idleRounds = 0 } func (p *adapterPool) release(w *worker) { w.L.SetTop(0) p.mu.Lock() + if p.inUse > 0 { + p.inUse-- + } if p.closed { + p.created-- + p.mu.Unlock() + w.close() + return + } + // A lowered ceiling is honoured on the spot rather than waiting for the + // janitor: returning a state the pool may no longer keep closes it. + if p.created > p.maxW { + p.created-- + p.cond.Signal() p.mu.Unlock() w.close() return @@ -155,11 +421,42 @@ func (p *adapterPool) release(w *worker) { p.mu.Unlock() } +// reclaim releases idle states down to the current keep target, using a step +// scaled by the live connection count. It returns how many states were closed. +func (p *adapterPool) reclaim() int { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return 0 + } + step := p.shrinkStepLocked() + if step <= 0 { + p.idleRounds = 0 + p.mu.Unlock() + return 0 + } + p.idleRounds++ + if p.idleRounds < shrinkGraceRounds { + p.mu.Unlock() + return 0 + } + victims := append([]*worker(nil), p.idle[len(p.idle)-step:]...) + p.idle = p.idle[:len(p.idle)-step] + p.created -= step + p.idleRounds = 0 + p.mu.Unlock() + for _, w := range victims { + w.close() + } + return step +} + func (p *adapterPool) shutdown() { p.mu.Lock() p.closed = true idle := p.idle p.idle = nil + p.created -= len(idle) p.cond.Broadcast() p.mu.Unlock() for _, w := range idle { @@ -173,13 +470,84 @@ type VM struct { mu sync.RWMutex dir string pools map[string]*adapterPool + + // janitor drives elastic shrinking. One goroutine serves every pool (not + // one per adapter), so an idle gateway costs a single sleeping goroutine. + janitorOnce sync.Once + janitorStop chan struct{} + janitorDone chan struct{} } func NewVM(dir string) *VM { - return &VM{dir: dir, pools: map[string]*adapterPool{}} + return &VM{ + dir: dir, + pools: map[string]*adapterPool{}, + janitorStop: make(chan struct{}), + janitorDone: make(chan struct{}), + } +} + +// startJanitor launches the single shrink loop. Idempotent. +func (v *VM) startJanitor() { + v.janitorOnce.Do(func() { + go v.janitorLoop(shrinkInterval) + }) +} + +func (v *VM) janitorLoop(every time.Duration) { + defer close(v.janitorDone) + t := time.NewTicker(every) + defer t.Stop() + for { + select { + case <-v.janitorStop: + return + case <-t.C: + v.reclaimIdle() + } + } +} + +// reclaimIdle runs one shrink round over every pool. Exported behaviour is +// tested through ReclaimIdleNow. +func (v *VM) reclaimIdle() int { + v.mu.RLock() + pools := make([]*adapterPool, 0, len(v.pools)) + for _, p := range v.pools { + pools = append(pools, p) + } + v.mu.RUnlock() + n := 0 + for _, p := range pools { + n += p.reclaim() + } + return n +} + +// ReclaimIdleNow forces one shrink round and reports how many worker states +// were closed. Intended for tests and for an explicit "release idle memory" +// action; normal operation relies on the janitor. +func (v *VM) ReclaimIdleNow() int { return v.reclaimIdle() } + +// PoolStats returns the elastic sizing snapshot of every loaded adapter, +// sorted by name. +func (v *VM) PoolStats() []PoolStats { + v.mu.RLock() + pools := make([]*adapterPool, 0, len(v.pools)) + for _, p := range v.pools { + pools = append(pools, p) + } + v.mu.RUnlock() + out := make([]PoolStats, 0, len(pools)) + for _, p := range pools { + out = append(out, p.stats()) + } + sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name }) + return out } func (v *VM) Start() error { + v.startJanitor() if v.dir == "" { return nil } @@ -214,6 +582,12 @@ func (v *VM) Start() error { } func (v *VM) Stop() { + select { + case <-v.janitorStop: + // already stopped + default: + close(v.janitorStop) + } v.mu.Lock() pools := v.pools v.pools = map[string]*adapterPool{} @@ -223,17 +597,17 @@ func (v *VM) Stop() { } } -// ConfigureConcurrency sets each adapter's worker target to the sum of -// max_concurrent of every source using it. Values below one are clamped to 1. +// ConfigureConcurrency sets each adapter's worker CEILING to the sum of +// max_concurrent of every source using it. It is a ceiling, not a preallocation: +// states are booted on demand and reclaimed when demand drops, so a config with +// a large aggregate concurrency no longer implies a large resident pool. +// Values below one are clamped to 1. func (v *VM) ConfigureConcurrency(adapterConcurrency map[string]int) { v.mu.RLock() defer v.mu.RUnlock() for name, n := range adapterConcurrency { if p, ok := v.pools[name]; ok { - if n < 1 { - n = 1 - } - p.setTarget(n) + p.setMax(n) } } } @@ -260,16 +634,19 @@ func (v *VM) LoadAdapter(path string) error { } // LoadAdapterSource registers an adapter from source code. Replacing an -// existing adapter shuts the old pool down and re-boots lazily. +// existing adapter shuts the old pool down and re-boots lazily, carrying the +// configured ceiling over so a hot-reload does not reset elastic sizing. func (v *VM) LoadAdapterSource(name, code string) error { static, err := inspectScript(name, code) if err != nil { return err } - target := 1 + maxW := 1 v.mu.RLock() if p, ok := v.pools[static.name]; ok { - target = p.target + p.mu.Lock() + maxW = p.maxW + p.mu.Unlock() } v.mu.RUnlock() @@ -278,9 +655,10 @@ func (v *VM) LoadAdapterSource(name, code string) error { p.shutdown() } p := newAdapterPool(static.name, code, *static) - p.setTarget(target) + p.setMax(maxW) v.pools[static.name] = p v.mu.Unlock() + v.startJanitor() return nil } diff --git a/internal/lua/vm_test.go b/internal/lua/vm_test.go index d6953c5..1128844 100644 --- a/internal/lua/vm_test.go +++ b/internal/lua/vm_test.go @@ -5,7 +5,9 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" + "time" ) // freshAdapterDir returns a path under a temp dir that does not exist yet, so @@ -484,8 +486,8 @@ func TestAnthropicToolRoundTrip(t *testing.T) { } var r1 struct { Tools []struct { - Name string `json:"name"` - InputSchema map[string]interface{} `json:"input_schema"` + Name string `json:"name"` + InputSchema map[string]interface{} `json:"input_schema"` } `json:"tools"` Messages []map[string]interface{} `json:"messages"` } @@ -512,12 +514,12 @@ func TestAnthropicToolRoundTrip(t *testing.T) { Messages []struct { Role string `json:"role"` Content []struct { - Type string `json:"type"` - ID string `json:"id"` - Name string `json:"name"` + Type string `json:"type"` + ID string `json:"id"` + Name string `json:"name"` Input map[string]interface{} `json:"input"` - ToolUseID string `json:"tool_use_id"` - ContentText string `json:"content"` + ToolUseID string `json:"tool_use_id"` + ContentText string `json:"content"` } `json:"content"` } `json:"messages"` } @@ -746,3 +748,494 @@ func TestAdaptersReportZeroCacheHit(t *testing.T) { t.Fatalf("openai fabricated cache details when upstream reported none: %s", out) } } + +// ---- elastic pool sizing (plan 阶段 2) ---- + +// poolOf reaches into the VM for one adapter's pool. Tests assert on internal +// counters because the whole point of the change is that they no longer grow +// monotonically. +func poolOf(t *testing.T, v *VM, name string) *adapterPool { + t.Helper() + p := v.pool(name) + if p == nil { + t.Fatalf("adapter %q not loaded", name) + } + return p +} + +func poolCounts(p *adapterPool) (created, idle, inUse int) { + p.mu.Lock() + defer p.mu.Unlock() + return p.created, len(p.idle), p.inUse +} + +// TestPoolStartsEmpty is the startup-memory claim: loading adapters must not +// boot a single Lua state, no matter how large the configured ceiling is. +func TestPoolStartsEmpty(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 64, "deepseek": 32}) + for _, st := range vm.PoolStats() { + if st.Created != 0 || st.Idle != 0 || st.InUse != 0 { + t.Fatalf("adapter %q booted eagerly: created=%d idle=%d in_use=%d", + st.Name, st.Created, st.Idle, st.InUse) + } + } +} + +// TestGrowStepFollowsMaxConcurrent pins requirement 3: the growth step is +// derived from the adapter's maximum concurrency. +func TestGrowStepFollowsMaxConcurrent(t *testing.T) { + p := newAdapterPool("t", "return {name='t'}", staticInfo{name: "t"}) + for _, tc := range []struct{ max, want int }{ + {1, 1}, {4, 1}, {8, 1}, {9, 2}, {16, 2}, {32, 4}, {64, 8}, {200, 8}, + } { + p.setMax(tc.max) + p.mu.Lock() + got := p.growStepLocked() + p.mu.Unlock() + if got != tc.want { + t.Fatalf("max=%d: grow step = %d, want %d", tc.max, got, tc.want) + } + } +} + +// TestShrinkStepFollowsConnections pins requirement 4: the shrink step is +// driven by the live connection count — no connections means collapse, busy +// means give up one state at a time. +func TestShrinkStepFollowsConnections(t *testing.T) { + p := newAdapterPool("t", "return {name='t'}", staticInfo{name: "t"}) + p.setMax(32) + p.used = true + + cases := []struct { + created, idle, inUse, want int + why string + }{ + {created: 9, idle: 9, inUse: 0, want: 8, why: "no connections: collapse to the resident floor in one round"}, + {created: 9, idle: 8, inUse: 1, want: 4, why: "1 connection: halve the slack"}, + {created: 9, idle: 6, inUse: 3, want: 2, why: "3 connections: quarter the slack"}, + {created: 9, idle: 2, inUse: 7, want: 1, why: "busy: release one state at a time"}, + {created: 4, idle: 0, inUse: 4, want: 0, why: "everything checked out: nothing to reclaim"}, + {created: 1, idle: 1, inUse: 0, want: 0, why: "at the resident floor: keep the warm state"}, + } + for _, tc := range cases { + p.mu.Lock() + p.created, p.inUse = tc.created, tc.inUse + p.idle = make([]*worker, tc.idle) + got := p.shrinkStepLocked() + p.mu.Unlock() + if got != tc.want { + t.Fatalf("%s: created=%d idle=%d inUse=%d -> step %d, want %d", + tc.why, tc.created, tc.idle, tc.inUse, got, tc.want) + } + } +} + +// TestPoolGrowsOnDemandAndReclaims is the end-to-end elasticity loop: boot on +// demand, stay under the ceiling, then collapse back to the resident floor. +func TestPoolGrowsOnDemandAndReclaims(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 16}) + p := poolOf(t, vm, "openai") + + if created, _, _ := poolCounts(p); created != 0 { + t.Fatalf("pool must start empty, created=%d", created) + } + + // concurrent load: hold several workers at once so the pool must grow + const hold = 6 + var wg sync.WaitGroup + release := make(chan struct{}) + for i := 0; i < hold; i++ { + wg.Add(1) + go func() { + defer wg.Done() + w, err := p.acquire() + if err != nil { + t.Errorf("acquire: %v", err) + return + } + <-release + p.release(w) + }() + } + // wait until all six are checked out + deadline := time.Now().Add(5 * time.Second) + for { + _, _, inUse := poolCounts(p) + if inUse >= hold { + break + } + if time.Now().After(deadline) { + t.Fatalf("workers never checked out, in_use=%d", inUse) + } + time.Sleep(5 * time.Millisecond) + } + created, _, inUse := poolCounts(p) + if inUse != hold { + t.Fatalf("in_use = %d, want %d", inUse, hold) + } + if created < hold { + t.Fatalf("created = %d, must cover %d concurrent checkouts", created, hold) + } + if created > 16 { + t.Fatalf("created = %d exceeds the ceiling 16", created) + } + close(release) + wg.Wait() + + // all connections closed: the janitor collapses the pool to the floor. + // shrinkGraceRounds rounds are required, and one extra round proves the + // floor is stable rather than shrinking to zero. + for i := 0; i < shrinkGraceRounds+2; i++ { + vm.ReclaimIdleNow() + } + created, idle, inUse := poolCounts(p) + if inUse != 0 { + t.Fatalf("in_use = %d after release, want 0", inUse) + } + if created != residentWorkers || idle != residentWorkers { + t.Fatalf("after reclaim: created=%d idle=%d, want %d/%d (resident floor)", + created, idle, residentWorkers, residentWorkers) + } +} + +// TestReclaimNeedsGraceRounds: a single quiet round must not tear the pool +// down — a gap between requests is not the end of a load period. +func TestReclaimNeedsGraceRounds(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 8}) + p := poolOf(t, vm, "openai") + + // build slack: 4 states created, all returned + var ws []*worker + for i := 0; i < 4; i++ { + w, err := p.acquire() + if err != nil { + t.Fatalf("acquire: %v", err) + } + ws = append(ws, w) + } + for _, w := range ws { + p.release(w) + } + created, _, _ := poolCounts(p) + if created != 4 { + t.Fatalf("created = %d, want 4", created) + } + + if n := vm.ReclaimIdleNow(); n != 0 { + t.Fatalf("first quiet round must not reclaim (grace), closed %d", n) + } + if c, _, _ := poolCounts(p); c != 4 { + t.Fatalf("created = %d after the grace round, want 4", c) + } + if n := vm.ReclaimIdleNow(); n == 0 { + t.Fatal("second quiet round must reclaim") + } +} + +// TestCheckoutResetsGrace: live traffic between janitor rounds must cancel a +// pending shrink decision. +func TestCheckoutResetsGrace(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 8}) + p := poolOf(t, vm, "openai") + + var ws []*worker + for i := 0; i < 4; i++ { + w, _ := p.acquire() + ws = append(ws, w) + } + for _, w := range ws { + p.release(w) + } + if n := vm.ReclaimIdleNow(); n != 0 { + t.Fatalf("grace round reclaimed %d", n) + } + // a request arrives: the grace counter resets + w, _ := p.acquire() + p.release(w) + if n := vm.ReclaimIdleNow(); n != 0 { + t.Fatalf("traffic must reset the grace counter, reclaimed %d", n) + } +} + +// TestPoolCeilingIsRespectedUnderContention: acquires beyond the ceiling block +// instead of booting unbounded states. +func TestPoolCeilingIsRespectedUnderContention(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 2}) + p := poolOf(t, vm, "openai") + + w1, err := p.acquire() + if err != nil { + t.Fatalf("acquire 1: %v", err) + } + w2, err := p.acquire() + if err != nil { + t.Fatalf("acquire 2: %v", err) + } + got := make(chan struct{}) + go func() { + w3, err := p.acquire() + if err == nil { + p.release(w3) + } + close(got) + }() + select { + case <-got: + t.Fatal("a third acquire must block at the ceiling of 2") + case <-time.After(100 * time.Millisecond): + } + if created, _, _ := poolCounts(p); created != 2 { + t.Fatalf("created = %d, must not exceed the ceiling 2", created) + } + p.release(w1) + select { + case <-got: + case <-time.After(2 * time.Second): + t.Fatal("releasing a state must unblock the waiter") + } + p.release(w2) +} + +// TestLoweredCeilingReleasesStates: shrinking max_concurrent in the config must +// actually give memory back rather than leaving orphaned states parked. +func TestLoweredCeilingReleasesStates(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 8}) + p := poolOf(t, vm, "openai") + + var ws []*worker + for i := 0; i < 6; i++ { + w, err := p.acquire() + if err != nil { + t.Fatalf("acquire: %v", err) + } + ws = append(ws, w) + } + if created, _, _ := poolCounts(p); created != 6 { + t.Fatalf("created = %d, want 6", created) + } + // operator lowers the ceiling while the states are still checked out + vm.ConfigureConcurrency(map[string]int{"openai": 2}) + for _, w := range ws { + p.release(w) + } + created, idle, _ := poolCounts(p) + if created > 2 || idle > 2 { + t.Fatalf("lowering the ceiling must drop the excess: created=%d idle=%d", created, idle) + } +} + +// TestPoolStatsSurfacesAlgorithm checks the observability payload the status API +// exposes. +func TestPoolStatsSurfacesAlgorithm(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 32}) + p := poolOf(t, vm, "openai") + w, err := p.acquire() + if err != nil { + t.Fatalf("acquire: %v", err) + } + + var st PoolStats + for _, s := range vm.PoolStats() { + if s.Name == "openai" { + st = s + } + } + if st.Max != 32 { + t.Fatalf("Max = %d, want 32", st.Max) + } + if st.InUse != 1 || st.Created != 1 { + t.Fatalf("InUse=%d Created=%d, want 1/1", st.InUse, st.Created) + } + if st.GrowStep != 4 { + t.Fatalf("GrowStep = %d, want ceil(32/8)=4", st.GrowStep) + } + if st.PeakInUse != 1 { + t.Fatalf("PeakInUse = %d, want 1", st.PeakInUse) + } + p.release(w) + + // stats must be sorted by name for a stable UI + names := []string{} + for _, s := range vm.PoolStats() { + names = append(names, s.Name) + } + for i := 1; i < len(names); i++ { + if names[i-1] > names[i] { + t.Fatalf("PoolStats not sorted: %v", names) + } + } +} + +// TestHotReloadKeepsCeiling: replacing an adapter's code must not reset its +// elastic ceiling back to 1. +func TestHotReloadKeepsCeiling(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 24}) + if err := vm.LoadAdapterSource("openai", "return {name='openai', version='9'}"); err != nil { + t.Fatalf("reload: %v", err) + } + p := poolOf(t, vm, "openai") + p.mu.Lock() + maxW := p.maxW + p.mu.Unlock() + if maxW != 24 { + t.Fatalf("ceiling after hot reload = %d, want 24", maxW) + } +} + +// TestConcurrentTransformStress exercises the grow/shrink paths together under +// real hook calls, and asserts the accounting stays consistent (no leaked +// states, no negative counters). +func TestConcurrentTransformStress(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 16}) + + var wg sync.WaitGroup + for i := 0; i < 40; i++ { + wg.Add(1) + go func() { + defer wg.Done() + raw := `{"model":"m","messages":[{"role":"user","content":"hi"}]}` + if _, err := vm.Transform("openai", "transform_request", raw); err != nil { + t.Errorf("transform: %v", err) + } + }() + } + wg.Wait() + + p := poolOf(t, vm, "openai") + created, idle, inUse := poolCounts(p) + if inUse != 0 { + t.Fatalf("in_use = %d after the stress run, want 0 (leaked checkout)", inUse) + } + if created != idle { + t.Fatalf("created=%d idle=%d must match when nothing is checked out", created, idle) + } + if created > 16 { + t.Fatalf("created = %d exceeds the ceiling", created) + } + for i := 0; i < shrinkGraceRounds+2; i++ { + vm.ReclaimIdleNow() + } + created, idle, _ = poolCounts(p) + if created != residentWorkers || idle != residentWorkers { + t.Fatalf("post-stress reclaim: created=%d idle=%d, want %d/%d", + created, idle, residentWorkers, residentWorkers) + } +} + +// TestStopIsIdempotent guards the janitor's stop channel against a double close. +func TestStopIsIdempotent(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + vm.Stop() + vm.Stop() // must not panic +} + +// TestSequentialTrafficDoesNotBatch: a single caller looping must keep reusing +// one warm state — batch prewarm is for genuine concurrency only. +func TestSequentialTrafficDoesNotBatch(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 64}) // grow step would be 8 + p := poolOf(t, vm, "openai") + for i := 0; i < 20; i++ { + w, err := p.acquire() + if err != nil { + t.Fatalf("acquire %d: %v", i, err) + } + p.release(w) + } + created, _, _ := poolCounts(p) + if created != 1 { + t.Fatalf("sequential traffic must reuse one state, created=%d", created) + } +} + +// TestContentionBatchPrewarms: when concurrent demand exceeds the warm set, the +// pool warms a whole grow step instead of one state per request. +func TestContentionBatchPrewarms(t *testing.T) { + vm := NewVM(freshAdapterDir(t)) + if err := vm.Start(); err != nil { + t.Fatalf("start: %v", err) + } + defer vm.Stop() + vm.ConfigureConcurrency(map[string]int{"openai": 32}) // grow step = 4 + p := poolOf(t, vm, "openai") + + // hold the first state, then miss while it is busy -> contended + w1, err := p.acquire() + if err != nil { + t.Fatalf("acquire 1: %v", err) + } + w2, err := p.acquire() + if err != nil { + t.Fatalf("acquire 2: %v", err) + } + // the batch is booted in the background; wait for it to land + deadline := time.Now().Add(5 * time.Second) + for { + created, _, _ := poolCounts(p) + if created >= 5 { // 2 checked out + 3 prewarmed + break + } + if time.Now().After(deadline) { + t.Fatalf("contended miss did not batch prewarm, created=%d", created) + } + time.Sleep(5 * time.Millisecond) + } + created, _, _ := poolCounts(p) + if created > 32 { + t.Fatalf("prewarm exceeded the ceiling, created=%d", created) + } + p.release(w1) + p.release(w2) +}