feat(meta): add async Lua spawn and sleep support with tests (#18445)

Add comprehensive async functionality to Lua environment:

- Implement async spawn() function using tokio::task::spawn_local
- Add async sleep() function with tokio::time::sleep
- Create LuaTask wrapper with proper join() method
- Use LocalSet for async Lua script execution in metactl

Test coverage includes:
- Basic spawn functionality with concurrent tasks
- Cross-task gRPC client operations
- Sleep timing verification
- Dynamic test discovery system

This enables concurrent Lua script execution within the async
runtime for improved performance and capability.

Author: drmingdrmer

src/meta/binaries/metactl/main.rs

@@ -245,7 +245,11 @@ impl App {
             }
         };
 
-        if let Err(e) = lua.load(&script).exec_async().await {
+        #[allow(clippy::disallowed_types)]
+        let local = tokio::task::LocalSet::new();
+        let res = local.run_until(lua.load(&script).exec_async()).await;
+
+        if let Err(e) = res {
             return Err(anyhow::anyhow!("Lua execution error: {}", e));
         }
         Ok(())

src/meta/control/src/lua_support.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::cell::RefCell;
+use std::rc::Rc;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -25,6 +27,7 @@ use mlua::LuaSerdeExt;
 use mlua::UserData;
 use mlua::UserDataMethods;
 use mlua::Value;
+use tokio::time;
 
 pub struct LuaGrpcClient {
     client: Arc<ClientHandle>,
@@ -64,6 +67,31 @@ impl UserData for LuaGrpcClient {
     }
 }
 
+pub struct LuaTask {
+    handle: Rc<RefCell<Option<tokio::task::JoinHandle<mlua::Value>>>>,
+}
+
+impl UserData for LuaTask {
+    fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
+        methods.add_async_method("join", |_lua, this, ()| async move {
+            let handle_opt = this.handle.borrow_mut().take();
+            match handle_opt {
+                Some(handle) => match handle.await {
+                    Ok(result) => Ok(result),
+                    Err(e) => {
+                        eprintln!("Join error: {}", e);
+                        Ok(mlua::Value::Nil)
+                    }
+                },
+                None => {
+                    eprintln!("Handle already consumed - task was already awaited");
+                    Ok(mlua::Value::Nil)
+                }
+            }
+        });
+    }
+}
+
 pub fn setup_lua_environment(lua: &Lua) -> anyhow::Result<()> {
     // Register new_grpc_client function
     let new_grpc_client = lua
@@ -91,6 +119,43 @@ pub fn setup_lua_environment(lua: &Lua) -> anyhow::Result<()> {
         .set("NULL", Value::NULL)
         .map_err(|e| anyhow::anyhow!("Failed to register NULL constant: {}", e))?;
 
+    // Register spawn function that delegates to tokio::task::spawn_local
+    let spawn_fn = lua
+        .create_function(|_lua, func: mlua::Function| {
+            #[allow(clippy::disallowed_methods)]
+            let handle = tokio::task::spawn_local(async move {
+                match func.call_async::<mlua::Value>(()).await {
+                    Ok(result) => result,
+                    Err(e) => {
+                        eprintln!("Spawned task error: {}", e);
+                        mlua::Value::Nil
+                    }
+                }
+            });
+
+            Ok(LuaTask {
+                handle: Rc::new(RefCell::new(Some(handle))),
+            })
+        })
+        .map_err(|e| anyhow::anyhow!("Failed to create spawn function: {}", e))?;
+
+    lua.globals()
+        .set("spawn", spawn_fn)
+        .map_err(|e| anyhow::anyhow!("Failed to register spawn function: {}", e))?;
+
+    // Register async sleep function
+    let sleep_fn = lua
+        .create_async_function(|_lua, seconds: f64| async move {
+            let duration = Duration::from_secs_f64(seconds);
+            time::sleep(duration).await;
+            Ok(())
+        })
+        .map_err(|e| anyhow::anyhow!("Failed to create sleep function: {}", e))?;
+
+    lua.globals()
+        .set("sleep", sleep_fn)
+        .map_err(|e| anyhow::anyhow!("Failed to register sleep function: {}", e))?;
+
     Ok(())
 }
 

tests/metactl/metactl_utils.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import sys
 import json
 import subprocess
 import time
@@ -11,6 +12,34 @@
 metactl_bin = f"./target/{BUILD_PROFILE}/databend-metactl"
 
 
+def load_lua_util():
+    """Load lua utility functions from lua_util.lua file."""
+    with open("tests/metactl/lua_util.lua", 'r') as f:
+        return f.read()
+
+
+def metactl_run_lua(lua_script=None, lua_filename=None):
+    cmds = [
+        metactl_bin, "lua"
+    ]
+
+    if lua_filename:
+        cmds += [lua_filename]
+        print("metactl_run_lua from_file:", lua_filename, file=sys.stderr)
+    else:
+        print("metactl_run_lua from_stdin:", lua_script, file=sys.stderr)
+
+    result = subprocess.run(cmds, input=lua_script, capture_output=True, text=True)
+
+    print("metactl_run_lua result:", result, file=sys.stderr)
+
+    if result.returncode != 0:
+        raise Exception(result)
+
+    output = result.stdout.strip()
+    return output
+
+
 def metactl_upsert(grpc_addr, key, value):
     """Upsert a key-value pair using the upsert subcommand."""
     result = run_command([

tests/metactl/subcommands/cmd_lua_grpc.py

@@ -84,77 +84,9 @@ def test_lua_grpc_client():
     shutil.rmtree(".databend", ignore_errors=True)
 
 
-def test_lua_grpc_from_file():
-    """Test lua subcommand with gRPC client using file input."""
-    print_title("Test lua subcommand with gRPC client from file")
-
-    # Setup meta service
-    grpc_addr = setup_test_environment()
-
-    lua_util_str = load_lua_util()
-
-    # Create a temporary Lua script file
-    with tempfile.NamedTemporaryFile(mode='w', suffix='.lua', delete=False) as f:
-        f.write(f'''
-{lua_util_str}
-
-local client = new_grpc_client("{grpc_addr}")
-
--- Upsert multiple key-value pairs
-local upsert1_result, upsert1_err = client:upsert("key1", "value1")
-if upsert1_err then
-    print("Key1 upsert error:", upsert1_err)
-end
-
-local upsert2_result, upsert2_err = client:upsert("key2", "value2")
-if upsert2_err then
-    print("Key2 upsert error:", upsert2_err)
-end
-
--- Get and print the values
-local result1, err1 = client:get("key1")
-if err1 then
-    print("Key1 error:", err1)
-else
-    print("Key1 result:", to_string(result1))
-end
-
-local result2, err2 = client:get("key2")
-if err2 then
-    print("Key2 error:", err2)
-else
-    print("Key2 result:", to_string(result2))
-end
-''')
-        lua_file = f.name
-
-    # Run metactl lua with file
-    result = subprocess.run([
-        metactl_bin, "lua",
-        "--file", lua_file
-    ], capture_output=True, text=True, check=True)
-
-    output = result.stdout.strip()
-    print("output:", output)
-
-    expected_output = '''Key1 result:\t{"data"="value1","seq"=1}
-Key2 result:\t{"data"="value2","seq"=2}'''
-
-    # Check if entire output matches expected value
-    assert output == expected_output, f"Expected:\n{expected_output}\n\nActual:\n{output}"
-
-    print("✓ Lua gRPC client file test passed")
-
-    # Only cleanup on success
-    os.unlink(lua_file)
-    kill_databend_meta()
-    shutil.rmtree(".databend", ignore_errors=True)
-
-
 def main():
     """Main function to run all lua gRPC tests."""
     test_lua_grpc_client()
-    test_lua_grpc_from_file()
 
 
 if __name__ == "__main__":

tests/metactl/subcommands/cmd_lua_sleep.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+
+import subprocess
+import time
+from metactl_utils import metactl_bin
+from utils import print_title
+
+def test_lua_sleep():
+    """Test sleep functionality."""
+    print_title("Test sleep functionality")
+
+    lua_script = '''
+print("Before sleep")
+sleep(0.5)
+print("After sleep")
+'''
+
+    start_time = time.time()
+
+    result = subprocess.run([
+        metactl_bin, "lua"
+    ], input=lua_script, capture_output=True, text=True, check=True)
+
+    execution_time = time.time() - start_time
+
+    expected_output = "Before sleep\nAfter sleep"
+    assert result.stdout.strip() == expected_output
+
+    # Verify timing
+    assert execution_time >= 0.5, f"Expected >= 0.5s, got {execution_time:.2f}s"
+    assert execution_time < 2.0, f"Expected < 2.0s, got {execution_time:.2f}s"
+
+    print(f"✓ Sleep test passed ({execution_time:.2f}s)")
+
+
+def main():
+    """Main function to run lua sleep test."""
+    test_lua_sleep()
+
+
+if __name__ == "__main__":
+    main()

tests/metactl/subcommands/cmd_lua_spawn_concurrent.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+
+import subprocess
+import tempfile
+import os
+import shutil
+from metactl_utils import metactl_bin, load_lua_util, metactl_run_lua
+from utils import print_title, kill_databend_meta, start_meta_node
+
+
+def test_spawn_basic():
+    """Test basic spawn functionality without gRPC."""
+    print_title("Test basic spawn functionality")
+
+    lua_util_str = load_lua_util()
+
+    lua_script = f'''
+{lua_util_str}
+
+print("Testing basic spawn functionality...")
+
+local task1 = spawn(function()
+    print("Task 1: Starting")
+    sleep(0.1)
+    print("Task 1: Finished")
+end)
+
+local task2 = spawn(function()
+    print("Task 2: Starting")
+    sleep(0.2)
+    print("Task 2: Finished")
+end)
+
+print("Created both tasks, now joining...")
+
+task1:join()
+print("Task 1 joined")
+
+task2:join()
+print("Task 2 joined")
+
+print("All tasks completed!")
+'''
+
+    output = metactl_run_lua(lua_script=lua_script)
+
+    expected_phrases = [
+        "Testing basic spawn functionality...",
+        "Task 1: Starting",
+        "Task 2: Starting", 
+        "Created both tasks, now joining...",
+        "Task 1: Finished",
+        "Task 1 joined",
+        "Task 2: Finished",
+        "Task 2 joined",
+        "All tasks completed!"
+    ]
+
+    for phrase in expected_phrases:
+        assert phrase in output, f"Expected phrase '{phrase}' not found in output:\n{output}"
+    print("✓ Basic spawn functionality test passed")
+
+
+def main():
+    """Test spawn functionality with multiple concurrent tasks."""
+    test_spawn_basic()
+
+
+if __name__ == "__main__":
+    main()