diff --git a/src/lua.rs b/src/lua.rs index 5cb06d2..78c7f2d 100644 --- a/src/lua.rs +++ b/src/lua.rs @@ -317,7 +317,12 @@ impl Lua { { init_gc_metatable_for::(state, None)?; init_gc_metatable_for::>>(state, None)?; - init_gc_metatable_for::(state, None)?; + init_gc_metatable_for::>(state, None)?; + + // Create empty Waker slot + push_gc_userdata::>(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::(state, -1).as_ref() { - waker = (*w).clone(); - } + let waker = match get_gc_userdata::>(state, -1).as_ref() { + Some(Some(waker)) => waker.clone(), + _ => noop_waker(), + }; ffi::lua_pop(state, 1); let mut ctx = Context::from_waker(&waker); diff --git a/src/thread.rs b/src/thread.rs index 238c040..96c932f 100644 --- a/src/thread.rs +++ b/src/thread.rs @@ -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); #[cfg(feature = "async")] impl WakerGuard { pub fn new(state: *mut ffi::lua_State, waker: Waker) -> Result { 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::>(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::>(state, -1).as_mut(); + mem::swap(mlua_expect!(waker_slot, "Waker is destroyed"), &mut self.1); } } }