From ecbb9dcef86e2b2a00f570810c5e38bc6a9f3085 Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Sun, 11 Feb 2024 10:14:01 +0100
Subject: [PATCH] Add api to spawn blocking tasks from lua

---
 Cargo.lock    |  1 +
 Cargo.toml    |  1 +
 lib/traits.rs | 89 +++++++++++++++++++++++++++++++++++++++++++++------
 3 files changed, 81 insertions(+), 10 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c3b0077..2027e5d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -376,6 +376,7 @@ dependencies = [
  "async-executor",
  "async-fs",
  "async-io",
+ "blocking",
  "concurrent-queue",
  "derive_more",
  "event-listener",
diff --git a/Cargo.toml b/Cargo.toml
index f3ef0a6..d883a2f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,6 +11,7 @@ categories = ["async"]
 
 [dependencies]
 async-executor = "1.8"
+blocking = "1.5"
 concurrent-queue = "2.4"
 derive_more = "0.99"
 event-listener = "4.0"
diff --git a/lib/traits.rs b/lib/traits.rs
index 13aba08..02c6105 100644
--- a/lib/traits.rs
+++ b/lib/traits.rs
@@ -5,9 +5,9 @@ use std::{
     cell::Cell, future::Future, process::ExitCode, rc::Weak as WeakRc, sync::Weak as WeakArc,
 };
 
-use mlua::prelude::*;
-
 use async_executor::{Executor, Task};
+use mlua::prelude::*;
+use tracing::trace;
 
 use crate::{
     exit::Exit,
@@ -183,7 +183,10 @@ pub trait LuaRuntimeExt<'lua> {
 
         [`Runtime`]: crate::Runtime
     */
-    fn spawn<T: Send + 'static>(&self, fut: impl Future<Output = T> + Send + 'static) -> Task<T>;
+    fn spawn<F, T>(&self, fut: F) -> Task<T>
+    where
+        F: Future<Output = T> + Send + 'static,
+        T: Send + 'static;
 
     /**
         Spawns the given thread-local future on the current executor.
@@ -224,7 +227,52 @@ pub trait LuaRuntimeExt<'lua> {
         }
         ```
     */
-    fn spawn_local(&self, fut: impl Future<Output = ()> + 'static);
+    fn spawn_local<F>(&self, fut: F)
+    where
+        F: Future<Output = ()> + 'static;
+
+    /**
+        Spawns the given blocking function and returns its [`Task`].
+
+        This function will run on a separate thread pool and not block the current executor.
+
+        # Panics
+
+        Panics if called outside of a running [`Runtime`].
+
+        # Example usage
+
+        ```rust
+        use async_io::block_on;
+
+        use mlua::prelude::*;
+        use mlua_luau_runtime::*;
+
+        fn main() -> LuaResult<()> {
+            let lua = Lua::new();
+
+            lua.globals().set(
+                "spawnBlockingTask",
+                lua.create_async_function(|lua, ()| async move {
+                    lua.spawn_blocking(|| {
+                        println!("Hello from blocking task!");
+                    }).await;
+                    Ok(())
+                })?
+            )?;
+
+            let rt = Runtime::new(&lua);
+            rt.push_thread_front(lua.load("spawnBlockingTask()"), ());
+            block_on(rt.run());
+
+            Ok(())
+        }
+        ```
+    */
+    fn spawn_blocking<F, T>(&self, f: F) -> Task<T>
+    where
+        F: FnOnce() -> T + Send + 'static,
+        T: Send + 'static;
 }
 
 impl<'lua> LuaRuntimeExt<'lua> for Lua {
@@ -278,23 +326,44 @@ impl<'lua> LuaRuntimeExt<'lua> for Lua {
         async move { map.listen(id).await }
     }
 
-    fn spawn<T: Send + 'static>(&self, fut: impl Future<Output = T> + Send + 'static) -> Task<T> {
+    fn spawn<F, T>(&self, fut: F) -> Task<T>
+    where
+        F: Future<Output = T> + Send + 'static,
+        T: Send + 'static,
+    {
         let exec = self
             .app_data_ref::<WeakArc<Executor>>()
-            .expect("futures can only be spawned within a runtime")
+            .expect("tasks can only be spawned within a runtime")
             .upgrade()
             .expect("executor was dropped");
-        tracing::trace!("spawning future on executor");
+        trace!("spawning future on executor");
         exec.spawn(fut)
     }
 
-    fn spawn_local(&self, fut: impl Future<Output = ()> + 'static) {
+    fn spawn_local<F>(&self, fut: F)
+    where
+        F: Future<Output = ()> + 'static,
+    {
         let queue = self
             .app_data_ref::<WeakRc<FuturesQueue>>()
-            .expect("futures can only be spawned within a runtime")
+            .expect("tasks can only be spawned within a runtime")
             .upgrade()
             .expect("executor was dropped");
-        tracing::trace!("spawning local future on executor");
+        trace!("spawning local task on executor");
         queue.push_item(fut);
     }
+
+    fn spawn_blocking<F, T>(&self, f: F) -> Task<T>
+    where
+        F: FnOnce() -> T + Send + 'static,
+        T: Send + 'static,
+    {
+        let exec = self
+            .app_data_ref::<WeakArc<Executor>>()
+            .expect("tasks can only be spawned within a runtime")
+            .upgrade()
+            .expect("executor was dropped");
+        trace!("spawning blocking task on executor");
+        exec.spawn(blocking::unblock(f))
+    }
 }