From feec72bcbd65dfa20136685e6e0c6adf7cfd72d9 Mon Sep 17 00:00:00 2001 From: Alex Orlenko Date: Thu, 6 Nov 2025 23:09:26 +0000 Subject: [PATCH] Reduce number of allocations when calling async function Instead of creating a uniq poller with upvalue on each async call, return future directly and pass it to the poller This also gives about 3-5% perf improvements --- src/state.rs | 10 +++++++--- src/state/raw.rs | 44 ++++++++++++++++++++------------------------ 2 files changed, 27 insertions(+), 27 deletions(-) diff --git a/src/state.rs b/src/state.rs index aa27f2f..7cf2a78 100644 --- a/src/state.rs +++ b/src/state.rs @@ -2176,9 +2176,13 @@ impl Lua { None => unsafe { let lua = self.lock(); let state = lua.state(); - let _sg = StackGuard::with_top(state, 0); - let nvals = ffi::lua_gettop(state); - Poll::Ready(R::from_stack_multi(nvals, &lua)) + let top = ffi::lua_gettop(state); + if top == 0 || ffi::lua_type(state, 1) != ffi::LUA_TUSERDATA { + // This must be impossible scenario if used correctly + return Poll::Ready(R::from_stack_multi(0, &lua)); + } + let _sg = StackGuard::with_top(state, 1); + Poll::Ready(R::from_stack_multi(top - 1, &lua)) }, }) .await diff --git a/src/state/raw.rs b/src/state/raw.rs index 508ac64..b19e871 100644 --- a/src/state/raw.rs +++ b/src/state/raw.rs @@ -1287,7 +1287,7 @@ impl RawLua { } } - unsafe extern "C-unwind" fn call_callback(state: *mut ffi::lua_State) -> c_int { + unsafe extern "C-unwind" fn get_future_callback(state: *mut ffi::lua_State) -> c_int { // Async functions cannot be scoped and therefore destroyed, // so the first upvalue is always valid let upvalue = get_userdata::(state, ffi::lua_upvalueindex(1)); @@ -1301,33 +1301,27 @@ impl RawLua { let extra = XRc::clone(&(*upvalue).extra); let protect = !rawlua.unlikely_memory_error(); push_internal_userdata(state, AsyncPollUpvalue { data: fut, extra }, protect)?; - if protect { - protect_lua!(state, 1, 1, fn(state) { - ffi::lua_pushcclosure(state, poll_future, 1); - })?; - } else { - ffi::lua_pushcclosure(state, poll_future, 1); - } Ok(1) }) } unsafe extern "C-unwind" fn poll_future(state: *mut ffi::lua_State) -> c_int { - let upvalue = get_userdata::(state, ffi::lua_upvalueindex(1)); - callback_error_ext(state, (*upvalue).extra.get(), true, |extra, nargs| { + // Future is always passed in the first argument + let future = get_userdata::(state, 1); + callback_error_ext(state, (*future).extra.get(), true, |extra, nargs| { // Lua ensures that `LUA_MINSTACK` stack spaces are available (after pushing arguments) // The lock must be already held as the future is polled let rawlua = (*extra).raw_lua(); - if nargs == 1 && ffi::lua_tolightuserdata(state, -1) == Lua::poll_terminate().0 { + if nargs == 2 && ffi::lua_tolightuserdata(state, -1) == Lua::poll_terminate().0 { // Destroy the future and terminate the Lua thread - (*upvalue).data.take(); + (*future).data.take(); ffi::lua_pushinteger(state, -1); return Ok(1); } - let fut = &mut (*upvalue).data; + let fut = &mut (*future).data; let mut ctx = Context::from_waker(rawlua.waker()); match fut.as_mut().map(|fut| fut.as_mut().poll(&mut ctx)) { Some(Poll::Pending) => { @@ -1366,7 +1360,7 @@ impl RawLua { } let state = self.state(); - let get_poll = unsafe { + let get_future = unsafe { let _sg = StackGuard::new(state); check_stack(state, 4)?; @@ -1376,10 +1370,10 @@ impl RawLua { push_internal_userdata(state, upvalue, protect)?; if protect { protect_lua!(state, 1, 1, fn(state) { - ffi::lua_pushcclosure(state, call_callback, 1); + ffi::lua_pushcclosure(state, get_future_callback, 1); })?; } else { - ffi::lua_pushcclosure(state, call_callback, 1); + ffi::lua_pushcclosure(state, get_future_callback, 1); } Function(self.pop_ref()) @@ -1398,15 +1392,17 @@ impl RawLua { let coroutine = lua.globals().get::("coroutine")?; // Prepare environment for the async poller - let env = lua.create_table_with_capacity(0, 3)?; - env.set("get_poll", get_poll)?; + let env = lua.create_table_with_capacity(0, 4)?; + env.set("get_future", get_future)?; + env.set("poll", unsafe { lua.create_c_function(poll_future)? })?; env.set("yield", coroutine.get::("yield")?)?; env.set("unpack", unsafe { lua.create_c_function(unpack)? })?; lua.load( r#" - local poll = get_poll(...) - local nres, res, res2 = poll() + local poll, yield = poll, yield + local future = get_future(...) + local nres, res, res2 = poll(future) while true do -- Poll::Ready branch, `nres` is the number of results if nres ~= nil then @@ -1430,13 +1426,13 @@ impl RawLua { -- `res` is a "pending" value -- `yield` can return a signal to drop the future that we should propagate -- to the poller - nres, res, res2 = poll(yield(res)) + nres, res, res2 = poll(future, yield(res)) elseif res2 == 0 then - nres, res, res2 = poll(yield()) + nres, res, res2 = poll(future, yield()) elseif res2 == 1 then - nres, res, res2 = poll(yield(res)) + nres, res, res2 = poll(future, yield(res)) else - nres, res, res2 = poll(yield(unpack(res, res2))) + nres, res, res2 = poll(future, yield(unpack(res, res2))) end end "#,