commit 6657afd7f24752a16272eab467a974571a07d8ca
parent bd783f6193bb61f4b7bd58a99fa99af5306508fb
Author: Sylvia Ivory <git@sivory.net>
Date: Wed, 2 Jul 2025 01:23:36 -0700
Allow cancelling system Lua execution
Diffstat:
5 files changed, 165 insertions(+), 32 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
@@ -728,11 +728,13 @@ dependencies = [
"emmylua_parser",
"lazy_static",
"mlua",
+ "nix",
"nu-ansi-term",
"reedline",
"rexpect",
"rowan",
"send_wrapper",
+ "tempfile",
"thiserror",
]
diff --git a/Cargo.toml b/Cargo.toml
@@ -23,9 +23,11 @@ directories = "6.0.0"
emmylua_parser = "0.10.8"
lazy_static = "1.5.0"
mlua = { version = "0.10.5", features = ["anyhow", "send", "async", "macros"] }
+nix = { version = "0.30.1", features = ["signal"] }
nu-ansi-term = "0.50.1"
reedline = "0.40.0"
rexpect = "0.6.2"
rowan = "0.16.1"
send_wrapper = "0.6.0"
+tempfile = "3.20.0"
thiserror = "2.0.12"
diff --git a/lua/rpc.lua b/lua/rpc.lua
@@ -232,15 +232,72 @@ function rpc.globals()
rpc.respond('globals', _G)
end
+local loadstring_luas = {
+ ['Lua 5.1'] = true,
+ ['Luau'] = true,
+}
+
function rpc.exec(code)
- local l = _VERSION == 'Lua 5.1' and loadstring or load
- local res = l(code, 'repl')
+ local load_fn = loadstring_luas[_VERSION] and loadstring or load
+ local fn = load_fn(code, 'repl')
+
+ if not fn then
+ fn = assert(load_fn('return (' .. code .. ')', 'repl'))
+ end
+
+ local function cleanup()
+ if debug and debug.sethook then
+ debug.sethook()
+ end
+
+ if io and io.open and rpc.cancel_file then
+ local f = io.open(rpc.cancel_file, 'w')
+
+ if f then
+ f:close()
+ end
+ end
+ end
+
+ if debug and debug.sethook and io and io.open and rpc.cancel_file then
+ local function cancel()
+ local f = assert(io.open(rpc.cancel_file, 'r'))
+
+ local data = f:read('*a')
- if not res then
- res = assert(l('return (' .. code .. ')', 'repl'))
+ f:close()
+
+ if data:find('stop') then
+ cleanup()
+ error('cancelled')
+ end
+ end
+
+ -- we can't do every line like in MluaExecutor due to the FS call
+ debug.sethook(cancel, "", 500000)
+ end
+
+ local success, res = pcall(fn)
+
+ cleanup()
+
+ if success then
+ rpc.respond('exec', res)
+ else
+ rpc.respond('error', res)
+ end
+end
+
+function rpc.prepare(file)
+ -- LuaJIT can't stop due to JIT
+ if not io or jit then
+ rpc.respond('cancel', false)
+ return
end
- rpc.respond('exec', res())
+ rpc.cancel_file = file
+
+ rpc.respond('cancel', true)
end
for line in io.stdin:lines() do
diff --git a/luau-shim b/luau-shim
@@ -1,8 +0,0 @@
-#!/bin/bash
-
-# luau-shim -e <rpc code>
-code="$2"
-
-luau << EOF
-$code
-EOF
diff --git a/src/lua.rs b/src/lua.rs
@@ -1,14 +1,20 @@
use std::{
+ io::Write,
process::Command,
sync::{
Arc, RwLock,
- atomic::{AtomicBool, Ordering},
+ atomic::{AtomicBool, AtomicI32, Ordering},
},
};
use mlua::prelude::*;
+use nix::{
+ sys::signal::{Signal, kill},
+ unistd::Pid,
+};
use rexpect::session::{PtySession, spawn_command};
use send_wrapper::SendWrapper;
+use tempfile::NamedTempFile;
use thiserror::Error;
pub trait LuaExecutor: Send + Sync {
@@ -57,9 +63,13 @@ impl LuaExecutor for MluaExecutor {
}
pub struct SystemLuaExecutor {
- process: RwLock<SendWrapper<PtySession>>,
+ session: RwLock<SendWrapper<PtySession>>,
program: String,
lua: Lua,
+
+ cancellation_file: RwLock<Option<NamedTempFile>>,
+ pid: AtomicI32,
+ is_stopping: AtomicBool,
}
#[derive(Debug, Error)]
@@ -68,6 +78,8 @@ pub enum SystemLuaError {
Lua(#[from] LuaError),
#[error("expect error: {0}")]
Expect(#[from] rexpect::error::Error),
+ #[error("io error: {0}")]
+ Io(#[from] std::io::Error),
#[error("restarted system Lua")]
Restarted,
#[error("runtime error")]
@@ -77,19 +89,15 @@ pub enum SystemLuaError {
enum RpcCommand {
Globals,
Exec(String),
+ Prepare(String),
}
impl RpcCommand {
pub fn to_lua(&self) -> String {
- let func = match self {
- Self::Globals => "globals",
- Self::Exec(_) => "exec",
- };
-
- if let Self::Exec(code) = self {
- format!("{func}:{code}")
- } else {
- func.to_string()
+ match self {
+ Self::Globals => String::from("globals"),
+ Self::Exec(code) => format!("exec:{code}"),
+ Self::Prepare(file) => format!("prepare:{file}"),
}
}
}
@@ -98,33 +106,88 @@ const RPC_CODE: &str = include_str!("../lua/rpc.lua");
impl SystemLuaExecutor {
pub fn new(program: &str) -> Result<Self, SystemLuaError> {
+ let (session, file) = Self::obtain_session(program)?;
+ let pid = session.process.child_pid.as_raw();
+
Ok(Self {
- process: RwLock::new(SendWrapper::new(Self::obtain_process(program)?)),
+ session: RwLock::new(SendWrapper::new(session)),
program: program.to_string(),
lua: unsafe { Lua::unsafe_new() },
+ cancellation_file: RwLock::new(file),
+ pid: AtomicI32::new(pid),
+ is_stopping: AtomicBool::new(false),
})
}
- fn obtain_process(program: &str) -> Result<PtySession, SystemLuaError> {
+ fn obtain_session(
+ program: &str,
+ ) -> Result<(PtySession, Option<NamedTempFile>), SystemLuaError> {
let mut cmd = Command::new(program);
cmd.arg("-e");
cmd.arg(RPC_CODE);
- Ok(spawn_command(cmd, None)?)
+ let mut session = spawn_command(cmd, None)?;
+
+ // TODO; should this be in our cache/run dir?
+ let file = NamedTempFile::new()?;
+
+ let prepare = RpcCommand::Prepare(file.path().to_string_lossy().to_string());
+
+ let cmd = prepare.to_lua();
+ session.send_line(&cmd)?;
+
+ let lua = Lua::new();
+
+ loop {
+ let code = session.read_line()?;
+
+ if let Ok(prepare_result) = lua.load(&code).eval::<LuaTable>() {
+ if prepare_result.get::<bool>("data")? {
+ return Ok((session, Some(file)));
+ } else {
+ return Ok((session, None));
+ }
+ }
+ }
+ }
+
+ fn restart_process(&self, session: &mut SendWrapper<PtySession>) -> Result<(), SystemLuaError> {
+ let (pty, file) = Self::obtain_session(&self.program)?;
+ self.pid
+ .store(pty.process.child_pid.as_raw(), Ordering::Relaxed);
+
+ *session = SendWrapper::new(pty);
+
+ let mut cancellation_file = self
+ .cancellation_file
+ .write()
+ .expect("write cancellation_file");
+ *cancellation_file = file;
+
+ Ok(())
}
fn request(&self, command: RpcCommand) -> Result<LuaTable, SystemLuaError> {
- let mut process = self.process.write().expect("write process");
+ self.is_stopping.store(false, Ordering::Relaxed);
+
+ let mut session = self.session.write().expect("write process");
let cmd = command.to_lua();
- process.send_line(&cmd)?;
+
+ if session.send_line(&cmd).is_err() {
+ // killed
+ self.restart_process(&mut session)?;
+
+ return Err(SystemLuaError::Restarted);
+ }
loop {
- let code = match process.read_line() {
+ let code = match session.read_line() {
Ok(code) => code,
Err(rexpect::error::Error::EOF { .. }) => {
- *process = SendWrapper::new(Self::obtain_process(&self.program)?);
+ self.restart_process(&mut session)?;
+
return Err(SystemLuaError::Restarted);
}
x => x?,
@@ -157,6 +220,23 @@ impl LuaExecutor for SystemLuaExecutor {
}
fn cancel(&self) {
- todo!()
+ let mut cancellation_file = self
+ .cancellation_file
+ .write()
+ .expect("write cancellation_file");
+
+ if !self.is_stopping.load(Ordering::Relaxed) {
+ self.is_stopping.store(true, Ordering::Relaxed);
+
+ if let Some(file) = cancellation_file.as_mut() {
+ if file.write_all(b"stop").is_ok() && file.flush().is_ok() {
+ return;
+ }
+ }
+ }
+
+ // Restart process
+ let pid = self.pid.load(Ordering::Relaxed);
+ let _ = kill(Pid::from_raw(pid), Signal::SIGKILL);
}
}
