From 6a2c2f588eefa6fde62ba8fafb5cc0ee328d2b09 Mon Sep 17 00:00:00 2001 From: Filip Tibell Date: Tue, 16 Jan 2024 22:23:11 +0100 Subject: [PATCH] Initial commit --- .gitignore | 1 + Cargo.lock | 468 +++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 9 + src/main.rs | 178 ++++++++++++++++++ src/thread_id.rs | 34 ++++ 5 files changed, 690 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 src/main.rs create mode 100644 src/thread_id.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ea8c4bf --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/target diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..6901c73 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,468 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "addr2line" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a30b2e23b9e17a9f90641c7ab1549cd9b44f296d3ccbf309d2863cfe398a0cb" +dependencies = [ + "gimli", +] + +[[package]] +name = "adler" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" + +[[package]] +name = "anyhow" +version = "1.0.79" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "080e9890a082662b09c1ad45f567faeeb47f22b5fb23895fbe1e651e718e25ca" + +[[package]] +name = "autocfg" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" + +[[package]] +name = "backtrace" +version = "0.3.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2089b7e3f35b9dd2d0ed921ead4f6d318c27680d4a5bd167b3ee120edb105837" +dependencies = [ + "addr2line", + "cc", + "cfg-if", + "libc", + "miniz_oxide", + "object", + "rustc-demangle", +] + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "bstr" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c48f0051a4b4c5e0b6d365cd04af53aeaa209e3cc15ec2cdb69e73cc87fbd0dc" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "bytes" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2bd12c1caf447e69cd4528f47f94d203fd2582878ecb9e9465484c4148a8223" + +[[package]] +name = "cc" +version = "1.0.83" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1174fb0b6ec23863f8b971027804a42614e347eafb0a95bf0b12cdae21fc4d0" +dependencies = [ + "libc", +] + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "gimli" +version = "0.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4271d37baee1b8c7e4b708028c57d816cf9d2434acb33a549475f78c181f6253" + +[[package]] +name = "hermit-abi" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d77f7ec81a6d05a3abb01ab6eb7590f6083d08449fe5a1c8b1e620283546ccb7" + +[[package]] +name = "libc" +version = "0.2.152" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13e3bf6590cbc649f4d1a3eefc9d5d6eb746f5200ffb04e5e142700b8faa56e7" + +[[package]] +name = "libloading" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c571b676ddfc9a8c12f1f3d3085a7b163966a8fd8098a90640953ce5f6170161" +dependencies = [ + "cfg-if", + "windows-sys", +] + +[[package]] +name = "lock_api" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c168f8615b12bc01f9c17e2eb0cc07dcae1940121185446edc3744920e8ef45" +dependencies = [ + "autocfg", + "scopeguard", +] + +[[package]] +name = "luau-scheduler-experiments" +version = "0.0.0" +dependencies = [ + "anyhow", + "mlua", + "tokio", +] + +[[package]] +name = "luau0-src" +version = "0.7.11+luau606" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07ffc4945ee953a33cb2b331e00b19e11275fc105c8ac8a977c810597d790f08" +dependencies = [ + "cc", +] + +[[package]] +name = "memchr" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "523dc4f511e55ab87b694dc30d0f820d60906ef06413f93d4d7a1385599cc149" + +[[package]] +name = "miniz_oxide" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7810e0be55b428ada41041c41f32c9f1a42817901b4ccf45fa3d4b6561e74c7" +dependencies = [ + "adler", +] + +[[package]] +name = "mio" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f3d0b296e374a4e6f3c7b0a1f5a51d748a0d34c85e7dc48fc3fa9a87657fe09" +dependencies = [ + "libc", + "wasi", + "windows-sys", +] + +[[package]] +name = "mlua" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "069264935e816c85884b99e88c8b408d6d92e40ae8760f726c983526a53546b5" +dependencies = [ + "bstr", + "libloading", + "mlua-sys", + "num-traits", + "once_cell", + "rustc-hash", +] + +[[package]] +name = "mlua-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4655631a02e3739d014951291ecfa08db49c4da3f7f8c6f3931ed236af5dd78e" +dependencies = [ + "cc", + "cfg-if", + "luau0-src", + "pkg-config", +] + +[[package]] +name = "num-traits" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c" +dependencies = [ + "autocfg", +] + +[[package]] +name = "num_cpus" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "object" +version = "0.32.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6a622008b6e321afc04970976f62ee297fdbaa6f95318ca343e3eebb9648441" +dependencies = [ + "memchr", +] + +[[package]] +name = "once_cell" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" + +[[package]] +name = "parking_lot" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c42a9226546d68acdd9c0a280d17ce19bfe27a46bf68784e4066115788d008e" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-targets", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8afb450f006bf6385ca15ef45d71d2288452bc3683ce2e2cacc0d18e4be60b58" + +[[package]] +name = "pkg-config" +version = "0.3.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69d3587f8a9e599cc7ec2c00e331f71c4e69a5f9a4b8a6efd5b07466b9736f9a" + +[[package]] +name = "proc-macro2" +version = "1.0.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95fc56cda0b5c3325f5fbbd7ff9fda9e02bb00bb3dac51252d2f1bfa1cb8cc8c" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291ec9ab5efd934aaf503a6466c5d5251535d108ee747472c3977cc5acc868ef" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" +dependencies = [ + "bitflags", +] + +[[package]] +name = "rustc-demangle" +version = "0.1.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.195" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63261df402c67811e9ac6def069e4786148c4563f4b50fd4bf30aa370d626b02" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.195" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46fe8f8603d81ba86327b23a2e9cdf49e1255fb94a4c5f297f6ee0547178ea2c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "signal-hook-registry" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" +dependencies = [ + "libc", +] + +[[package]] +name = "smallvec" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2593d31f82ead8df961d8bd23a64c2ccf2eb5dd34b0a34bfb4dd54011c72009e" + +[[package]] +name = "socket2" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5fac59a5cb5dd637972e5fca70daf0523c9067fcdc4842f053dae04a18f8e9" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "syn" +version = "2.0.48" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f3531638e407dfc0814761abb7c00a5b54992b849452a0646b7f65c9f770f3f" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tokio" +version = "1.35.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89b4efa943be685f629b149f53829423f8f5531ea21249408e8e2f8671ec104" +dependencies = [ + "backtrace", + "bytes", + "libc", + "mio", + "num_cpus", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b8a1e28f2deaa14e508979454cb3a223b10b938b45af148bc0986de36f1923b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "unicode-ident" +version = "1.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" + +[[package]] +name = "wasi" +version = "0.11.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" + +[[package]] +name = "windows-sys" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-targets" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" + +[[package]] +name = "windows_i686_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" + +[[package]] +name = "windows_i686_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..567d40d --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "luau-scheduler-experiments" +version = "0.0.0" +edition = "2021" + +[dependencies] +anyhow = "1.0" +tokio = { version = "1.0", features = ["full"] } +mlua = { version = "0.9", features = ["luau", "luau-jit"] } diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..b9e3db9 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,178 @@ +use std::{collections::HashMap, time::Duration}; + +use mlua::prelude::*; + +mod thread_id; +use thread_id::ThreadId; +use tokio::{ + runtime::Runtime as TokioRuntime, + select, spawn, + sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender}, + task::LocalSet, + time::{interval, sleep, Instant, MissedTickBehavior}, +}; + +const NUM_TEST_BATCHES: usize = 20; +const NUM_TEST_THREADS: usize = 50_000; + +const MAIN_CHUNK: &str = r#" +wait(0.001 * math.random()) +"#; + +const WAIT_IMPL: &str = r#" +__scheduler__resumeAfter(...) +coroutine.yield() +"#; + +type RuntimeSender = UnboundedSender; +type RuntimeReceiver = UnboundedReceiver; + +#[derive(Debug, Clone, Copy)] +enum RuntimeMessage { + Resume(ThreadId), + Cancel(ThreadId), + Yield(ThreadId, Duration), +} + +fn main() { + let rt = TokioRuntime::new().unwrap(); + let set = LocalSet::new(); + let _guard = set.enter(); + + let (msg_tx, lua_rx) = unbounded_channel::(); + let (lua_tx, msg_rx) = unbounded_channel::(); + + set.block_on(&rt, async { + select! { + _ = set.spawn_local(lua_main(lua_rx, lua_tx)) => {}, + _ = spawn(sched_main(msg_rx, msg_tx)) => {}, + } + }); +} + +async fn lua_main(mut rx: RuntimeReceiver, tx: RuntimeSender) -> LuaResult<()> { + let lua = Lua::new(); + let g = lua.globals(); + + lua.enable_jit(true); + lua.set_app_data(tx.clone()); + + let send_message = |lua: &Lua, msg: RuntimeMessage| { + lua.app_data_ref::() + .unwrap() + .send(msg) + .unwrap(); + }; + + g.set( + "__scheduler__resumeAfter", + LuaFunction::wrap(move |lua, duration: f64| { + let thread_id = ThreadId::from(lua.current_thread()); + let duration = Duration::from_secs_f64(duration); + send_message(lua, RuntimeMessage::Yield(thread_id, duration)); + Ok(()) + }), + )?; + + g.set( + "__scheduler__cancel", + LuaFunction::wrap(move |lua, thread: LuaThread| { + let thread_id = ThreadId::from(thread); + send_message(lua, RuntimeMessage::Cancel(thread_id)); + Ok(()) + }), + )?; + + g.set("wait", lua.load(WAIT_IMPL).into_function()?)?; + + let mut yielded_threads: HashMap = HashMap::new(); + let mut runnable_threads: HashMap = HashMap::new(); + + let before = Instant::now(); + + let mut throttle = interval(Duration::from_millis(5)); + throttle.set_missed_tick_behavior(MissedTickBehavior::Delay); + + for n in 1..=NUM_TEST_BATCHES { + println!("Running batch {n} of {NUM_TEST_BATCHES}"); + + let main_fn = lua.load(MAIN_CHUNK).into_function()?; + for _ in 0..NUM_TEST_THREADS { + let thread = lua.create_thread(main_fn.clone())?; + runnable_threads.insert(ThreadId::from(&thread), thread); + } + + loop { + // Runnable / yielded threads may be empty because of cancellation + if runnable_threads.is_empty() && yielded_threads.is_empty() { + break; + } + + // Limit this loop to a maximum of 200hz, this lets us improve performance + // by batching more work and not switching between running threads and waiting + // for the next message as often. It may however add another 5 milliseconds of + // latency to something like a web server, but the tradeoff is worth it. + throttle.tick().await; + + // Resume as many threads as possible + for (thread_id, thread) in runnable_threads.drain() { + thread.resume(())?; + if thread.status() == LuaThreadStatus::Resumable { + yielded_threads.insert(thread_id, thread); + } + } + + if yielded_threads.is_empty() { + break; // All threads ran and we don't have any async task that can spawn more + } + + // Wait for at least one message, but try to receive as many as possible + let mut process_message = |message| match message { + RuntimeMessage::Resume(thread_id) => { + if let Some(thread) = yielded_threads.remove(&thread_id) { + runnable_threads.insert(thread_id, thread); + } + } + RuntimeMessage::Cancel(thread_id) => { + yielded_threads.remove(&thread_id); + runnable_threads.remove(&thread_id); + } + _ => unreachable!(), + }; + if let Some(message) = rx.recv().await { + process_message(message); + while let Ok(message) = rx.try_recv() { + process_message(message); + } + } else { + break; // Scheduler exited + } + } + } + + let after = Instant::now(); + println!( + "Ran {} threads in {:?}", + NUM_TEST_BATCHES * NUM_TEST_THREADS, + after - before + ); + + Ok(()) +} + +async fn sched_main(mut rx: RuntimeReceiver, tx: RuntimeSender) -> LuaResult<()> { + while let Some(message) = rx.recv().await { + match message { + RuntimeMessage::Yield(thread_id, duration) => { + let tx = tx.clone(); + spawn(async move { + sleep(duration).await; + let _ = tx.send(RuntimeMessage::Resume(thread_id)); + }); + } + _ => unreachable!(), + } + } + + Ok(()) +} diff --git a/src/thread_id.rs b/src/thread_id.rs new file mode 100644 index 0000000..3201437 --- /dev/null +++ b/src/thread_id.rs @@ -0,0 +1,34 @@ +use mlua::prelude::*; + +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +pub struct ThreadId(usize); + +impl ThreadId { + fn new(value: &LuaThread) -> Self { + // HACK: We rely on the debug format of mlua + // thread refs here, but currently this is the + // only way to get a proper unique id using mlua + let addr_string = format!("{value:?}"); + let addr = addr_string + .strip_prefix("Thread(Ref(0x") + .expect("Invalid thread address format - unknown prefix") + .split_once(')') + .map(|(s, _)| s) + .expect("Invalid thread address format - missing ')'"); + let id = usize::from_str_radix(addr, 16) + .expect("Failed to parse thread address as hexadecimal into usize"); + Self(id) + } +} + +impl From> for ThreadId { + fn from(value: LuaThread) -> Self { + Self::new(&value) + } +} + +impl From<&LuaThread<'_>> for ThreadId { + fn from(value: &LuaThread) -> Self { + Self::new(value) + } +}