Use scheduler blocking threadpool instead of tokio

lune-org · Feb 16, 2024 · ede4682 · ede4682
1 parent 3e80e79
commit ede4682
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 27 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -79,6 +79,7 @@ urlencoding = "2.1"
 
 ### RUNTIME
 
+blocking = "1.5"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 tokio = { version = "1.24", features = ["full", "tracing"] }

diff --git a/src/lune/builtins/roblox/mod.rs b/src/lune/builtins/roblox/mod.rs
@@ -1,4 +1,5 @@
 use mlua::prelude::*;
+use mlua_luau_scheduler::LuaSpawnExt;
 use once_cell::sync::OnceCell;
 
 use crate::{
@@ -11,8 +12,6 @@ use crate::{
     },
 };
 
-use tokio::task;
-
 static REFLECTION_DATABASE: OnceCell<ReflectionDatabase> = OnceCell::new();
 
 pub fn create(lua: &Lua) -> LuaResult<LuaTable> {
@@ -41,41 +40,41 @@ async fn deserialize_place<'lua>(
     contents: LuaString<'lua>,
 ) -> LuaResult<LuaValue<'lua>> {
     let bytes = contents.as_bytes().to_vec();
-    let fut = task::spawn_blocking(move || {
+    let fut = lua.spawn_blocking(move || {
         let doc = Document::from_bytes(bytes, DocumentKind::Place)?;
         let data_model = doc.into_data_model_instance()?;
         Ok::<_, DocumentError>(data_model)
     });
-    fut.await.into_lua_err()??.into_lua(lua)
+    fut.await.into_lua_err()?.into_lua(lua)
 }
 
 async fn deserialize_model<'lua>(
     lua: &'lua Lua,
     contents: LuaString<'lua>,
 ) -> LuaResult<LuaValue<'lua>> {
     let bytes = contents.as_bytes().to_vec();
-    let fut = task::spawn_blocking(move || {
+    let fut = lua.spawn_blocking(move || {
         let doc = Document::from_bytes(bytes, DocumentKind::Model)?;
         let instance_array = doc.into_instance_array()?;
         Ok::<_, DocumentError>(instance_array)
     });
-    fut.await.into_lua_err()??.into_lua(lua)
+    fut.await.into_lua_err()?.into_lua(lua)
 }
 
 async fn serialize_place<'lua>(
     lua: &'lua Lua,
     (data_model, as_xml): (LuaUserDataRef<'lua, Instance>, Option<bool>),
 ) -> LuaResult<LuaString<'lua>> {
     let data_model = (*data_model).clone();
-    let fut = task::spawn_blocking(move || {
+    let fut = lua.spawn_blocking(move || {
         let doc = Document::from_data_model_instance(data_model)?;
         let bytes = doc.to_bytes_with_format(match as_xml {
             Some(true) => DocumentFormat::Xml,
             _ => DocumentFormat::Binary,
         })?;
         Ok::<_, DocumentError>(bytes)
     });
-    let bytes = fut.await.into_lua_err()??;
+    let bytes = fut.await.into_lua_err()?;
     lua.create_string(bytes)
 }
 
@@ -84,15 +83,15 @@ async fn serialize_model<'lua>(
     (instances, as_xml): (Vec<LuaUserDataRef<'lua, Instance>>, Option<bool>),
 ) -> LuaResult<LuaString<'lua>> {
     let instances = instances.iter().map(|i| (*i).clone()).collect();
-    let fut = task::spawn_blocking(move || {
+    let fut = lua.spawn_blocking(move || {
         let doc = Document::from_instance_array(instances)?;
         let bytes = doc.to_bytes_with_format(match as_xml {
             Some(true) => DocumentFormat::Xml,
             _ => DocumentFormat::Binary,
         })?;
         Ok::<_, DocumentError>(bytes)
     });
-    let bytes = fut.await.into_lua_err()??;
+    let bytes = fut.await.into_lua_err()?;
     lua.create_string(bytes)
 }
 

diff --git a/src/lune/builtins/serde/compress_decompress.rs b/src/lune/builtins/serde/compress_decompress.rs
@@ -1,9 +1,7 @@
-use lz4_flex::{compress_prepend_size, decompress_size_prepended};
 use mlua::prelude::*;
-use tokio::{
-    io::{copy, BufReader},
-    task,
-};
+
+use lz4_flex::{compress_prepend_size, decompress_size_prepended};
+use tokio::io::{copy, BufReader};
 
 use async_compression::{
     tokio::bufread::{
@@ -100,9 +98,7 @@ pub async fn compress<'lua>(
 ) -> LuaResult<Vec<u8>> {
     if let CompressDecompressFormat::LZ4 = format {
         let source = source.as_ref().to_vec();
-        return task::spawn_blocking(move || compress_prepend_size(&source))
-            .await
-            .into_lua_err();
+        return Ok(blocking::unblock(move || compress_prepend_size(&source)).await);
     }
 
     let mut bytes = Vec::new();
@@ -133,9 +129,8 @@ pub async fn decompress<'lua>(
 ) -> LuaResult<Vec<u8>> {
     if let CompressDecompressFormat::LZ4 = format {
         let source = source.as_ref().to_vec();
-        return task::spawn_blocking(move || decompress_size_prepended(&source))
+        return blocking::unblock(move || decompress_size_prepended(&source))
             .await
-            .into_lua_err()?
             .into_lua_err();
     }
 

diff --git a/src/lune/builtins/stdio/mod.rs b/src/lune/builtins/stdio/mod.rs
@@ -1,10 +1,8 @@
 use mlua::prelude::*;
 
 use dialoguer::{theme::ColorfulTheme, Confirm, Input, MultiSelect, Select};
-use tokio::{
-    io::{self, AsyncWriteExt},
-    task,
-};
+use mlua_luau_scheduler::LuaSpawnExt;
+use tokio::io::{self, AsyncWriteExt};
 
 use crate::lune::util::{
     formatting::{
@@ -55,10 +53,10 @@ async fn stdio_ewrite(_: &Lua, s: LuaString<'_>) -> LuaResult<()> {
     Ok(())
 }
 
-async fn stdio_prompt(_: &Lua, options: PromptOptions) -> LuaResult<PromptResult> {
-    task::spawn_blocking(move || prompt(options))
+async fn stdio_prompt(lua: &Lua, options: PromptOptions) -> LuaResult<PromptResult> {
+    lua.spawn_blocking(move || prompt(options))
         .await
-        .into_lua_err()?
+        .into_lua_err()
 }
 
 fn prompt(options: PromptOptions) -> LuaResult<PromptResult> {