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
This commit is contained in:
Alex Orlenko
2025-11-06 23:09:26 +00:00
parent 0611906c6a
commit feec72bcbd
2 changed files with 27 additions and 27 deletions
+7 -3
View File
@@ -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
+20 -24
View File
@@ -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::<AsyncCallbackUpvalue>(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::<AsyncPollUpvalue>(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::<AsyncPollUpvalue>(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::<Table>("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::<Function>("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
"#,