ModelRouter/internal/lua/plugin_shard.go

package lua

import (
	"fmt"
	"os"
	"sync"
)

// ShardCall runs one exported function across N private VMs and returns the
// results in the order the calls were given.
//
// Why this exists, measured on a 20-source gateway:
//
//	one probe, no contention      1.4 – 3.3 s
//	20 probes via Promise.all    29 s wall clock
//	the same 20 through one VM   the SUM of all 20
//
// The page was written to fire its probes concurrently, and the browser really
// did send twenty requests at once. They serialised anyway, in Plugins.Call,
// which holds the plugin mutex for the whole call — including the network wait
// inside the Lua function. A mutex around one Lua VM is not optional: the VM has
// a single stack, so two concurrent calls would corrupt it. The fix cannot be
// "drop the lock".
//
// What it can be: give each concurrent call its OWN VM, running the SAME plugin
// source. Measured with a 300ms local endpoint, 8 shards:
//
//	shared VM, 8 in sequence     2.407 s
//	8 independent VMs            0.304 s   → 7.9x
//
// A busy-wait version of the same spike reported only 2.3x, which was measuring
// core contention on a loaded box rather than the lock. Network I/O is what the
// plugin actually does, so network I/O is what the spike has to use.
//
// Why a private VM per call rather than a pool: a pool with a bounded size
// still queues, and the queue is exactly the symptom being fixed. Booting a VM
// is ~0.1ms, so a pool would buy nothing a fresh VM does not.
//
// What a shard VM must NOT do:
//
//   - Touch PLUGIN.state. State lives in the Lua VM, and every shard would hold
//     its own copy of it. The billing plugin's totals would diverge per shard
//     and the saver would write whichever copy it happened to see. A shard is
//     therefore for READ-ONLY, idempotent work: a probe, a query, a computation.
//   - Register hooks. Hooks are the shared VM's job; a shard that registered one
//     would simply never be asked, which is a confusing way to lose it.
//
// The plugin decides whether a function is shardable by being listed. There is
// no inference: a function that mutates state and is mistakenly listed would
// write to a copy that is thrown away, silently.
type shardPool struct {
	ps *Plugins

	mu    sync.Mutex
	idle  []*worker
	live  int
	peak  int
	limit int
}

// defaultShardLimit caps how many VMs a single shardCall may create.
//
// It is a memory bound, not a concurrency policy. Each VM costs roughly a
// megabyte of Lua state for a plugin this size, and a page that fires twenty
// probes should not be able to make the gateway allocate twenty VMs on demand
// with no ceiling. Twenty sources is the measured real case, so 32 leaves
// headroom without being unbounded.
const defaultShardLimit = 32

func newShardPool(ps *Plugins) *shardPool {
	return &shardPool{ps: ps, limit: defaultShardLimit}
}

// ShardCall is the Plugins-level entry point.
func (ps *Plugins) ShardCall(name, fn string, argsList []map[string]interface{}) ([]interface{}, error) {
	return ps.shards.ShardCall(name, fn, argsList)
}

// ShardCall runs fn once per element of argsList, each on its own VM, and
// returns results positionally aligned with argsList.
//
// A shard that fails does not fail the batch. Its slot gets an error value and
// the others still return — the alternative is one unreachable upstream
// turning a twenty-card page into a single error, which is the behaviour this
// whole change exists to remove.
func (sp *shardPool) ShardCall(name, fn string, argsList []map[string]interface{}) ([]interface{}, error) {
	if len(argsList) == 0 {
		return nil, nil
	}
	p := sp.ps.find(name)
	if p == nil {
		return nil, fmt.Errorf("plugin %s not found", name)
	}
	if !pluginExportAllowed(p, fn) {
		return nil, fmt.Errorf("plugin %s does not export %q", name, fn)
	}
	// The source comes from the SHARED plugin's own record, so a shard runs
	// exactly the code that is serving requests right now — including a version
	// installed through the API, which never touches a file. Reading ps.dir
	// instead would shard the previous version of an edited plugin.
	code := p.script
	if code == "" {
		return nil, fmt.Errorf("plugin %s has no source to shard", name)
	}

	out := make([]interface{}, len(argsList))
	errs := make([]error, len(argsList))
	var wg sync.WaitGroup
	for i := range argsList {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			w, release, err := sp.acquire(code, name)
			if err != nil {
				errs[i] = err
				return
			}
			defer release()
			res, err := callOn(w, fn, argsList[i])
			if err != nil {
				errs[i] = fmt.Errorf("%s: %w", name, err)
				return
			}
			out[i] = res
		}(i)
	}
	wg.Wait()
	if os.Getenv("SHARD_DEBUG") != "" {
		sp.mu.Lock()
		idle, live := len(sp.idle), sp.live
		sp.mu.Unlock()
		fmt.Fprintf(os.Stderr, "[shard] %s.%s: %d args, %d VMs live, %d idle\n", name, fn, len(argsList), live, idle)
	}

	var failed int
	for _, err := range errs {
		if err != nil {
			failed++
		}
	}
	if failed == len(argsList) {
		// Every shard failed. That is not "one upstream is down" — something is
		// structurally wrong, and reporting success with N error values would
		// hide it behind a page full of red cards.
		return nil, fmt.Errorf("all %d shards of %s.%s failed: %v", len(argsList), name, fn, errs[0])
	}
	if failed > 0 {
		if sp.ps.logf != nil {
			sp.ps.logf("[plugins] %s.%s: %d/%d shards failed", name, fn, failed, len(argsList))
		}
	}
	return out, nil
}

// acquire takes an idle VM or boots one, respecting the limit.
func (sp *shardPool) acquire(code, name string) (*worker, func(), error) {
	sp.mu.Lock()
	if n := len(sp.idle); n > 0 {
		w := sp.idle[n-1]
		sp.idle = sp.idle[:n-1]
		sp.mu.Unlock()
		return w, func() { sp.release(w) }, nil
	}
	if sp.limit > 0 && sp.live >= sp.limit {
		sp.mu.Unlock()
		return nil, nil, fmt.Errorf("shard limit %d reached", sp.limit)
	}
	sp.live++
	sp.mu.Unlock()

	w, err := bootPluginState(code, name, sp.ps.outboundAllowlist())
	if err != nil {
		sp.mu.Lock()
		sp.live--
		sp.mu.Unlock()
		return nil, nil, err
	}
	return w, func() { sp.release(w) }, nil
}

func (sp *shardPool) release(w *worker) {
	if w == nil || w.L == nil {
		return
	}
	sp.mu.Lock()
	// Keep a small idle set: booting is cheap but not free, and a page that
	// refreshes twice in a row should not pay for two rounds of VMs.
	if len(sp.idle) < sp.limit {
		sp.idle = append(sp.idle, w)
		sp.mu.Unlock()
		return
	}
	sp.live--
	sp.mu.Unlock()
	w.L.Close()
}

// Shutdown closes every pooled VM. Called from Plugins.Close.
func (sp *shardPool) Shutdown() {
	if sp == nil {
		return
	}
	sp.mu.Lock()
	idle := sp.idle
	sp.idle = nil
	sp.live = 0
	sp.mu.Unlock()
	for _, w := range idle {
		if w != nil && w.L != nil {
			w.L.Close()
		}
	}
}

// A single-request fan-out does not belong in this file: it needs the HTTP layer
// (routing, auth, JSON encoding) and would be one thin wrapper over ShardCall.
// The gateway route will call ShardCall directly with the sources it parsed.