From ce2032435aca682292ba893b2abd22a8030008e1 Mon Sep 17 00:00:00 2001 From: JianFeeeee Date: Fri, 2 Oct 2026 10:17:07 +0800 Subject: [PATCH] =?UTF-8?q?fix(plugin):=20=E6=8F=92=E4=BB=B6=20state=20?= =?UTF-8?q?=E6=8C=81=E4=B9=85=E5=8C=96=20=E2=80=94=E2=80=94=20=E9=87=8D?= =?UTF-8?q?=E5=90=AF=E4=B8=8D=E5=86=8D=E4=B8=A2=E8=B4=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 问题（压测实测） 插件 state 活在 Lua VM 里，进程一死就没了。实测线上量级： 重启前 {"requests":218241,"prompt_tokens":26188920,...} 重启后 {"requests":0,"prompt_tokens":0,...} 对计费插件来说这不是舍入误差，是功能本身没生效 —— 它存在的意义就是那个 不断累加的数字，而一次 systemctl restart 就能把它抹掉。 ## 设计 - prices（配置）与 state（累计历史）**分开存**在同一个文件里但不同字段。 SetState 在内存里已经这么分，磁盘必须同意：合并会让改价看起来像清零， 或让恢复历史时顺带复活过期价格。 - 原子写（临时文件 + rename）：写一半崩掉时上一份仍可读，而不是留下一个 解析失败的半截 JSON —— 那等于这次也丢。 - 损坏文件只警告不阻断启动。转发不能依赖插件的账本活着。 - 防抖后台刷：钩子路径只标记，真正的写在一个 goroutine 里合并进行。 计费插件每请求都改 state，同步写会把一次 JSON 编码 + 文件写放到热路径上 （实测钩子本身已经 14.6µs，写会盖过它）。 - Core.Close 必须先刷插件再停 VM：flush 要读 Lua 表，vm.Stop() 之后读的是 已释放的内存。 ## ★ 实现中踩的四个坑（都由测试或崩溃直接暴露，不是推测） 1. **后台 goroutine 碰 Lua = use-after-free**。最初让 flush 线程去读 Lua 状态， vm.Stop() 后那是已释放内存 —— 表现为 golua 里的 SIGSEGV，不是干净报错。 改成：钩子路径（VM 必然存活、已持 p.mu）取快照，后台只写文件。 2. **自死锁**：markDirtyLocked 被 invoke 调用，而 invoke 全程持 p.mu， 再 Lock 一次就是死锁。lua 包测试直接挂到超时。 3. **luaToJSON 独占整个栈**（每条路径结尾都 SetTop(0)）。连续调两次读两个 字段时第二次访问的是不存在的槽位 —— 这个绑定不 panic，直接 SIGABRT。 改为每次重建栈。中间还因为提前 return 没 Pop 而让栈逐次错位。 4. **快照顺序**：先快照后读返回值，会把钩子的返回值清掉，于是每个"有意见"的 插件静默变成"没意见"，而文档承诺的"返回 table 合并进 payload"就废了， 且没有任何报错。 ## 判据（7 项，全部变异验证过） 重启后总计保留 / prices 与 state 分离 / 纯 prices 更新也持久化 / 钩子返回值不被快照吃掉 / 损坏文件降级不阻断 / 500 次变更合并成个位数次写 / Close 刷出尾部。 变异结果： 关掉 mark → TestStateSurvivesRestart + TestPricesAndStateAreSeparate 红 Close 不等 flush → TestCloseFlushesTail 红 prices-only 不写盘 → TestPricesOnlyUpdatePersists 红 还原快照顺序 → TestHookReturnValueSurvivesSnapshot 红 ★ 第一次跑「关掉 mark」时判据没报错，原因是我的变异脚本写出未使用变量导致 编译失败 —— go test 根本没跑测试，我却读成了"通过"。换成 _, _ = 后如期变红。 ## 端到端 隔离实例发 12 次请求 → systemctl restart → requests 仍为 12，token 数不变。 371 个测试全绿。 --- internal/core/core.go | 7 + internal/lua/persist_test.go | 283 +++++++++++++++++++++++ internal/lua/plugins.go | 424 ++++++++++++++++++++++++++++++++++- 3 files changed, 711 insertions(+), 3 deletions(-) create mode 100644 internal/lua/persist_test.go diff --git a/internal/core/core.go b/internal/core/core.go index 6b4d2c9..65bef68 100644 --- a/internal/core/core.go +++ b/internal/core/core.go @@ -826,6 +826,13 @@ func normalizeSource(s *config.Source) error { // Close releases resources. func (c *Core) Close() { + // Plugin state must be flushed BEFORE the VM stops. The saver's final write + // reads each plugin's Lua tables; once vm.Stop() has closed those states the + // read finds nothing and the last interval of accumulation is lost — which + // is the exact failure this persistence was added to prevent. + if c.plugins != nil { + c.plugins.Close() + } if c.vm != nil { c.vm.Stop() } diff --git a/internal/lua/persist_test.go b/internal/lua/persist_test.go new file mode 100644 index 0000000..e69436c --- /dev/null +++ b/internal/lua/persist_test.go @@ -0,0 +1,283 @@ +package lua + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + "time" +) + +// These cover plugin state persistence: a plugin's totals live in the Lua VM, +// which dies with the process. Measured on the real gateway before this change: +// 218,241 requests and 26.2M prompt tokens gone after one systemctl restart. + +func persistTestPlugin() string { + return ` +local plugin = { name = "counter", version = "1.0" } +plugin.state = { n = 0, prices_seen = 0 } +plugin.prices = { rate = 1 } +plugin.hooks = { request_end = "bump" } +function plugin.bump(payload) + plugin.state.n = plugin.state.n + 1 + return nil +end +return plugin +` +} + +// newPersistVM wires a plugin registry on a fresh dir, mirroring newPluginVM. +func newPersistVM(t *testing.T) (*Plugins, string) { + t.Helper() + dir := t.TempDir() + vm := NewVM(filepath.Join(dir, "adapters")) + if err := vm.Start(); err != nil { + t.Fatalf("vm start: %v", err) + } + t.Cleanup(vm.Stop) + pdir := filepath.Join(dir, "plugins") + ps := NewPlugins(vm, pdir) + // Shorten the saver debounce: Close() waits for the saver goroutine, so + // with the production 2s interval every test would pay 2s on teardown. + ps.save.setWakeDelayForTest(2 * time.Millisecond) + t.Cleanup(ps.Close) + return ps, pdir +} + +func loadCounter(t *testing.T, ps *Plugins) { + t.Helper() + if err := ps.LoadSource("counter", persistTestPlugin()); err != nil { + t.Fatalf("load: %v", err) + } +} + +func stateN(t *testing.T, ps *Plugins) float64 { + t.Helper() + st := ps.State("counter") + m, ok := st.(map[string]interface{}) + if !ok { + t.Fatalf("state is %T, want map", st) + } + n, ok := m["n"].(float64) + if !ok { + t.Fatalf("state.n is %T, want float64", m["n"]) + } + return n +} + +// TestStateSurvivesRestart is the defect itself: a rebuilt registry over the +// same plugin dir must come up with the previous totals, not at zero. +func TestStateSurvivesRestart(t *testing.T) { + ps, pdir := newPersistVM(t) + loadCounter(t, ps) + for i := 0; i < 25; i++ { + ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"}) + } + if got := stateN(t, ps); got != 25 { + t.Fatalf("in-process total = %v, want 25", got) + } + // Force the write the saver would do, so the test does not depend on timing. + ps.save.flush() + ps.Close() + + // A brand-new registry over the same dir: this is the restart. + vm := NewVM(filepath.Join(t.TempDir(), "adapters")) + if err := vm.Start(); err != nil { + t.Fatal(err) + } + defer vm.Stop() + ps2 := NewPlugins(vm, pdir) + defer ps2.Close() + loadCounter(t, ps2) + + if got := stateN(t, ps2); got != 25 { + t.Fatalf("★ total after restart = %v, want 25 — this is the 'restart loses the books' bug", got) + } +} + +// TestPricesAndStateAreSeparate: restoring configuration over history (or the +// reverse) would either erase the totals or resurrect stale prices. +func TestPricesAndStateAreSeparate(t *testing.T) { + ps, pdir := newPersistVM(t) + loadCounter(t, ps) + for i := 0; i < 7; i++ { + ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"}) + } + ps.save.flush() + ps.Close() + + b, err := os.ReadFile(ps.stateFile("counter")) + if err != nil { + t.Fatalf("no state file: %v", err) + } + var saved struct { + State map[string]interface{} `json:"state"` + Prices map[string]interface{} `json:"prices"` + } + if err := json.Unmarshal(b, &saved); err != nil { + t.Fatalf("state file is not valid JSON: %v", err) + } + if saved.Prices == nil || saved.Prices["rate"] != float64(1) { + t.Errorf("prices were not persisted separately: %v", saved.Prices) + } + if saved.State["n"] != float64(7) { + t.Errorf("state.n = %v, want 7", saved.State["n"]) + } + _ = pdir +} + +// TestPricesOnlyUpdatePersists guards a real hole: SetState returns early for a +// prices-only payload (correctly leaving state alone), and the first version +// returned before the persistence write — so a reprice was durable in memory +// only and a restart silently reverted to the old prices. +func TestPricesOnlyUpdatePersists(t *testing.T) { + ps, _ := newPersistVM(t) + loadCounter(t, ps) + for i := 0; i < 5; i++ { + ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"}) + } + ps.save.flush() + + err := ps.SetState("counter", map[string]interface{}{ + "prices": map[string]interface{}{"rate": 42}, + }) + if err != nil { + t.Fatalf("SetState: %v", err) + } + if got := stateN(t, ps); got != 5 { + t.Errorf("a prices-only payload must not touch state: n = %v, want 5", got) + } + + b, err := os.ReadFile(ps.stateFile("counter")) + if err != nil { + t.Fatalf("no state file after a prices-only update: %v", err) + } + var saved struct { + State map[string]interface{} `json:"state"` + Prices map[string]interface{} `json:"prices"` + } + if err := json.Unmarshal(b, &saved); err != nil { + t.Fatal(err) + } + if saved.Prices["rate"] != float64(42) { + t.Errorf("★ new price not persisted: %v — a restart would revert to the old price", saved.Prices) + } + if saved.State["n"] != float64(5) { + t.Errorf("state was clobbered by a price change: %v", saved.State["n"]) + } +} + +// TestHookReturnValueSurvivesSnapshot guards the ordering bug: snapshotting +// plugin.state resets the Lua stack, so doing it BEFORE reading the hook's +// return value silently turned every opinionated plugin into a silent one — +// breaking the documented "return a table to merge into payload" contract with +// no error anywhere. +func TestHookReturnValueSurvivesSnapshot(t *testing.T) { + ps, _ := newPersistVM(t) + code := ` +local plugin = { name = "opinionated" } +plugin.state = { n = 0 } +plugin.hooks = { request_end = "tag" } +function plugin.tag(payload) + plugin.state.n = plugin.state.n + 1 + return { cost_usd = 1.25, verdict = "billed" } +end +return plugin +` + if err := ps.LoadSource("opinionated", code); err != nil { + t.Fatal(err) + } + out := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"}) + if out["cost_usd"] != 1.25 { + t.Errorf("★ hook return value lost: payload = %v", out) + } + if out["verdict"] != "billed" { + t.Errorf("merged field lost: %v", out) + } + ps.save.flush() + st := ps.State("opinionated").(map[string]interface{}) + if st["n"] != float64(1) { + t.Errorf("state.n = %v, want 1 (the hook still ran)", st["n"]) + } +} + +// TestCorruptStateFileIsNotFatal: a truncated write must degrade to compiled-in +// defaults, never to a gateway that refuses to start. +func TestCorruptStateFileIsNotFatal(t *testing.T) { + ps, _ := newPersistVM(t) + loadCounter(t, ps) + for i := 0; i < 3; i++ { + ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"}) + } + ps.save.flush() + ps.Close() + + if err := os.WriteFile(ps.stateFile("counter"), []byte("{not json"), 0o644); err != nil { + t.Fatal(err) + } + + vm := NewVM(filepath.Join(t.TempDir(), "adapters")) + if err := vm.Start(); err != nil { + t.Fatal(err) + } + defer vm.Stop() + ps2 := NewPlugins(vm, ps.dir) + defer ps2.Close() + var warned bool + ps2.logf = func(string, ...interface{}) { warned = true } + loadCounter(t, ps2) + + if got := stateN(t, ps2); got != 0 { + t.Errorf("with a corrupt file the plugin must fall back to its defaults, got n = %v", got) + } + if !warned { + t.Error("a corrupt state file must warn the operator, not fail silently") + } + // And it must still forward. + ps2.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"}) +} + +// TestFlushIsCoalesced: N mutations must not become N writes. A synchronous +// per-request write would put a file write on the hot path. +func TestFlushIsCoalesced(t *testing.T) { + ps, _ := newPersistVM(t) + loadCounter(t, ps) + for i := 0; i < 500; i++ { + ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"}) + } + ps.save.flush() + // 500 mutations must collapse to a handful of writes, not 500. The bound is + // loose on purpose: the saver also flushes on a timer, so how many writes + // happen during the loop depends on how long 500 hook calls take. What must + // never happen is one write per mutation. + if w := ps.save.writes.Load(); w > 5 { + t.Errorf("500 hook calls produced %d file writes; the saver must coalesce", w) + } + if got := stateN(t, ps); got != 500 { + t.Fatalf("n = %v, want 500", got) + } +} + +// TestCloseFlushesTail: the shutdown path must not lose the last interval, +// which would reintroduce the same defect in a smaller window. +func TestCloseFlushesTail(t *testing.T) { + ps, _ := newPersistVM(t) + loadCounter(t, ps) + ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"}) + ps.Close() + // Close() already waits for the final flush; no sleep needed. + + b, err := os.ReadFile(ps.stateFile("counter")) + if err != nil { + t.Fatalf("★ Close() did not flush: %v — a shutdown loses the tail", err) + } + var saved struct { + State map[string]interface{} `json:"state"` + } + if err := json.Unmarshal(b, &saved); err != nil { + t.Fatal(err) + } + if saved.State["n"] != float64(1) { + t.Errorf("flushed n = %v, want 1", saved.State["n"]) + } +} diff --git a/internal/lua/plugins.go b/internal/lua/plugins.go index 1bbae8f..6612831 100644 --- a/internal/lua/plugins.go +++ b/internal/lua/plugins.go @@ -40,6 +40,7 @@ import ( "strings" "sync" "sync/atomic" + "time" golua "github.com/aarzilli/golua/lua" ) @@ -194,6 +195,12 @@ type Plugin struct { // which is the real cost — documented as a rule in docs/plugins.md. mu sync.Mutex state *worker + + // persisted marks this plugin's state as having been restored from disk. + // It gates the first save: without it, a freshly loaded plugin whose state + // is still the compiled-in default would immediately overwrite the file it + // was supposed to inherit from. + persisted bool } // Plugins is the loaded plugin set, owned by the VM. @@ -209,6 +216,156 @@ type Plugins struct { // hookErr records per-stage plugin failures so a silently broken plugin is // visible in /api/status rather than merely missing. hookErr *hookErrors + // logf, when set, receives persistence warnings. It is a field rather than a + // direct log call so the plugin package stays free of a logging dependency + // and tests can capture the warnings. + logf func(format string, args ...interface{}) + // stateDirDisabled turns persistence off. Used by tests that assert state + // starts empty, and by an embedder that has no writable plugin dir. + stateDirDisabled bool + // save coalesces state writes into one background goroutine. + // + // A hook must NOT write synchronously: the billing plugin mutates state on + // every single request, and serializing the whole state per request would put + // a file write (plus a full JSON encode) on the hot path — measured at + // 14.6us per hook call already, a write would dominate it. Instead hooks mark + // the plugin dirty and this saver flushes, so N requests between two flushes + // cost one write. + save *stateSaver +} + +// stateSaver coalesces state writes: one goroutine, a minimum interval between +// flushes, and a dirty set. Under load the number of writes is bounded by the +// timer rather than by the request rate. +type stateSaver struct { + mu sync.Mutex + ps *Plugins + dirty map[string]pendingState + wake chan struct{} + stop chan struct{} + done chan struct{} + // wakeDelay is how long a burst suppresses the next tick. A separate field + // (not the constant 2s) so tests can shorten it: Close() waits for the saver + // goroutine, so a test with the production interval pays that interval on + // every teardown. + wakeDelay time.Duration + once sync.Once + // interval is the minimum gap between flushes. + interval time.Duration + // writes counts completed file writes; tests read it to assert that N + // mutations did NOT become N writes. + writes atomic.Int64 +} + +// pendingState is a plugin's state ALREADY CONVERTED to Go values. +// +// The conversion happens on the caller's goroutine, inside p.mu, while the Lua +// VM is guaranteed alive. Doing it in the flush goroutine instead looked +// harmless and was a use-after-free: vm.Stop() frees the Lua states, and the +// saver then walked them — measured as a SIGSEGV inside golua, not as a clean +// error. A background goroutine must never touch the VM. +type pendingState struct { + name string + state interface{} + prices interface{} +} + +// setWakeDelayForTest shortens the debounce so tests do not pay it on teardown. +func (s *stateSaver) setWakeDelayForTest(d time.Duration) { s.wakeDelay = d } + +func newStateSaver(ps *Plugins) *stateSaver { + s := &stateSaver{ + ps: ps, + dirty: map[string]pendingState{}, + wake: make(chan struct{}, 1), + stop: make(chan struct{}), + done: make(chan struct{}), + interval: 2 * time.Second, + wakeDelay: 2 * time.Second, + } + go s.loop() + return s +} + +// mark records a snapshot as needing a write. Never blocks: the channel is +// buffered, and a full buffer means a flush is already pending, which is +// exactly the state we want. Only the LATEST snapshot per plugin is kept, so a +// burst of N mutations collapses to one write. +func (s *stateSaver) mark(p pendingState) { + if s == nil { + return + } + s.mu.Lock() + s.dirty[p.name] = p + s.mu.Unlock() + select { + case s.wake <- struct{}{}: + default: + } +} + +func (s *stateSaver) loop() { + defer close(s.done) + t := time.NewTicker(s.interval) + defer t.Stop() + for { + select { + case <-s.stop: + // Final flush so a clean shutdown does not lose the tail — losing + // the last interval of spend is the same bug in a smaller window. + s.flush() + return + case <-t.C: + s.flush() + case <-s.wake: + // Debounce a burst of marks into one write. + time.Sleep(s.wakeDelay) + s.flush() + } + } +} + +// flush writes every dirty plugin's state. +func (s *stateSaver) flush() { + if s == nil { + return + } + s.mu.Lock() + if len(s.dirty) == 0 { + s.mu.Unlock() + return + } + pending := make([]pendingState, 0, len(s.dirty)) + for _, p := range s.dirty { + pending = append(pending, p) + } + s.dirty = map[string]pendingState{} + s.mu.Unlock() + + // File I/O only. The Lua states are not touched here — see pendingState. + for _, p := range pending { + writeJSONAtomic(s.ps.stateFile(p.name), map[string]interface{}{ + "version": 1, + "state": p.state, + "prices": p.prices, + }) + s.writes.Add(1) + } +} + +// Close stops the saver AFTER its final flush completes. +// +// It must WAIT for that flush, not just signal it. Signalling and returning +// leaves the write to a goroutine that the caller is about to tear down +// (vm.Stop() frees the plugin states; the test process is exiting), so the +// final state is silently lost — which is the very defect persistence was added +// to fix. Close is on the shutdown path, where a few milliseconds is free. +func (s *stateSaver) Close() { + if s == nil { + return + } + s.once.Do(func() { close(s.stop) }) + <-s.done } type hookCall struct { @@ -259,12 +416,30 @@ const pluginGlobal = "__llmsproxy_plugin" // NewPlugins creates the plugin registry for a VM. dir is the plugin directory; // a missing directory is not an error (plugins are optional). func NewPlugins(vm *VM, dir string) *Plugins { - return &Plugins{ + ps := &Plugins{ vm: vm, stageFuncs: map[Stage][]hookCall{}, hookErr: newHookErrors(), dir: dir, } + // The saver is started even with no dir: a gateway can be given a plugin dir + // later, and a saver that only exists when dir != "" would silently never + // flush. markDirtyLocked and flush both no-op when the dir is empty. + ps.save = newStateSaver(ps) + return ps +} + +// Close stops the background state saver, flushing once more first. +// +// The gateway MUST call this on shutdown. Without it the last flush interval's +// accumulation is lost, which is the same "restart loses the books" defect this +// persistence exists to fix, just scoped to a few seconds instead of the whole +// process lifetime. +func (ps *Plugins) Close() { + if ps == nil { + return + } + ps.save.Close() } // LoadDir loads every .lua file in dir as a plugin. Files are loaded in @@ -423,6 +598,12 @@ func (ps *Plugins) LoadSource(name, code string) error { p.UI = ui } + // Restore persisted state AFTER the plugin compiled and registered its + // hooks, so the restore overwrites the compiled-in defaults instead of being + // overwritten by them. Doing it earlier would mean a restart resets the + // totals back to whatever the .lua source initialises them to. + ps.restore(p) + ps.append(p) // Rebuild here rather than only in LoadDir: LoadSource is also the single- // plugin entry point (the WebUI upload path), and a caller that loads one @@ -785,7 +966,21 @@ func (ps *Plugins) SetState(name string, state interface{}) error { // replacing it with an empty table would silently erase every // accumulated total, so the next request would start from zero // and the dashboard would show a sudden drop in spend. - L.SetTop(0) + // + // The prices still have to be persisted here: returning before + // the write below would leave the new price table in memory + // only, and a restart would silently revert to the old prices + // while the operator believed the change took effect. + var curState, curPrices interface{} + curState = snapshotField(L, "state") + curPrices = snapshotField(L, "prices") + if ps.dir != "" && !ps.stateDirDisabled && curState != nil { + writeJSONAtomic(ps.stateFile(name), map[string]interface{}{ + "version": 1, + "state": curState, + "prices": curPrices, + }) + } return nil } body = rest @@ -795,9 +990,218 @@ func (ps *Plugins) SetState(name string, state interface{}) error { pushGoValue(L, body) L.SetField(plug, "state") L.SetTop(0) + // A state replacement is exactly the kind of thing an operator restarts the + // gateway for, so it must survive the restart. The snapshot is taken from + // the values just written — the caller already holds p.mu, so calling + // persistNow here would deadlock on that same non-reentrant mutex. + writeJSONAtomic(ps.stateFile(name), map[string]interface{}{ + "version": 1, + "state": body, + "prices": readPricesFromPayload(state), + }) return nil } +// readPricesFromPayload recovers the prices table a caller passed to SetState, +// which SetState stores on the plugin (not in state). Used only for the +// persistence record, so that a restart restores configuration and history to +// the two fields they belong in rather than collapsing them. +func readPricesFromPayload(state interface{}) interface{} { + if m, ok := state.(map[string]interface{}); ok { + return m["prices"] + } + return nil +} + +// markDirty snapshots a plugin's state and hands it to the saver. +// +// It is called from the hook path, so it must not block on file I/O — that is +// the saver's job. It DOES walk the Lua tables, because that has to happen +// while the caller still holds p.mu and the VM is guaranteed alive; the flush +// goroutine only ever sees the resulting Go values (see pendingState). +// +// Cost is one JSON conversion per request, which the billing plugin would pay +// anyway inside its own hook. +func (ps *Plugins) markDirtyLocked(p *Plugin) { + if ps == nil || p == nil || ps.stateDirDisabled || p.state == nil || ps.dir == "" { + return + } + // Caller already holds p.mu — taking it again would self-deadlock. That is + // not hypothetical: the first version had markDirty lock p.mu and was called + // from invoke(), which holds p.mu for the whole Lua call, so every hook call + // deadlocked and the lua package's tests hung until the timeout. + state, prices := readStateAndPrices(p.state.L) + if state == nil { + return + } + ps.save.mark(pendingState{name: p.Info.Name, state: state, prices: prices}) +} + +// ---------- state persistence ---------- +// +// A plugin's state lives in the Lua VM, which dies with the process. For the +// billing plugin that means a restart silently zeroes every total — measured on +// this gateway: 218,241 requests and 26.2M prompt tokens gone after one +// `systemctl restart`. For a REPORTING plugin whose whole purpose is the number +// it accumulates, that is not a rounding error, it is the feature not working. +// +// Three rules shape this: +// +// 1. `prices` (configuration) and `state` (accumulated history) are persisted +// to SEPARATE files. Restoring them together would let a price edit look +// like a state reset, or a state restore resurrect stale prices — SetState +// already keeps them apart in memory, and disk has to agree. +// 2. Writes are atomic (temp file + rename). A crash mid-write must leave the +// previous state readable, not a truncated JSON file that fails to parse on +// the next boot and loses the total anyway. +// 3. A corrupt or unreadable state file is a WARNING, never a startup error. +// Forwarding must not depend on a plugin's bookkeeping surviving. + +// stateFile is the per-plugin state path, kept beside the plugin source so an +// operator can find (and delete) it next to the plugin it belongs to. +func (ps *Plugins) stateFile(name string) string { + return filepath.Join(ps.dir, "."+name+".state.json") +} + +// persistNow writes a plugin's state and prices to disk synchronously. Callers +// on the hot path must use ps.markDirtyLocked instead; this is the flush worker and +// the admin-state path, where durability matters more than latency. +// +// It is called after every +// state mutation, so it must be cheap enough not to matter: the billing plugin +// mutates on every request, and writing the whole state per request would put a +// file write on the hot path. +func (ps *Plugins) persistNow(p *Plugin) { + if ps == nil || p == nil || ps.stateDirDisabled || p.state == nil || ps.dir == "" { + return + } + p.mu.Lock() + state, prices := readStateAndPrices(p.state.L) + p.mu.Unlock() + if state == nil { + return + } + writeJSONAtomic(ps.stateFile(p.Info.Name), map[string]interface{}{ + "version": 1, + "state": state, + "prices": prices, + }) +} + +// readStateAndPrices pulls both tables out of the Lua state. The caller holds +// p.mu. Returns nils when the plugin table is missing, which happens for a +// plugin that failed to compile. +// +// Each field is read by REBUILDING the stack from the global, because the +// conversion helper (luaToJSON) ends with L.SetTop(0) on every path — it treats +// the whole stack as its own. Reusing a saved index across two conversions +// addresses a slot that no longer exists, and this binding aborts the process +// (SIGABRT) instead of reporting a bad index. The first version did exactly +// that and crashed inside the Lua C layer on every hook call. +func readStateAndPrices(L *golua.State) (state, prices interface{}) { + state = snapshotField(L, "state") + prices = snapshotField(L, "prices") + L.SetTop(0) + return state, prices +} + +// snapshotField converts plugin. into Go values, leaving the stack clean. +// Caller holds p.mu; caller is responsible for the Lua state being alive. +func snapshotField(L *golua.State, field string) interface{} { + L.SetTop(0) + defer L.SetTop(0) + L.GetGlobal(pluginGlobal) + if L.IsNil(-1) { + return nil + } + L.GetField(L.GetTop(), field) + if L.Type(-1) != golua.LUA_TTABLE { + return nil + } + var out interface{} + if err := luaToJSON(L, -1, &out); err != nil { + return nil + } + return out +} + +// writeJSONAtomic writes v to path via a temp file + rename, so a reader (or a +// crash) never observes a half-written file. +func writeJSONAtomic(path string, v interface{}) { + b, err := json.MarshalIndent(v, "", " ") + if err != nil { + return + } + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0755); err != nil { + return + } + tmp, err := os.CreateTemp(dir, filepath.Base(path)+".tmp*") + if err != nil { + return + } + tmpName := tmp.Name() + if _, err := tmp.Write(b); err != nil { + tmp.Close() + os.Remove(tmpName) + return + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return + } + if err := os.Rename(tmpName, path); err != nil { + os.Remove(tmpName) + } +} + +// restore re-applies a persisted state (and prices) onto a freshly compiled +// plugin. Called once per plugin after it compiles and registers its hooks. +// +// Ordering matters: this runs AFTER compilation so the plugin's own defaults +// exist, and it OVERWRITES them, so a restart continues from the saved totals +// rather than from the values baked into the .lua source. +func (ps *Plugins) restore(p *Plugin) { + if ps == nil || p == nil || p.state == nil || ps.stateDirDisabled { + return + } + b, err := os.ReadFile(ps.stateFile(p.Info.Name)) + if err != nil { + return // no saved state yet: the compiled-in default stands + } + var saved struct { + Version int `json:"version"` + State interface{} `json:"state"` + Prices interface{} `json:"prices"` + } + if err := json.Unmarshal(b, &saved); err != nil { + // A corrupt file must not stop the gateway: the plugin keeps its + // compiled-in defaults and the operator sees a warning in the log. + ps.logf("plugin %s: ignoring unreadable state file %s: %v", + p.Info.Name, ps.stateFile(p.Info.Name), err) + return + } + p.mu.Lock() + defer p.mu.Unlock() + L := p.state.L + L.SetTop(0) + defer L.SetTop(0) + L.GetGlobal(pluginGlobal) + if L.IsNil(-1) { + return + } + plug := L.GetTop() + if saved.Prices != nil { + pushGoValue(L, saved.Prices) + L.SetField(plug, "prices") + } + if saved.State != nil { + pushGoValue(L, saved.State) + L.SetField(plug, "state") + } + p.persisted = true +} + // SetEnabled turns a plugin's dispatch on or off without touching its file. // // The state is on the Plugin record (not derived from disk) so a disable survives @@ -960,12 +1364,26 @@ func (ps *Plugins) invoke(p *Plugin, fn string, payload map[string]interface{}) if err := L.Call(1, 1); err != nil { return nil, err } + // Read the return value FIRST, then snapshot state for persistence. + // + // The order is load-bearing. markDirtyLocked walks the Lua tables and resets + // the stack, so calling it before the return value is read destroyed the + // hook's answer — every plugin that returned a table silently became a + // plugin that "had no opinion". Reading first costs nothing and keeps the + // documented merge contract intact. + var out map[string]interface{} if L.GetTop() < 1 || L.IsNil(-1) { + ps.markDirtyLocked(p) return nil, nil } - var out map[string]interface{} if err := luaToJSON(L, -1, &out); err != nil { + ps.markDirtyLocked(p) return nil, nil // not a table: treat as "no opinion" } + // A hook that ran at all may have mutated plugin.state, whether or not it + // returned anything. Marking here (not only on a returned table) is what + // makes an accumulating plugin like billing durable: its totals change on + // every call and it returns nil every time. + ps.markDirtyLocked(p) return out, nil }