Improve instrumentation and spans, add tracy example

Cargo.lock

@@ -268,6 +268,19 @@ dependencies = [
  "slab",
 ]
 
+[[package]]
+name = "generator"
+version = "0.7.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5cc16584ff22b460a382b7feec54b23d2908d858152e5739a120b949293bd74e"
+dependencies = [
+ "cc",
+ "libc",
+ "log",
+ "rustversion",
+ "windows",
+]
+
 [[package]]
 name = "lazy_static"
 version = "1.4.0"
@@ -302,6 +315,19 @@ version = "0.4.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f"
 
+[[package]]
+name = "loom"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7e045d70ddfbc984eacfa964ded019534e8f6cbf36f6410aee0ed5cefa5a9175"
+dependencies = [
+ "cfg-if",
+ "generator",
+ "scoped-tls",
+ "tracing",
+ "tracing-subscriber",
+]
+
 [[package]]
 name = "luau0-src"
 version = "0.8.0+luau609"
@@ -358,6 +384,7 @@ dependencies = [
  "rustc-hash",
  "tracing",
  "tracing-subscriber",
+ "tracing-tracy",
 ]
 
 [[package]]
@@ -550,6 +577,18 @@ dependencies = [
  "windows-sys 0.52.0",
 ]
 
+[[package]]
+name = "rustversion"
+version = "1.0.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ffc183a10b4478d04cbbbfc96d0873219d962dd5accaff2ffbd4ceb7df837f4"
+
+[[package]]
+name = "scoped-tls"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294"
+
 [[package]]
 name = "semver"
 version = "1.0.21"
@@ -703,6 +742,37 @@ dependencies = [
  "tracing-log",
 ]
 
+[[package]]
+name = "tracing-tracy"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6024d04f84a69fd0d1dc1eee3a2b070bd246530a0582f9982ae487cb6c703614"
+dependencies = [
+ "tracing-core",
+ "tracing-subscriber",
+ "tracy-client",
+]
+
+[[package]]
+name = "tracy-client"
+version = "0.17.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "59fb931a64ff88984f86d3e9bcd1ae8843aa7fe44dd0f8097527bc172351741d"
+dependencies = [
+ "loom",
+ "once_cell",
+ "tracy-client-sys",
+]
+
+[[package]]
+name = "tracy-client-sys"
+version = "0.22.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9d104d610dfa9dd154535102cc9c6164ae1fa37842bc2d9e83f9ac82b0ae0882"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "unicode-ident"
 version = "1.0.12"
@@ -737,6 +807,15 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
 
+[[package]]
+name = "windows"
+version = "0.48.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e686886bc078bc1b0b600cac0147aadb815089b6e4da64016cbd754b6342700f"
+dependencies = [
+ "windows-targets 0.48.5",
+]
+
 [[package]]
 name = "windows-sys"
 version = "0.48.0"

Cargo.toml

@@ -29,6 +29,7 @@ mlua = { git = "https://github.com/mlua-rs/mlua.git", rev = "1754226c7440ec6c194
 async-fs = "2.1"
 async-io = "2.3"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+tracing-tracy = "0.11"
 
 [lints.clippy]
 all = { level = "deny", priority = -3 }
@@ -61,3 +62,7 @@ test = true
 [[example]]
 name = "scheduler_ordering"
 test = true
+
+[[example]]
+name = "tracy"
+test = false

examples/lua/lots_of_threads.luau

@@ -1,8 +1,8 @@
 --!nocheck
 --!nolint UnknownGlobal
 
-local NUM_BATCHES = 100
+local NUM_BATCHES = 10
-local NUM_THREADS = 10_000
+local NUM_THREADS = 100_000
 
 print(`Spawning {NUM_BATCHES * NUM_THREADS} threads split into {NUM_BATCHES} batches\n`)
 
@@ -14,7 +14,7 @@ for i = 1, NUM_BATCHES do
 	local counter = 0
 	for j = 1, NUM_THREADS do
 		spawn(function()
-			sleep()
+			sleep(0.1)
 			counter += 1
 			if counter == NUM_THREADS then
 				spawn(thread)

examples/tracy.rs

@@ -0,0 +1,61 @@
+/*
+    NOTE: This example is the same as "lots_of_threads", but with tracy set up for performance profiling.
+
+    How to run:
+
+    1. Install tracy
+       - Follow the instructions at https://github.com/wolfpld/tracy
+       - Or install via something like homebrew: `brew install tracy`
+    2. Run the server (`tracy`) in a terminal
+    3. Run the example in another terminal
+       - `export RUST_LOG=trace`
+       - `cargo run --example tracy`
+*/
+
+#![allow(clippy::missing_errors_doc)]
+
+use std::time::Duration;
+
+use async_io::{block_on, Timer};
+
+use mlua::prelude::*;
+use mlua_luau_runtime::{Functions, Runtime};
+
+const MAIN_SCRIPT: &str = include_str!("./lua/lots_of_threads.luau");
+
+const ONE_NANOSECOND: Duration = Duration::from_nanos(1);
+
+pub fn main() -> LuaResult<()> {
+    use tracing_subscriber::layer::SubscriberExt;
+    tracing::subscriber::set_global_default(
+        tracing_subscriber::registry().with(tracing_tracy::TracyLayer::default()),
+    )
+    .unwrap();
+
+    let _client = tracing_tracy::client::Client::start();
+
+    // Set up persistent Lua environment
+    let lua = Lua::new();
+    let rt = Runtime::new(&lua);
+    let fns = Functions::new(&lua)?;
+
+    lua.globals().set("spawn", fns.spawn)?;
+    lua.globals().set(
+        "sleep",
+        lua.create_async_function(|_, ()| async move {
+            // Obviously we can't sleep for a single nanosecond since
+            // this uses OS scheduling under the hood, but we can try
+            Timer::after(ONE_NANOSECOND).await;
+            Ok(())
+        })?,
+    )?;
+
+    // Load the main script into the runtime
+    let main = lua.load(MAIN_SCRIPT);
+    rt.push_thread_front(main, ())?;
+
+    // Run until completion
+    block_on(rt.run());
+
+    Ok(())
+}

lib/functions.rs

@@ -112,6 +112,7 @@ impl<'lua> Functions<'lua> {
         let resume_map = result_map.clone();
         let resume =
             lua.create_function(move |lua, (thread, args): (LuaThread, LuaMultiValue)| {
+                let _span = tracing::trace_span!("lua::resume").entered();
                 match thread.resume::<_, LuaMultiValue>(args.clone()) {
                     Ok(v) => {
                         if v.get(0).map(is_poll_pending).unwrap_or_default() {
@@ -163,6 +164,7 @@ impl<'lua> Functions<'lua> {
         let spawn_map = result_map.clone();
         let spawn = lua.create_function(
             move |lua, (tof, args): (LuaThreadOrFunction, LuaMultiValue)| {
+                let _span = tracing::trace_span!("lua::spawn").entered();
                 let thread = tof.into_thread(lua)?;
                 if thread.status() == LuaThreadStatus::Resumable {
                     // NOTE: We need to resume the thread once instantly for correct behavior,
@@ -199,6 +201,7 @@ impl<'lua> Functions<'lua> {
 
         let defer = lua.create_function(
             move |lua, (tof, args): (LuaThreadOrFunction, LuaMultiValue)| {
+                let _span = tracing::trace_span!("lua::defer").entered();
                 let thread = tof.into_thread(lua)?;
                 if thread.status() == LuaThreadStatus::Resumable {
                     defer_queue.push_item(lua, &thread, args)?;
@@ -213,6 +216,7 @@ impl<'lua> Functions<'lua> {
             .get::<_, LuaFunction>("close")?;
         let close_key = lua.create_registry_value(close)?;
         let cancel = lua.create_function(move |lua, thread: LuaThread| {
+            let _span = tracing::trace_span!("lua::cancel").entered();
             let close: LuaFunction = lua.registry_value(&close_key)?;
             match close.call(thread) {
                 Err(LuaError::CoroutineInactive) | Ok(()) => Ok(()),
@@ -224,6 +228,7 @@ impl<'lua> Functions<'lua> {
             (
                 "exit",
                 lua.create_function(|lua, code: Option<u8>| {
+                    let _span = tracing::trace_span!("lua::exit").entered();
                     let code = code.map(ExitCode::from).unwrap_or_default();
                     lua.set_exit_code(code);
                     Ok(())

lib/runtime.rs

@@ -12,7 +12,7 @@ use futures_lite::prelude::*;
 use mlua::prelude::*;
 
 use async_executor::{Executor, LocalExecutor};
-use tracing::{debug, debug_span, instrument, Instrument};
+use tracing::{debug, instrument, trace, trace_span, Instrument};
 
 use crate::{
     error_callback::ThreadErrorCallback,
@@ -272,7 +272,7 @@ impl<'lua> Runtime<'lua> {
         Panics if the given Lua state already has a runtime attached to it.
     */
     #[allow(clippy::too_many_lines)]
-    #[instrument(skip(self))]
+    #[instrument(level = "debug", name = "runtime::run", skip(self))]
     pub async fn run(&self) {
         /*
             Create new executors to use - note that we do not need create multiple executors
@@ -323,8 +323,6 @@ impl<'lua> Runtime<'lua> {
         let fut = async {
             let result_map = self.result_map.clone();
             let process_thread = |thread: LuaThread<'lua>, args| {
-                let span = debug_span!("process_thread");
-                let _guard = span.enter();
                 // NOTE: Thread may have been cancelled from Lua
                 // before we got here, so we need to check it again
                 if thread.status() == LuaThreadStatus::Resumable {
@@ -371,7 +369,7 @@ impl<'lua> Runtime<'lua> {
 
                 // 5
                 let mut num_processed = 0;
-                let span_tick = tracing::debug_span!("tick_executor");
+                let span_tick = trace_span!("runtime::tick");
                 let fut_tick = async {
                     local_exec.tick().await;
                     // NOTE: Try to do as much work as possible instead of just a single tick()
@@ -391,27 +389,34 @@ impl<'lua> Runtime<'lua> {
 
                 // Check if we should exit
                 if self.exit.get().is_some() {
-                    tracing::trace!("exited with code");
+                    debug!("exit signal received");
                     break;
                 }
 
-                // Process spawned threads first, then deferred threads
+                // Process spawned threads first, then deferred threads, then futures
                 let mut num_spawned = 0;
                 let mut num_deferred = 0;
+                let mut num_futures = 0;
+                {
+                    let _span = trace_span!("runtime::drain_spawned").entered();
                     for (thread, args) in self.queue_spawn.drain_items(self.lua) {
                         process_thread(thread, args);
                         num_spawned += 1;
                     }
+                }
+                {
+                    let _span = trace_span!("runtime::drain_deferred").entered();
                     for (thread, args) in self.queue_defer.drain_items(self.lua) {
                         process_thread(thread, args);
                         num_deferred += 1;
                     }
-
+                }
-                // Process spawned futures
+                {
-                let mut num_futs = 0;
+                    let _span = trace_span!("runtime::drain_futures").entered();
                     for fut in fut_queue.drain_items() {
                         local_exec.spawn(fut).detach();
-                    num_futs += 1;
+                        num_futures += 1;
+                    }
                 }
 
                 // Empty executor = we didn't spawn any new Lua tasks
@@ -419,8 +424,8 @@ impl<'lua> Runtime<'lua> {
                 let completed = local_exec.is_empty()
                     && self.queue_spawn.is_empty()
                     && self.queue_defer.is_empty();
-                debug!(
+                trace!(
-                    futures_spawned = num_futs,
+                    futures_spawned = num_futures,
                     futures_processed = num_processed,
                     lua_threads_spawned = num_spawned,
                     lua_threads_deferred = num_deferred,
@@ -434,8 +439,7 @@ impl<'lua> Runtime<'lua> {
 
         // Run the executor inside a span until all lua threads complete
         self.set_status(Status::Running);
-        let span = debug_span!("executor");
+        main_exec.run(fut).await;
-        main_exec.run(fut).instrument(span.or_current()).await;
         self.set_status(Status::Completed);
 
         // Clean up

lib/util.rs

@@ -1,5 +1,6 @@
 use futures_lite::StreamExt;
 use mlua::prelude::*;
+use tracing::instrument;
 
 /**
     Runs a Lua thread until it manually yields (using coroutine.yield), errors, or completes.
@@ -8,6 +9,7 @@ use mlua::prelude::*;
 
     Otherwise returns the values yielded by the thread, or the error that caused it to stop.
 */
+#[instrument(level = "trace", name = "runtime::run_until_yield", skip_all)]
 pub(crate) async fn run_until_yield<'lua>(
     thread: LuaThread<'lua>,
     args: LuaMultiValue<'lua>,