mlua_swarm/middleware/
lua_layer.rs

1//! LuaSpawnerLayer — inserts middleware written in Lua into the `SpawnerStack`.
2//!
3//! Shape:
4//!   - `before_src` = a Lua source that evaluates to a function. Called
5//!     immediately before the spawn as `function(ctx_table) ... end`. If it
6//!     raises, the spawn is rejected with
7//!     `SpawnError::RejectedByMiddleware` (a gate).
8//!   - `after_src` = a Lua source that evaluates to
9//!     `function(ctx_table, result_table) return result' end`, called after
10//!     the worker finishes; the return value becomes the new result flowing
11//!     downstream (a transform).
12//!
13//! # Implementation axis
14//!
15//! `mlua::Lua` is `!Send`, so the initial form built a fresh `Lua::new()`
16//! per call and dropped it (per-call VM). Under high-frequency spawn the
17//! overhead becomes visible. This version switches to
18//! `mlua-isle::AsyncIslePool` (thread-isolated Lua VM + async + pool) and
19//! reuses VMs from the pool. The fully-async chain stays the same shape
20//! (no `block_on` / `spawn_blocking`); the `!Send` constraint is resolved
21//! inside the isle so it rides on the caller's tokio runtime directly.
22
23use crate::core::ctx::Ctx;
24use crate::core::engine::Engine;
25use crate::middleware::SpawnerLayer;
26use crate::types::{CapToken, StepId};
27use crate::worker::adapter::{SpawnError, SpawnerAdapter, WorkerError};
28use crate::worker::{wrap_join, Worker};
29use async_trait::async_trait;
30use mlua::LuaSerdeExt;
31use mlua_isle::{AsyncIslePool, IsleError, LuaErrorKind, LuaFailure, PoolConfig, PoolStrategy};
32use serde_json::Value;
33use std::sync::{Arc, OnceLock};
34
35/// Default pool config (4 warm VMs reused). Callers that want their own
36/// pool can swap it in via `LuaMiddleware::with_pool`.
37fn default_pool_config() -> PoolConfig {
38    PoolConfig {
39        max_size: 4,
40        strategy: PoolStrategy::Warm,
41    }
42}
43
44fn build_default_pool() -> Arc<AsyncIslePool> {
45    Arc::new(
46        AsyncIslePool::new(|_lua| Ok(()), default_pool_config())
47            .expect("AsyncIslePool::new (no-op factory) must succeed"),
48    )
49}
50
51/// `SpawnerLayer` that runs Lua source as a before-gate and/or an
52/// after-transform around a spawn, executed on a pooled `AsyncIslePool`
53/// VM. See the module doc for the exact function shapes expected of
54/// `before_src` / `after_src`.
55#[derive(Clone, Default)]
56pub struct LuaMiddleware {
57    before_src: Option<String>,
58    after_src: Option<String>,
59    pool: Option<Arc<AsyncIslePool>>,
60}
61
62impl LuaMiddleware {
63    /// Empty layer — no before/after hooks, default pool.
64    pub fn new() -> Self {
65        Self::default()
66    }
67    /// Sets the before-hook source (`function(ctx_table) ... end`).
68    /// Raising from this function rejects the spawn.
69    pub fn before(mut self, src: impl Into<String>) -> Self {
70        self.before_src = Some(src.into());
71        self
72    }
73    /// Sets the after-hook source
74    /// (`function(ctx_table, result_table) return result end`). Its
75    /// return value replaces the worker's result.
76    pub fn after(mut self, src: impl Into<String>) -> Self {
77        self.after_src = Some(src.into());
78        self
79    }
80    /// Inject an externally-built AsyncIslePool. Useful when the caller
81    /// wants to share a pool with another Lua layer or carry VM
82    /// initialisation across calls.
83    pub fn with_pool(mut self, pool: Arc<AsyncIslePool>) -> Self {
84        self.pool = Some(pool);
85        self
86    }
87}
88
89impl SpawnerLayer for LuaMiddleware {
90    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
91        // Pool resolution: caller injection wins over the process-wide default (built lazily, once).
92        static DEFAULT_POOL: OnceLock<Arc<AsyncIslePool>> = OnceLock::new();
93        let pool = self
94            .pool
95            .clone()
96            .unwrap_or_else(|| DEFAULT_POOL.get_or_init(build_default_pool).clone());
97        Arc::new(LuaWrapped {
98            inner,
99            before_src: self.before_src.clone(),
100            after_src: self.after_src.clone(),
101            pool,
102        })
103    }
104}
105
106struct LuaWrapped {
107    inner: Arc<dyn SpawnerAdapter>,
108    before_src: Option<String>,
109    after_src: Option<String>,
110    pool: Arc<AsyncIslePool>,
111}
112
113/// Serializable view of `Ctx` handed to Lua functions (`Arc<dyn>` fields are excluded since they cannot serde).
114fn ctx_to_serializable(ctx: &Ctx) -> Value {
115    serde_json::json!({
116        "task_id": ctx.task_id.as_str(),
117        "attempt": ctx.attempt,
118        "agent": ctx.agent,
119        "operator": {
120            "kind": format!("{:?}", ctx.operator.kind),
121            "id": ctx.operator.id,
122        },
123    })
124}
125
126/// A JSON (de)serialization failure on the hook's ctx / result, reported as
127/// a Lua-side conversion failure so it surfaces the same way a Lua error does.
128fn json_failure(e: serde_json::Error) -> IsleError {
129    IsleError::Lua(LuaFailure::new(LuaErrorKind::Conversion, e.to_string()))
130}
131
132/// Helper that returns a closure running the "before hook" inside AsyncIsle.exec.
133/// Return value is fixed to the literal string "ok" — success / failure is
134/// expressed through IsleError.
135fn make_before_exec(
136    src: String,
137    ctx_json: String,
138) -> impl FnOnce(&mlua::Lua) -> Result<String, IsleError> + Send + 'static {
139    move |lua| {
140        let ctx_val: Value = serde_json::from_str(&ctx_json).map_err(json_failure)?;
141        let ctx_lua: mlua::Value = lua.to_value(&ctx_val)?;
142        let f: mlua::Function = lua.load(&src).eval()?;
143        let _: mlua::Value = f.call(ctx_lua)?;
144        Ok("ok".to_string())
145    }
146}
147
148/// Returns a closure that runs the "after hook" inside AsyncIsle.exec.
149/// Return value is the new result as a JSON string (`{"value": ..., "ok": bool}`).
150fn make_after_exec(
151    src: String,
152    ctx_json: String,
153    result_json: String,
154) -> impl FnOnce(&mlua::Lua) -> Result<String, IsleError> + Send + 'static {
155    move |lua| {
156        let ctx_val: Value = serde_json::from_str(&ctx_json).map_err(json_failure)?;
157        let result_val: Value = serde_json::from_str(&result_json).map_err(json_failure)?;
158        let ctx_lua: mlua::Value = lua.to_value(&ctx_val)?;
159        let result_lua: mlua::Value = lua.to_value(&result_val)?;
160        let f: mlua::Function = lua.load(&src).eval()?;
161        let returned: mlua::Value = f.call((ctx_lua, result_lua))?;
162        let new_result: Value = lua.from_value(returned)?;
163        serde_json::to_string(&new_result).map_err(json_failure)
164    }
165}
166
167#[async_trait]
168impl SpawnerAdapter for LuaWrapped {
169    async fn spawn(
170        &self,
171        engine: &Engine,
172        ctx: &Ctx,
173        task_id: StepId,
174        attempt: u32,
175        token: CapToken,
176    ) -> Result<Box<dyn Worker>, SpawnError> {
177        // ─── before hook (= pool checkout → exec → return) ────────────────
178        if let Some(src) = &self.before_src {
179            let ctx_json = serde_json::to_string(&ctx_to_serializable(ctx))
180                .map_err(|e| SpawnError::Internal(format!("ctx serialize: {e}")))?;
181            let isle = self
182                .pool
183                .checkout()
184                .await
185                .map_err(|e| SpawnError::Internal(format!("isle pool checkout: {e}")))?;
186            let f = make_before_exec(src.clone(), ctx_json);
187            isle.exec(f)
188                .await
189                .map_err(|e| SpawnError::RejectedByMiddleware(format!("lua before: {e}")))?;
190            // The isle is dropped here — either returned to the pool (Warm) or shut down (Cold).
191        }
192
193        let engine_clone = engine.clone();
194        let token_clone = token.clone();
195        let task_id_clone = task_id.clone();
196        let handle = self
197            .inner
198            .spawn(engine, ctx, task_id, attempt, token)
199            .await?;
200
201        // ─── after hook ───────────────────────────────────────────────────
202        let Some(after_src) = self.after_src.clone() else {
203            return Ok(handle);
204        };
205        let ctx_val = ctx_to_serializable(ctx);
206        Ok(wrap_completion_with_lua_pool(
207            handle,
208            after_src,
209            ctx_val,
210            self.pool.clone(),
211            engine_clone,
212            token_clone,
213            task_id_clone,
214            attempt,
215        ))
216    }
217}
218
219/// Helper that wraps the completion signal and drives the Lua after-hook
220/// through the pool. Follows the signal-only design: pulls the value from
221/// `engine.output_tail`, and pushes the post-Lua `{value, ok}` (Lua-wire
222/// JSON) as an override Final via `engine.submit_output`.
223#[allow(clippy::too_many_arguments)]
224fn wrap_completion_with_lua_pool(
225    handle: Box<dyn Worker>,
226    after_src: String,
227    ctx_val: Value,
228    pool: Arc<AsyncIslePool>,
229    engine: Engine,
230    token: crate::types::CapToken,
231    task_id: StepId,
232    attempt: u32,
233) -> Box<dyn Worker> {
234    wrap_join(handle, move |signal| async move {
235        match signal {
236            Ok(()) => {
237                let _ = apply_lua_after_pool(
238                    &after_src, &ctx_val, &pool, &engine, &token, &task_id, attempt,
239                )
240                .await;
241                Ok(())
242            }
243            Err(e) => Err(e),
244        }
245    })
246}
247
248async fn apply_lua_after_pool(
249    after_src: &str,
250    ctx_val: &Value,
251    pool: &AsyncIslePool,
252    engine: &Engine,
253    token: &crate::types::CapToken,
254    task_id: &StepId,
255    attempt: u32,
256) -> Result<(), WorkerError> {
257    // Pull the existing Final from the tail (only Inline is fed through Lua; FileRef passes through as-is).
258    let tail = engine.output_tail(task_id, attempt).await;
259    let (value, ok) = match tail.iter().rev().find_map(|ev| match ev {
260        crate::worker::output::OutputEvent::Final {
261            content: crate::worker::output::ContentRef::Inline { value },
262            ok,
263        } => Some((value.clone(), *ok)),
264        _ => None,
265    }) {
266        Some(v) => v,
267        None => return Ok(()), // No Inline Final: do nothing.
268    };
269
270    let ctx_json = serde_json::to_string(ctx_val)
271        .map_err(|e| WorkerError::Failed(format!("ctx serialize: {e}")))?;
272    let result_json = serde_json::to_string(&serde_json::json!({"value": value, "ok": ok}))
273        .map_err(|e| WorkerError::Failed(format!("result serialize: {e}")))?;
274
275    let isle = pool
276        .checkout()
277        .await
278        .map_err(|e| WorkerError::Failed(format!("isle pool checkout: {e}")))?;
279    let f = make_after_exec(after_src.to_string(), ctx_json, result_json);
280    let new_json = isle
281        .exec(f)
282        .await
283        .map_err(|e| WorkerError::Failed(format!("lua after: {e}")))?;
284    let new_result: Value = serde_json::from_str(&new_json)
285        .map_err(|e| WorkerError::Failed(format!("lua after decode: {e}")))?;
286    let new_value = new_result.get("value").cloned().unwrap_or(Value::Null);
287    let new_ok = new_result.get("ok").and_then(|v| v.as_bool()).unwrap_or(ok);
288
289    // Push the override Final to the engine so the downstream dispatch's rev().find picks up the latest.
290    let _ = engine
291        .submit_output(
292            token,
293            task_id,
294            attempt,
295            crate::worker::output::OutputEvent::Final {
296                content: crate::worker::output::ContentRef::Inline { value: new_value },
297                ok: new_ok,
298            },
299        )
300        .await;
301    Ok(())
302}
mlua_swarm/middleware/lua_layer.rs

mlua_swarm/middleware/
lua_layer.rs
