Implement traceback helper to use in scheduler

packages/lib/src/globals/task.rs

@@ -98,8 +98,7 @@ pub fn create(lua: &'static Lua) -> LuaResult<LuaTable<'static>> {
     // We want the task scheduler to be transparent,
     // but it does not return real lua threads, so
     // we need to override some globals to fake it
-    let globals = lua.globals();
+    let type_original: LuaFunction = lua.named_registry_value("type")?;
-    let type_original: LuaFunction = globals.get("type")?;
     let type_proxy = lua.create_function(move |_, value: LuaValue| {
         if let LuaValue::UserData(u) = &value {
             if u.is::<TaskReference>() {
@@ -108,7 +107,7 @@ pub fn create(lua: &'static Lua) -> LuaResult<LuaTable<'static>> {
         }
         type_original.call(value)
     })?;
-    let typeof_original: LuaFunction = globals.get("typeof")?;
+    let typeof_original: LuaFunction = lua.named_registry_value("typeof")?;
     let typeof_proxy = lua.create_function(move |_, value: LuaValue| {
         if let LuaValue::UserData(u) = &value {
             if u.is::<TaskReference>() {
@@ -117,6 +116,7 @@ pub fn create(lua: &'static Lua) -> LuaResult<LuaTable<'static>> {
         }
         typeof_original.call(value)
     })?;
+    let globals = lua.globals();
     globals.set("type", type_proxy)?;
     globals.set("typeof", typeof_proxy)?;
     // Functions in the built-in coroutine library also need to be

packages/lib/src/lib.rs

@@ -14,6 +14,7 @@ mod tests;
 use crate::utils::formatting::pretty_format_luau_error;
 
 pub use globals::LuneGlobal;
+pub use lua::create_lune_lua;
 
 #[derive(Clone, Debug, Default)]
 pub struct Lune {
@@ -88,18 +89,8 @@ impl Lune {
         script_name: &str,
         script_contents: &str,
     ) -> Result<ExitCode, LuaError> {
-        let lua = Lua::new().into_static();
+        // Create our special lune-flavored Lua object with extra registry values
-        // Store original lua global functions in the registry so we can use
+        let lua = create_lune_lua().expect("Failed to create Lua object");
-        // them later without passing them around and dealing with lifetimes
-        lua.set_named_registry_value("require", lua.globals().get::<_, LuaFunction>("require")?)?;
-        lua.set_named_registry_value("print", lua.globals().get::<_, LuaFunction>("print")?)?;
-        lua.set_named_registry_value("error", lua.globals().get::<_, LuaFunction>("error")?)?;
-        let coroutine: LuaTable = lua.globals().get("coroutine")?;
-        lua.set_named_registry_value("co.thread", coroutine.get::<_, LuaFunction>("running")?)?;
-        lua.set_named_registry_value("co.yield", coroutine.get::<_, LuaFunction>("yield")?)?;
-        lua.set_named_registry_value("co.close", coroutine.get::<_, LuaFunction>("close")?)?;
-        let debug: LuaTable = lua.globals().raw_get("debug")?;
-        lua.set_named_registry_value("dbg.info", debug.get::<_, LuaFunction>("info")?)?;
         // Create our task scheduler and schedule the main thread on it
         let sched = TaskScheduler::new(lua)?.into_static();
         lua.set_app_data(sched);

packages/lib/src/lua/create.rs

@@ -0,0 +1,93 @@
+use mlua::prelude::*;
+
+/*
+    - Level 0 is the call to info
+    - Level 1 is the load call in create() below where we load this into a function
+    - Level 2 is the call to the scheduler, probably, but we can't know for sure so we start at 2
+*/
+const TRACE_IMPL_LUA: &str = r#"
+local lines = {}
+for level = 2, 2^8 do
+    local source, line, name = info(level, "sln")
+    if source then
+        if line then
+            if name and #name > 0 then
+                push(lines, format("    Script '%s', Line %d - function %s", source, line, name))
+            else
+                push(lines, format("    Script '%s', Line %d", source, line))
+            end
+        elseif name and #name > 0 then
+            push(lines, format("    Script '%s' - function %s", source, name))
+        else
+            push(lines, format("    Script '%s'", source))
+        end
+    elseif name then
+        push(lines, format("[Lune] - function %s", source, name))
+    else
+        break
+    end
+end
+if #lines > 0 then
+    push(lines, 1, "Stack Begin")
+    push(lines, "Stack End")
+    return concat(lines, "\n")
+else
+    return nil
+end
+"#;
+
+/**
+    Creates a [`mlua::Lua`] object with certain globals stored in the Lua registry.
+
+    These globals can then be modified safely after constructing Lua using this function.
+
+    ---
+    * `"require"` -> `require`
+    ---
+    * `"print"` -> `print`
+    * `"error"` -> `error`
+    ---
+    * `"type"` -> `type`
+    * `"typeof"` -> `typeof`
+    ---
+    * `"co.thread"` -> `coroutine.running`
+    * `"co.yield"` -> `coroutine.yield`
+    * `"co.close"` -> `coroutine.close`
+    ---
+    * `"dbg.info"` -> `debug.info`
+    * `"dbg.trace"` -> `debug.traceback`
+    ---
+*/
+pub fn create() -> LuaResult<&'static Lua> {
+    let lua = Lua::new().into_static();
+    let globals = &lua.globals();
+    let debug: LuaTable = globals.raw_get("debug")?;
+    let table: LuaTable = globals.raw_get("table")?;
+    let string: LuaTable = globals.raw_get("string")?;
+    let coroutine: LuaTable = globals.get("coroutine")?;
+    // Store original lua global functions in the registry so we can use
+    // them later without passing them around and dealing with lifetimes
+    lua.set_named_registry_value("require", globals.get::<_, LuaFunction>("require")?)?;
+    lua.set_named_registry_value("print", globals.get::<_, LuaFunction>("print")?)?;
+    lua.set_named_registry_value("error", globals.get::<_, LuaFunction>("error")?)?;
+    lua.set_named_registry_value("type", globals.get::<_, LuaFunction>("type")?)?;
+    lua.set_named_registry_value("typeof", globals.get::<_, LuaFunction>("typeof")?)?;
+    lua.set_named_registry_value("co.thread", coroutine.get::<_, LuaFunction>("running")?)?;
+    lua.set_named_registry_value("co.yield", coroutine.get::<_, LuaFunction>("yield")?)?;
+    lua.set_named_registry_value("co.close", coroutine.get::<_, LuaFunction>("close")?)?;
+    lua.set_named_registry_value("dbg.info", debug.get::<_, LuaFunction>("info")?)?;
+    // Create a trace function that can be called to obtain a full stack trace from
+    // lua, this is not possible to do from rust when using our manual scheduler
+    let trace_env = lua.create_table_with_capacity(0, 1)?;
+    trace_env.set("info", debug.get::<_, LuaFunction>("info")?)?;
+    trace_env.set("push", table.get::<_, LuaFunction>("insert")?)?;
+    trace_env.set("concat", table.get::<_, LuaFunction>("concat")?)?;
+    trace_env.set("format", string.get::<_, LuaFunction>("format")?)?;
+    let trace_fn = lua
+        .load(TRACE_IMPL_LUA)
+        .set_environment(trace_env)?
+        .into_function()?;
+    lua.set_named_registry_value("dbg.trace", trace_fn)?;
+    // All done
+    Ok(lua)
+}

packages/lib/src/lua/mod.rs

@@ -1,3 +1,7 @@
+mod create;
+
 pub mod net;
 pub mod stdio;
 pub mod task;
+
+pub use create::create as create_lune_lua;

packages/lib/src/lua/task/result.rs

@@ -94,9 +94,12 @@ impl TaskSchedulerState {
         Returns `true` if the task scheduler is done,
         meaning it has no lua threads left to run, and
         no spawned tasks are running in the background.
+
+        Also returns `true` if a task has requested to exit the process.
     */
     pub fn is_done(&self) -> bool {
-        self.num_blocking == 0 && self.num_futures == 0 && self.num_background == 0
+        self.exit_code.is_some()
+            || (self.num_blocking == 0 && self.num_futures == 0 && self.num_background == 0)
     }
 }
 

tests/task/defer.luau

@@ -39,6 +39,26 @@ end)
 task.wait()
 assert(flag3 == 3, "Defer should run after spawned threads")
 
+-- Delay should be able to be nested
+
+local flag4: boolean = false
+task.delay(0.05, function()
+	local function nested3()
+		task.delay(0.05, function()
+			flag4 = true
+		end)
+	end
+	local function nested2()
+		task.delay(0.05, nested3)
+	end
+	local function nested1()
+		task.delay(0.05, nested2)
+	end
+	nested1()
+end)
+task.wait(0.25)
+assert(flag4, "Defer should work with nesting")
+
 -- Varargs should get passed correctly
 
 local fcheck = require("./fcheck")

tests/task/delay.luau

@@ -26,6 +26,34 @@ assert(flag, "Delay should work with yielding (1)")
 task.wait(0.1)
 assert(not flag2, "Delay should work with yielding (2)")
 
+-- Defer should be able to be nested
+
+local flag4: boolean = false
+task.defer(function()
+	local function nested3()
+		task.defer(function()
+			task.wait(0.05)
+			flag4 = true
+		end)
+	end
+	local function nested2()
+		task.defer(function()
+			task.wait(0.05)
+			nested3()
+		end)
+	end
+	local function nested1()
+		task.defer(function()
+			task.wait(0.05)
+			nested2()
+		end)
+	end
+	task.wait(0.05)
+	nested1()
+end)
+task.wait(0.25)
+assert(flag4, "Defer should work with nesting")
+
 -- Varargs should get passed correctly
 
 local fcheck = require("./fcheck")

tests/task/spawn.luau

@@ -31,6 +31,34 @@ end)
 task.spawn(thread2)
 assert(flag3, "Spawn should run threads made from coroutine.create")
 
+-- Spawn should be able to be nested
+
+local flag4: boolean = false
+task.spawn(function()
+	local function nested3()
+		task.spawn(function()
+			task.wait(0.05)
+			flag4 = true
+		end)
+	end
+	local function nested2()
+		task.spawn(function()
+			task.wait(0.05)
+			nested3()
+		end)
+	end
+	local function nested1()
+		task.spawn(function()
+			task.wait(0.05)
+			nested2()
+		end)
+	end
+	task.wait(0.05)
+	nested1()
+end)
+task.wait(0.25)
+assert(flag4, "Spawn should work with nesting")
+
 -- Varargs should get passed correctly
 
 local fcheck = require("./fcheck")