rootspring / worker.rs

// NOTE: Because the mlua crate is not Sync, we can use tokio spawn_local to run the Lua VM in an async context // but pinned to a single thread // // This is highly experimental pub struct LuaWorker { /// A handle that allows stopping the VM inside its tokio localset /// /// This is wrapped in an option to allow destroying the handle when the LuaWorker is dropped pub tx_stop: Option<tokio::sync::oneshot::Sender<()>>, /// A channel used for sending requests to the VM pub tx_msg_recv: tokio::sync::mpsc::Sender<LuaWorkerRequest>, /// A channel that can be used to listen for a response from the VM pub rx_msg_resp: tokio::sync::mpsc::Receiver<LuaWorkerResponse>, } impl LuaWorker { // Executes a Lua script in the Lua VM pub async fn exec( vm: &Lua, template: &str, args: serde_json::Value, ) -> Result<serde_json::Value, base_data::Error> { let f: LuaFunction = vm .load(template) .eval_async() .await .map_err(|e| LuaError::external(e.to_string()))?; let _args = vm .create_table() .map_err(|e| LuaError::external(e.to_string()))?; let args = vm .to_value(&args) .map_err(|e| LuaError::external(e.to_string()))?; _args .set("args", args) .map_err(|e| LuaError::external(e.to_string()))?; let v: LuaValue = f .call_async(_args) .await .map_err(|e| LuaError::external(e.to_string()))?; let v = serde_json::to_value(v).map_err(|e| LuaError::external(e.to_string()))?; Ok(v) } // Spawns a new LuaWorker with the given Lua VM pub fn new(lua: Lua) -> Self { let (tx_stop, rx_stop) = tokio::sync::oneshot::channel(); let (tx_msg_recv, rx_msg_recv) = tokio::sync::mpsc::channel(32); let (tx_msg_resp, rx_msg_resp) = tokio::sync::mpsc::channel(32); let worker = LuaWorker { tx_stop: Some(tx_stop), tx_msg_recv, rx_msg_resp, }; let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); std::thread::spawn(move || { rt.block_on(async move { let lua = Arc::new(lua); let mut rx_stop = rx_stop; let mut rx_msg_recv = rx_msg_recv; let mut tx_msg_resp = tx_msg_resp; loop { let lua = lua.clone(); let tx_msg_resp = tx_msg_resp.clone(); tokio::select! { /*_ = rx_stop => { break; },*/ Some(msg) = rx_msg_recv.recv() => { tokio::task::spawn_local(async move { let lua = lua.clone(); let res = LuaWorker::exec(&lua, &msg.template, msg.args).await; let _ = tx_msg_resp.send(match res { Ok(v) => LuaWorkerResponse::Ok(v), Err(e) => LuaWorkerResponse::Err(e.to_string()), }); }); } } } }); }); worker } } impl Drop for LuaWorker { fn drop(&mut self) { if let Some(sender) = self.tx_stop.take() { let _ = sender.send(()); } } } pub struct LuaWorkerRequest { pub template: String, pub args: serde_json::Value, } pub enum LuaWorkerResponse { Ok(serde_json::Value), Err(String), }