Allocate Waker slot in Registry in re-use it (instead of creating new userdata for it)

This commit is contained in:
Alex Orlenko 2021-05-02 13:04:02 +01:00
parent 67bc0b1196
commit 26d8d899f2
2 changed files with 22 additions and 14 deletions

View File

@ -317,7 +317,12 @@ impl Lua {
{
init_gc_metatable_for::<AsyncCallback>(state, None)?;
init_gc_metatable_for::<LocalBoxFuture<Result<MultiValue>>>(state, None)?;
init_gc_metatable_for::<Waker>(state, None)?;
init_gc_metatable_for::<Option<Waker>>(state, None)?;
// Create empty Waker slot
push_gc_userdata::<Option<Waker>>(state, None)?;
let waker_key = &WAKER_REGISTRY_KEY as *const u8 as *const c_void;
ffi::safe::lua_rawsetp(state, ffi::LUA_REGISTRYINDEX, waker_key)?;
}
// Init serde metatables
@ -1760,14 +1765,14 @@ impl Lua {
}
let lua = &mut *lua;
let mut waker = noop_waker();
// Try to get an outer poll waker
let waker_key = &WAKER_REGISTRY_KEY as *const u8 as *const c_void;
ffi::lua_rawgetp(state, ffi::LUA_REGISTRYINDEX, waker_key);
if let Some(w) = get_gc_userdata::<Waker>(state, -1).as_ref() {
waker = (*w).clone();
}
let waker = match get_gc_userdata::<Option<Waker>>(state, -1).as_ref() {
Some(Some(waker)) => waker.clone(),
_ => noop_waker(),
};
ffi::lua_pop(state, 1);
let mut ctx = Context::from_waker(&waker);

View File

@ -12,13 +12,14 @@ use {
crate::{
error::ExternalError,
lua::{ASYNC_POLL_PENDING, WAKER_REGISTRY_KEY},
util::push_gc_userdata,
util::get_gc_userdata,
value::Value,
},
futures_core::{future::Future, stream::Stream},
std::{
cell::RefCell,
marker::PhantomData,
mem,
os::raw::c_void,
pin::Pin,
task::{Context, Poll, Waker},
@ -311,20 +312,21 @@ fn is_poll_pending(val: &MultiValue) -> bool {
}
#[cfg(feature = "async")]
struct WakerGuard(*mut ffi::lua_State);
struct WakerGuard(*mut ffi::lua_State, Option<Waker>);
#[cfg(feature = "async")]
impl WakerGuard {
pub fn new(state: *mut ffi::lua_State, waker: Waker) -> Result<WakerGuard> {
unsafe {
let _sg = StackGuard::new(state);
check_stack(state, 5)?;
check_stack(state, 3)?;
push_gc_userdata(state, waker)?;
let waker_key = &WAKER_REGISTRY_KEY as *const u8 as *const c_void;
ffi::safe::lua_rawsetp(state, ffi::LUA_REGISTRYINDEX, waker_key)?;
ffi::lua_rawgetp(state, ffi::LUA_REGISTRYINDEX, waker_key);
let waker_slot = get_gc_userdata::<Option<Waker>>(state, -1).as_mut();
let old = mlua_expect!(waker_slot, "Waker is destroyed").replace(waker);
Ok(WakerGuard(state))
Ok(WakerGuard(state, old))
}
}
}
@ -335,11 +337,12 @@ impl Drop for WakerGuard {
let state = self.0;
unsafe {
let _sg = StackGuard::new(state);
assert_stack(state, 1);
assert_stack(state, 3);
ffi::lua_pushnil(state);
let waker_key = &WAKER_REGISTRY_KEY as *const u8 as *const c_void;
ffi::lua_rawsetp(state, ffi::LUA_REGISTRYINDEX, waker_key); // TODO: make safe
ffi::lua_rawgetp(state, ffi::LUA_REGISTRYINDEX, waker_key);
let waker_slot = get_gc_userdata::<Option<Waker>>(state, -1).as_mut();
mem::swap(mlua_expect!(waker_slot, "Waker is destroyed"), &mut self.1);
}
}
}