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

View File

@ -12,13 +12,14 @@ use {
crate::{ crate::{
error::ExternalError, error::ExternalError,
lua::{ASYNC_POLL_PENDING, WAKER_REGISTRY_KEY}, lua::{ASYNC_POLL_PENDING, WAKER_REGISTRY_KEY},
util::push_gc_userdata, util::get_gc_userdata,
value::Value, value::Value,
}, },
futures_core::{future::Future, stream::Stream}, futures_core::{future::Future, stream::Stream},
std::{ std::{
cell::RefCell, cell::RefCell,
marker::PhantomData, marker::PhantomData,
mem,
os::raw::c_void, os::raw::c_void,
pin::Pin, pin::Pin,
task::{Context, Poll, Waker}, task::{Context, Poll, Waker},
@ -311,20 +312,21 @@ fn is_poll_pending(val: &MultiValue) -> bool {
} }
#[cfg(feature = "async")] #[cfg(feature = "async")]
struct WakerGuard(*mut ffi::lua_State); struct WakerGuard(*mut ffi::lua_State, Option<Waker>);
#[cfg(feature = "async")] #[cfg(feature = "async")]
impl WakerGuard { impl WakerGuard {
pub fn new(state: *mut ffi::lua_State, waker: Waker) -> Result<WakerGuard> { pub fn new(state: *mut ffi::lua_State, waker: Waker) -> Result<WakerGuard> {
unsafe { unsafe {
let _sg = StackGuard::new(state); 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; 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; let state = self.0;
unsafe { unsafe {
let _sg = StackGuard::new(state); 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; 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);
} }
} }
} }