mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 17:07:59 +00:00
Production oscillated between ~46 MB and ~56 MB RSS with the openai pool cycling 1 -> 10..12 -> 1 states every couple of minutes, while peak_in_use never went above 2. Cause: the batch prewarm sized itself purely on the adapter's ceiling. With max_concurrent summing to 76, growStep is 8, so any two overlapping requests warmed 8 states — 6 more than anything was waiting for. A minute later the janitor correctly reclaimed the surplus, the next pair of overlapping requests warmed 8 again, and the pool churned boot/discard forever. The elasticity was working; the growth signal was simply wrong. Prewarm is now bounded by BOTH limits: the ceiling still caps the step, but the batch never exceeds p.waiting, the number of goroutines actually blocked on the pool. Overlapping-but-not-queued traffic (the common case) creates exactly the states it uses; a genuinely queued burst still ramps in one jump. TestContentionBatchPrewarms is rewritten to queue real waiters instead of relying on the ceiling to imply demand, and TestNoPrewarmWithoutWaiters pins the production shape: two overlapping requests against a 76-wide adapter must create exactly 2 states.
1114 lines
28 KiB
Go
1114 lines
28 KiB
Go
// Package lua implements the adapter runtime: bundles/loads *.lua adapter
|
|
// scripts, exposes json/string helpers, and lets Go call the protocol
|
|
// transform functions plus a build_headers signature hook.
|
|
//
|
|
// 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) 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 (
|
|
"crypto/hmac"
|
|
"crypto/sha256"
|
|
"embed"
|
|
"encoding/base64"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"sort"
|
|
"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).
|
|
golua "github.com/aarzilli/golua/lua"
|
|
)
|
|
|
|
//go:embed adapters/*.lua
|
|
var bundledAdapters embed.FS
|
|
|
|
// adapterGlobal is the reserved global holding the adapter table after the
|
|
// worker boots the script. A reserved name avoids clashing with user code.
|
|
const adapterGlobal = "__llmsproxy_adapter"
|
|
|
|
type APIAdapter struct {
|
|
Name string `json:"name"`
|
|
Version string `json:"version"`
|
|
}
|
|
|
|
// worker wraps one LuaJIT interpreter state. Not goroutine-safe; each worker
|
|
// is used by exactly one goroutine at a time, checked out from its pool.
|
|
type worker struct {
|
|
L *golua.State
|
|
}
|
|
|
|
func (w *worker) close() {
|
|
if w.L != nil {
|
|
w.L.Close()
|
|
}
|
|
}
|
|
|
|
// staticInfo caches the immutable fields of an adapter script, extracted once
|
|
// at load time. Endpoint() and the static headers fallback read from here so
|
|
// they never churn a pooled worker.
|
|
type staticInfo struct {
|
|
name string
|
|
version string
|
|
endpoint string
|
|
headers map[string]string
|
|
}
|
|
|
|
// adapterPool owns the worker states for one adapter. Idle workers are held in
|
|
// 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
|
|
static staticInfo
|
|
|
|
mu sync.Mutex
|
|
cond *sync.Cond
|
|
idle []*worker
|
|
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 the growth step CAP:
|
|
// step <= clamp(ceil(maxW / growDivisor), 1, growStepCap). The actual batch
|
|
// is additionally bounded by how many callers are queued (see growPlanLocked),
|
|
// so a wide adapter may ramp in big jumps but never warms states nobody is
|
|
// waiting for.
|
|
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. It guards against tearing a pool down between
|
|
// two bursts — NOT against ordinary traffic: a pool serving one request at a
|
|
// time still has reclaimable slack, and requiring "no traffic at all" would
|
|
// pin a burst's leftover states forever on any busy gateway. The counter is
|
|
// therefore reset by GROWTH (real contention), not by a mere checkout.
|
|
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, maxW: 1}
|
|
p.cond = sync.NewCond(&p.mu)
|
|
return p
|
|
}
|
|
|
|
// 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.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) {
|
|
L := golua.NewState()
|
|
L.OpenLibs()
|
|
setupGlobals(L)
|
|
if err := L.DoString(p.script); err != nil {
|
|
L.Close()
|
|
return nil, fmt.Errorf("compile adapter %s: %w", p.name, err)
|
|
}
|
|
if L.Type(-1) != golua.LUA_TTABLE {
|
|
L.Close()
|
|
return nil, fmt.Errorf("adapter %s must return a table", p.name)
|
|
}
|
|
L.SetGlobal(adapterGlobal)
|
|
L.SetTop(0)
|
|
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).
|
|
//
|
|
// Two independent limits apply, and BOTH matter:
|
|
//
|
|
// - the adapter's max concurrency caps the step (growStepLocked), so a
|
|
// high-concurrency adapter is allowed to ramp in bigger jumps than a
|
|
// narrow one;
|
|
// - the number of goroutines actually BLOCKED waiting for a state bounds it
|
|
// to real demand.
|
|
//
|
|
// Sizing on the ceiling alone over-provisions badly: an adapter with
|
|
// max_concurrent=76 has a step of 8, so two concurrent requests would warm 8
|
|
// states, and the janitor would throw 7 of them away a minute later — boot,
|
|
// discard, repeat, with RSS oscillating for no benefit. Prewarming only for
|
|
// goroutines that are genuinely queued keeps the ramp cheap without the churn.
|
|
//
|
|
// Caller holds p.mu.
|
|
func (p *adapterPool) growPlanLocked(contended bool) int {
|
|
if !contended || p.waiting <= 0 {
|
|
return 0
|
|
}
|
|
if time.Since(p.lastGrow) < growCooldown {
|
|
return 0
|
|
}
|
|
extra := p.growStepLocked() - 1
|
|
if extra > p.waiting {
|
|
extra = p.waiting // never warm more than the queue needs
|
|
}
|
|
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 {
|
|
if p.closed {
|
|
p.mu.Unlock()
|
|
return nil, fmt.Errorf("adapter %s pool closed", p.name)
|
|
}
|
|
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.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++
|
|
// Growing is the signal that capacity is genuinely short, so any
|
|
// pending shrink decision is stale.
|
|
p.idleRounds = 0
|
|
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--
|
|
}
|
|
}
|
|
|
|
// 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
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|
|
p.idle = append(p.idle, w)
|
|
p.cond.Signal()
|
|
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 {
|
|
w.close()
|
|
}
|
|
}
|
|
|
|
// VM aggregates per-adapter worker pools. All exported methods are safe for
|
|
// concurrent use.
|
|
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{},
|
|
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
|
|
}
|
|
// First-run seeding: a brand-new adapter dir is created and populated with
|
|
// the bundled adapters. If the dir already exists it is treated as
|
|
// authoritative and never rewritten — deleting a file there is a real delete.
|
|
firstRun := false
|
|
if _, err := os.Stat(v.dir); os.IsNotExist(err) {
|
|
firstRun = true
|
|
if err := os.MkdirAll(v.dir, 0755); err != nil {
|
|
return fmt.Errorf("mkdir adapter dir: %w", err)
|
|
}
|
|
}
|
|
if firstRun {
|
|
if err := v.writeBundledAdapters(); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
entries, err := os.ReadDir(v.dir)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, e := range entries {
|
|
if filepath.Ext(e.Name()) != ".lua" {
|
|
continue
|
|
}
|
|
if err := v.LoadAdapter(filepath.Join(v.dir, e.Name())); err != nil {
|
|
fmt.Printf("[lua] preload %s: %v\n", e.Name(), err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
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{}
|
|
v.mu.Unlock()
|
|
for _, p := range pools {
|
|
p.shutdown()
|
|
}
|
|
}
|
|
|
|
// 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 {
|
|
p.setMax(n)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (v *VM) ListAdapters() []APIAdapter {
|
|
v.mu.RLock()
|
|
defer v.mu.RUnlock()
|
|
list := make([]APIAdapter, 0, len(v.pools))
|
|
for name, p := range v.pools {
|
|
list = append(list, APIAdapter{Name: name, Version: p.static.version})
|
|
}
|
|
sort.Slice(list, func(i, j int) bool { return list[i].Name < list[j].Name })
|
|
return list
|
|
}
|
|
|
|
// LoadAdapter compiles and registers an adapter script from a file.
|
|
func (v *VM) LoadAdapter(path string) error {
|
|
data, err := os.ReadFile(path)
|
|
if err != nil {
|
|
return fmt.Errorf("read adapter: %w", err)
|
|
}
|
|
name := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
|
|
return v.LoadAdapterSource(name, string(data))
|
|
}
|
|
|
|
// LoadAdapterSource registers an adapter from source code. Replacing an
|
|
// 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
|
|
}
|
|
maxW := 1
|
|
v.mu.RLock()
|
|
if p, ok := v.pools[static.name]; ok {
|
|
p.mu.Lock()
|
|
maxW = p.maxW
|
|
p.mu.Unlock()
|
|
}
|
|
v.mu.RUnlock()
|
|
|
|
v.mu.Lock()
|
|
if p, ok := v.pools[static.name]; ok {
|
|
p.shutdown()
|
|
}
|
|
p := newAdapterPool(static.name, code, *static)
|
|
p.setMax(maxW)
|
|
v.pools[static.name] = p
|
|
v.mu.Unlock()
|
|
v.startJanitor()
|
|
return nil
|
|
}
|
|
|
|
// RemoveAdapter evicts an adapter and shuts its workers down.
|
|
func (v *VM) RemoveAdapter(name string) {
|
|
v.mu.Lock()
|
|
p := v.pools[name]
|
|
delete(v.pools, name)
|
|
v.mu.Unlock()
|
|
if p != nil {
|
|
p.shutdown()
|
|
}
|
|
}
|
|
|
|
func (v *VM) pool(name string) *adapterPool {
|
|
v.mu.RLock()
|
|
defer v.mu.RUnlock()
|
|
return v.pools[name]
|
|
}
|
|
|
|
func (v *VM) writeBundledAdapters() error {
|
|
// Derive the set from the embedded FS instead of a hardcoded name list so
|
|
// adding an adapter never requires maintaining a second list.
|
|
entries, err := bundledAdapters.ReadDir("adapters")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, e := range entries {
|
|
if e.IsDir() || !strings.HasSuffix(e.Name(), ".lua") {
|
|
continue
|
|
}
|
|
dst := filepath.Join(v.dir, e.Name())
|
|
if _, err := os.Stat(dst); err == nil {
|
|
continue
|
|
}
|
|
data, err := bundledAdapters.ReadFile("adapters/" + e.Name())
|
|
if err != nil {
|
|
continue
|
|
}
|
|
if err := os.WriteFile(dst, data, 0644); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Transform runs adapter.Fn(raw) and returns the resulting string.
|
|
func (v *VM) Transform(name, fn, raw string) (string, error) {
|
|
p := v.pool(name)
|
|
if p == nil {
|
|
return "", fmt.Errorf("adapter %s not loaded", name)
|
|
}
|
|
w, err := p.acquire()
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer p.release(w)
|
|
L := w.L
|
|
|
|
L.SetTop(0)
|
|
defer L.SetTop(0)
|
|
|
|
L.GetGlobal(adapterGlobal)
|
|
if L.IsNil(-1) {
|
|
return "", fmt.Errorf("adapter %s has no table", name)
|
|
}
|
|
L.GetField(-1, fn)
|
|
if L.IsNil(-1) || !L.IsFunction(-1) {
|
|
return "", fmt.Errorf("adapter %s missing %s", name, fn)
|
|
}
|
|
L.PushString(raw)
|
|
if err := L.Call(1, 1); err != nil {
|
|
return "", fmt.Errorf("%s: %w", fn, err)
|
|
}
|
|
return L.ToString(-1), nil
|
|
}
|
|
|
|
// BuildHeaders executes adapter.build_headers(meta). If the adapter defines no
|
|
// build_headers hook, it falls back to the static adapter.headers table.
|
|
func (v *VM) BuildHeaders(name string, meta map[string]interface{}) (map[string]string, error) {
|
|
p := v.pool(name)
|
|
if p == nil {
|
|
return nil, fmt.Errorf("adapter %s not loaded", name)
|
|
}
|
|
|
|
w, err := p.acquire()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer p.release(w)
|
|
L := w.L
|
|
|
|
L.SetTop(0)
|
|
defer L.SetTop(0)
|
|
|
|
L.GetGlobal(adapterGlobal)
|
|
if L.IsNil(-1) {
|
|
return nil, fmt.Errorf("adapter %s has no table", name)
|
|
}
|
|
L.GetField(-1, "build_headers")
|
|
if !L.IsFunction(-1) {
|
|
// no hook: fall back to the static headers captured at load time
|
|
p.mu.Lock()
|
|
hdr := make(map[string]string, len(p.static.headers))
|
|
for k, h := range p.static.headers {
|
|
hdr[k] = h
|
|
}
|
|
p.mu.Unlock()
|
|
return hdr, nil
|
|
}
|
|
|
|
L.SetTop(0)
|
|
L.GetGlobal(adapterGlobal)
|
|
L.GetField(-1, "build_headers")
|
|
pushGoValue(L, meta)
|
|
if err := L.Call(1, 1); err != nil {
|
|
return nil, fmt.Errorf("build_headers: %w", err)
|
|
}
|
|
if L.Type(-1) != golua.LUA_TTABLE {
|
|
return nil, fmt.Errorf("build_headers returned non-table")
|
|
}
|
|
headers := map[string]string{}
|
|
luaTableToMap(L, -1, headers)
|
|
return headers, nil
|
|
}
|
|
|
|
// TransformError executes the optional adapter.transform_error(status, body)
|
|
// hook: per-source condensing of an upstream error response into a short
|
|
// client-facing reason. ok=false means the adapter defines no hook (or it
|
|
// failed) — the caller falls back to the generic Go-side condenser.
|
|
func (v *VM) TransformError(name string, status int, body string) (reason string, ok bool, err error) {
|
|
p := v.pool(name)
|
|
if p == nil {
|
|
return "", false, fmt.Errorf("adapter %s not loaded", name)
|
|
}
|
|
w, err := p.acquire()
|
|
if err != nil {
|
|
return "", false, err
|
|
}
|
|
defer p.release(w)
|
|
L := w.L
|
|
|
|
L.SetTop(0)
|
|
defer L.SetTop(0)
|
|
L.GetGlobal(adapterGlobal)
|
|
if L.IsNil(-1) {
|
|
return "", false, nil
|
|
}
|
|
L.GetField(-1, "transform_error")
|
|
if !L.IsFunction(-1) {
|
|
return "", false, nil
|
|
}
|
|
L.SetTop(0)
|
|
L.GetGlobal(adapterGlobal)
|
|
L.GetField(-1, "transform_error")
|
|
L.PushInteger(int64(status))
|
|
L.PushString(body)
|
|
if callErr := L.Call(2, 1); callErr != nil {
|
|
return "", false, fmt.Errorf("transform_error: %w", callErr)
|
|
}
|
|
if L.Type(-1) != golua.LUA_TSTRING {
|
|
return "", false, nil
|
|
}
|
|
return L.ToString(-1), true, nil
|
|
}
|
|
|
|
func (v *VM) Endpoint(name string) string {
|
|
p := v.pool(name)
|
|
if p == nil {
|
|
return ""
|
|
}
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
return p.static.endpoint
|
|
}
|
|
|
|
// ---- static inspection (compile-once at load time) ----
|
|
|
|
func inspectScript(name, code string) (*staticInfo, error) {
|
|
L := golua.NewState()
|
|
L.OpenLibs()
|
|
setupGlobals(L)
|
|
if err := L.DoString(code); err != nil {
|
|
L.Close()
|
|
return nil, fmt.Errorf("compile adapter: %w", err)
|
|
}
|
|
if L.Type(-1) != golua.LUA_TTABLE {
|
|
L.Close()
|
|
return nil, fmt.Errorf("adapter script must return a table")
|
|
}
|
|
info := &staticInfo{name: name, headers: map[string]string{}}
|
|
if n := readString(L, -1, "name"); n != "" {
|
|
info.name = n
|
|
}
|
|
info.version = readString(L, -1, "version")
|
|
info.endpoint = readString(L, -1, "endpoint")
|
|
if getFieldBatch(L, "headers") {
|
|
luaTableToMap(L, -1, info.headers)
|
|
L.Pop(1)
|
|
} else {
|
|
L.Pop(1)
|
|
}
|
|
L.Close()
|
|
return info, nil
|
|
}
|
|
|
|
// getFieldBatch fetches field of the table at the top of the stack and returns
|
|
// true when the fetched value (now at the top) is a table.
|
|
func getFieldBatch(L *golua.State, field string) bool {
|
|
L.GetField(-1, field)
|
|
return L.Type(-1) == golua.LUA_TTABLE
|
|
}
|
|
|
|
func readString(L *golua.State, tableIdx int, field string) string {
|
|
L.GetField(tableIdx, field)
|
|
defer L.Pop(1)
|
|
if L.Type(-1) == golua.LUA_TNIL || L.Type(-1) == golua.LUA_TNONE {
|
|
return ""
|
|
}
|
|
return normalizeString(L, -1)
|
|
}
|
|
|
|
func normalizeString(L *golua.State, idx int) string {
|
|
if L.Type(idx) == golua.LUA_TNUMBER {
|
|
n := L.ToNumber(idx)
|
|
return strconv.FormatFloat(n, 'f', -1, 64)
|
|
}
|
|
return L.ToString(idx)
|
|
}
|
|
|
|
// ---- shared globals installed into every worker state ----
|
|
|
|
// restorePcall puts the real pcall/xpcall back in place. golua's OpenLibs
|
|
// replaces them with unsafe_ variants (it captures errors via its own panic
|
|
// handler); the built-in adapters rely on pcall for soft error handling.
|
|
func restorePcall(L *golua.State) {
|
|
for _, name := range []string{"pcall", "xpcall"} {
|
|
L.GetGlobal("unsafe_" + name)
|
|
if !L.IsNil(-1) {
|
|
L.SetGlobal(name)
|
|
} else {
|
|
L.Pop(1)
|
|
}
|
|
}
|
|
}
|
|
|
|
func setupGlobals(L *golua.State) {
|
|
restorePcall(L)
|
|
buildJSONTable(L)
|
|
|
|
registerFn(L, "hmac_sha256_hex", func(L *golua.State) int {
|
|
key := L.ToString(1)
|
|
data := L.ToString(2)
|
|
m := hmac.New(sha256.New, []byte(key))
|
|
m.Write([]byte(data))
|
|
L.PushString(hex.EncodeToString(m.Sum(nil)))
|
|
return 1
|
|
})
|
|
registerFn(L, "sha256_hex", func(L *golua.State) int {
|
|
h := sha256.Sum256([]byte(L.ToString(1)))
|
|
L.PushString(hex.EncodeToString(h[:]))
|
|
return 1
|
|
})
|
|
registerFn(L, "base64_encode", func(L *golua.State) int {
|
|
L.PushString(base64.StdEncoding.EncodeToString([]byte(L.ToString(1))))
|
|
return 1
|
|
})
|
|
registerFn(L, "tohex", func(L *golua.State) int {
|
|
L.PushString(hex.EncodeToString([]byte(L.ToString(1))))
|
|
return 1
|
|
})
|
|
registerFn(L, "log", func(L *golua.State) int {
|
|
fmt.Printf("[adapter/%s] %s\n", L.ToString(1), L.ToString(2))
|
|
return 0
|
|
})
|
|
}
|
|
|
|
func buildJSONTable(L *golua.State) {
|
|
L.NewTable()
|
|
for _, name := range []string{"encode", "decode"} {
|
|
var fn golua.LuaGoFunction
|
|
switch name {
|
|
case "encode":
|
|
fn = func(L *golua.State) int {
|
|
b, err := jsonEncode(luaValueToGo(L, 1))
|
|
if err != nil {
|
|
L.PushString("null")
|
|
return 1
|
|
}
|
|
L.PushString(string(b))
|
|
return 1
|
|
}
|
|
case "decode":
|
|
fn = func(L *golua.State) int {
|
|
v, err := jsonDecode(L.ToString(1))
|
|
if err != nil {
|
|
L.PushNil()
|
|
return 1
|
|
}
|
|
pushGoValue(L, v)
|
|
return 1
|
|
}
|
|
}
|
|
L.PushGoClosure(fn)
|
|
L.SetField(-2, name)
|
|
}
|
|
L.SetGlobal("json")
|
|
}
|
|
|
|
func registerFn(L *golua.State, name string, fn golua.LuaGoFunction) {
|
|
L.PushGoClosure(fn)
|
|
L.SetGlobal(name)
|
|
}
|
|
|
|
// ---- value conversions ----
|
|
|
|
func luaValueToGo(L *golua.State, idx int) interface{} {
|
|
switch L.Type(idx) {
|
|
case golua.LUA_TSTRING:
|
|
return L.ToString(idx)
|
|
case golua.LUA_TNUMBER:
|
|
return L.ToNumber(idx)
|
|
case golua.LUA_TBOOLEAN:
|
|
return L.ToBoolean(idx)
|
|
case golua.LUA_TTABLE:
|
|
return luaTableToGoValue(L, idx)
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// luaTableToGo value converts a Lua table to either []interface{} (array 1..n)
|
|
// or map[string]interface{}.
|
|
func luaTableToGoValue(L *golua.State, idx int) interface{} {
|
|
abs := idx
|
|
if idx < 0 {
|
|
abs = L.GetTop() + idx + 1
|
|
}
|
|
m := map[string]interface{}{}
|
|
allInt, n, maxK := true, 0, 0
|
|
L.PushNil()
|
|
for L.Next(abs) != 0 {
|
|
kt := L.Type(-2)
|
|
if kt == golua.LUA_TNUMBER {
|
|
k := int(L.ToNumber(-2))
|
|
if k > maxK {
|
|
maxK = k
|
|
}
|
|
} else {
|
|
allInt = false
|
|
}
|
|
m[keyToString(L, -2)] = luaValueToGo(L, -1)
|
|
L.Pop(1)
|
|
n++
|
|
}
|
|
if allInt && maxK > 0 && n == maxK {
|
|
arr := make([]interface{}, maxK)
|
|
for i := 1; i <= maxK; i++ {
|
|
arr[i-1] = m[strconv.Itoa(i)]
|
|
}
|
|
return arr
|
|
}
|
|
return m
|
|
}
|
|
|
|
func luaTableToMap(L *golua.State, idx int, dst map[string]string) {
|
|
abs := idx
|
|
if idx < 0 {
|
|
abs = L.GetTop() + idx + 1
|
|
}
|
|
L.PushNil()
|
|
for L.Next(abs) != 0 {
|
|
dst[keyToString(L, -2)] = normalizeString(L, -1)
|
|
L.Pop(1)
|
|
}
|
|
}
|
|
|
|
// keyToString reads the stack value at idx as a string without modifying the
|
|
// original (golua's ToString mutates numeric stack values in place, which
|
|
// would break the next lua_next call), then pops the temporary copy.
|
|
func keyToString(L *golua.State, idx int) string {
|
|
L.PushValue(idx)
|
|
defer L.Pop(1)
|
|
return normalizeString(L, -1)
|
|
}
|
|
|
|
func pushGoValue(L *golua.State, v interface{}) {
|
|
switch x := v.(type) {
|
|
case nil:
|
|
L.PushNil()
|
|
case string:
|
|
L.PushString(x)
|
|
case bool:
|
|
L.PushBoolean(x)
|
|
case int:
|
|
L.PushNumber(float64(x))
|
|
case int64:
|
|
L.PushNumber(float64(x))
|
|
case float64:
|
|
L.PushNumber(x)
|
|
case []interface{}:
|
|
L.NewTable()
|
|
for i, it := range x {
|
|
pushGoValue(L, it)
|
|
L.RawSeti(-2, i+1)
|
|
}
|
|
case map[string]interface{}:
|
|
L.NewTable()
|
|
for k, it := range x {
|
|
pushGoValue(L, it)
|
|
L.SetField(-2, k)
|
|
}
|
|
default:
|
|
if b, err := json.Marshal(x); err == nil {
|
|
var iv interface{}
|
|
if json.Unmarshal(b, &iv) == nil {
|
|
pushGoValue(L, iv)
|
|
return
|
|
}
|
|
}
|
|
L.PushNil()
|
|
}
|
|
}
|
|
|
|
func jsonEncode(v interface{}) ([]byte, error) { return json.Marshal(v) }
|
|
func jsonDecode(s string) (interface{}, error) {
|
|
var v interface{}
|
|
if err := json.Unmarshal([]byte(s), &v); err != nil {
|
|
return nil, err
|
|
}
|
|
return v, nil
|
|
}
|