From c091b05f6c447ce5313b5cb0055dc764f225702a Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Wed, 23 Apr 2025 19:18:08 +0200
Subject: [PATCH] Guarantee that wait/sleep functions always yield in both task
 library and scheduler examples

---
 crates/lune-std-task/src/lib.rs               | 14 ++++++--
 .../examples/basic_sleep.rs                   |  5 +++
 .../examples/lots_of_threads.rs               |  4 +++
 test.luau                                     | 35 +++++++++++++++++++
 4 files changed, 56 insertions(+), 2 deletions(-)
 create mode 100644 test.luau

diff --git a/crates/lune-std-task/src/lib.rs b/crates/lune-std-task/src/lib.rs
index aa8baaa..2e0523d 100644
--- a/crates/lune-std-task/src/lib.rs
+++ b/crates/lune-std-task/src/lib.rs
@@ -5,7 +5,10 @@ use std::time::Duration;
 use mlua::prelude::*;
 use mlua_luau_scheduler::Functions;
 
-use tokio::time::{sleep, Instant};
+use tokio::{
+    task::yield_now,
+    time::{sleep, Instant},
+};
 
 use lune_utils::TableBuilder;
 
@@ -49,7 +52,14 @@ return defer(function(...)
 end, ...)
 ";
 
-async fn wait(_: Lua, secs: Option<f64>) -> LuaResult<f64> {
+async fn wait(lua: Lua, secs: Option<f64>) -> LuaResult<f64> {
+    // NOTE: We must guarantee that the task.wait API always yields
+    // from a lua perspective, even if sleep/timer completes instantly
+    yield_now().await;
+    wait_inner(lua, secs).await
+}
+
+async fn wait_inner(_: Lua, secs: Option<f64>) -> LuaResult<f64> {
     let duration = Duration::from_secs_f64(secs.unwrap_or_default());
 
     let before = Instant::now();
diff --git a/crates/mlua-luau-scheduler/examples/basic_sleep.rs b/crates/mlua-luau-scheduler/examples/basic_sleep.rs
index 6db84b2..71eda28 100644
--- a/crates/mlua-luau-scheduler/examples/basic_sleep.rs
+++ b/crates/mlua-luau-scheduler/examples/basic_sleep.rs
@@ -4,6 +4,7 @@
 use std::time::{Duration, Instant};
 
 use async_io::{block_on, Timer};
+use futures_lite::future::yield_now;
 
 use mlua::prelude::*;
 use mlua_luau_scheduler::Scheduler;
@@ -22,6 +23,10 @@ pub fn main() -> LuaResult<()> {
     lua.globals().set(
         "sleep",
         lua.create_async_function(|_, duration: f64| async move {
+            // Guarantee that the coroutine that calls this sleep function
+            // always yields, even if the timer completes without doing so
+            yield_now().await;
+            // We may then sleep as normal
             let before = Instant::now();
             let after = Timer::after(Duration::from_secs_f64(duration)).await;
             Ok((after - before).as_secs_f64())
diff --git a/crates/mlua-luau-scheduler/examples/lots_of_threads.rs b/crates/mlua-luau-scheduler/examples/lots_of_threads.rs
index 48e07db..6a96e0d 100644
--- a/crates/mlua-luau-scheduler/examples/lots_of_threads.rs
+++ b/crates/mlua-luau-scheduler/examples/lots_of_threads.rs
@@ -4,6 +4,7 @@
 use std::time::Duration;
 
 use async_io::{block_on, Timer};
+use futures_lite::future::yield_now;
 
 use mlua::prelude::*;
 use mlua_luau_scheduler::{Functions, Scheduler};
@@ -28,6 +29,9 @@ pub fn main() -> LuaResult<()> {
     lua.globals().set(
         "sleep",
         lua.create_async_function(|_, ()| async move {
+            // Guarantee that the coroutine that calls this sleep function
+            // always yields, even if the timer completes without doing so
+            yield_now().await;
             // 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;
diff --git a/test.luau b/test.luau
new file mode 100644
index 0000000..e095075
--- /dev/null
+++ b/test.luau
@@ -0,0 +1,35 @@
+local task = require("@lune/task")
+
+local started = os.clock()
+
+local amount = 400000
+local batches = 5
+local per_batch = amount / batches
+
+for current = 1, batches do
+	local thread = coroutine.running()
+
+	print(`Batch {current} / {batches}`)
+
+	for i = 1, per_batch do
+		--print("Spawning thread #" .. i)
+		task.spawn(function()
+			print("[BEFORE] Thread number", i)
+			task.wait(0.1)
+			--_TEST_ASYNC_WORK(0.1)
+			print("[AFTER] Thread number", i)
+			if i == per_batch then
+				print("Last thread in batch #" .. current)
+				assert(
+					coroutine.status(thread) == "suspended",
+					`Thread {i} has status {coroutine.status(thread)}`
+				)
+				task.spawn(thread)
+			end
+		end)
+	end
+
+	coroutine.yield()
+end
+local took = os.clock() - started
+print(`Spawned {amount} sleeping threads in {took}s`)