From 8c18fa1764735ea8c1f8efb47939c4db871ec5fc Mon Sep 17 00:00:00 2001 From: Alex Orlenko Date: Fri, 3 Mar 2023 15:19:54 +0000 Subject: [PATCH] Async: store pointer to Waker in extra data rather than in Lua ref thread --- src/lua.rs | 51 ++++++++++++--------------------------------------- src/thread.rs | 13 +++++++------ 2 files changed, 19 insertions(+), 45 deletions(-) diff --git a/src/lua.rs b/src/lua.rs index e040ec5..bf9ab13 100644 --- a/src/lua.rs +++ b/src/lua.rs @@ -54,12 +54,12 @@ use crate::{chunk::Compiler, types::VmState}; #[cfg(feature = "async")] use { crate::types::{AsyncCallback, AsyncCallbackUpvalue, AsyncPollUpvalue}, - futures_core::{ + futures_task::noop_waker_ref, + futures_util::future::{self, TryFutureExt}, + std::{ future::Future, task::{Context, Poll, Waker}, }, - futures_task::noop_waker, - futures_util::future::{self, TryFutureExt}, }; #[cfg(feature = "serialize")] @@ -113,9 +113,9 @@ pub(crate) struct ExtraData { // Address of `WrappedFailure` metatable wrapped_failure_mt_ptr: *const c_void, - // Index of `Option` userdata on the ref thread + // Waker for polling futures #[cfg(feature = "async")] - ref_waker_idx: c_int, + waker: NonNull, #[cfg(not(feature = "luau"))] hook_callback: Option, @@ -246,13 +246,6 @@ impl Drop for LuaInner { ffi::lua_replace(extra.ref_thread, index); extra.ref_free.push(index); } - #[cfg(feature = "async")] - { - // Destroy Waker slot - ffi::lua_pushnil(extra.ref_thread); - ffi::lua_replace(extra.ref_thread, extra.ref_waker_idx); - extra.ref_free.push(extra.ref_waker_idx); - } #[cfg(feature = "luau")] { (*ffi::lua_callbacks(self.state())).userdata = ptr::null_mut(); @@ -555,17 +548,6 @@ impl Lua { ptr }; - // Create empty Waker slot on the ref thread - #[cfg(feature = "async")] - let ref_waker_idx = { - mlua_expect!( - push_gc_userdata::>(ref_thread, None, true), - "Error while creating Waker slot" - ); - ffi::lua_gettop(ref_thread) - }; - let ref_stack_top = ffi::lua_gettop(ref_thread); - // Create ExtraData let extra = Arc::new(UnsafeCell::new(ExtraData { inner: None, @@ -579,7 +561,7 @@ impl Lua { ref_thread, // We need 1 extra stack space to move values in and out of the ref stack. ref_stack_size: ffi::LUA_MINSTACK - 1, - ref_stack_top, + ref_stack_top: ffi::lua_gettop(ref_thread), ref_free: Vec::new(), wrapped_failure_pool: Vec::with_capacity(WRAPPED_FAILURE_POOL_SIZE), multivalue_pool: Vec::with_capacity(MULTIVALUE_POOL_SIZE), @@ -587,7 +569,7 @@ impl Lua { thread_pool: Vec::new(), wrapped_failure_mt_ptr, #[cfg(feature = "async")] - ref_waker_idx, + waker: NonNull::from(noop_waker_ref()), #[cfg(not(feature = "luau"))] hook_callback: None, #[cfg(feature = "lua54")] @@ -2818,11 +2800,8 @@ impl Lua { let lua: &Lua = mem::transmute((*extra).inner.as_ref().unwrap()); let _guard = StateGuard::new(&lua.0, state); - // Try to get an outer poll waker - let waker = lua.waker().unwrap_or_else(noop_waker); - let mut ctx = Context::from_waker(&waker); - let fut = &mut (*upvalue).data; + let mut ctx = Context::from_waker(lua.waker()); match fut.as_mut().poll(&mut ctx) { Poll::Pending => { check_stack(state, 1)?; @@ -2907,20 +2886,14 @@ impl Lua { #[cfg(feature = "async")] #[inline] - pub(crate) unsafe fn waker(&self) -> Option { - let extra = &*self.extra.get(); - (*get_userdata::>(extra.ref_thread, extra.ref_waker_idx)).clone() + pub(crate) unsafe fn waker(&self) -> &Waker { + (*self.extra.get()).waker.as_ref() } #[cfg(feature = "async")] #[inline] - pub(crate) unsafe fn set_waker(&self, waker: Option) -> Option { - let extra = &*self.extra.get(); - let waker_slot = &mut *get_userdata::>(extra.ref_thread, extra.ref_waker_idx); - match waker { - Some(waker) => waker_slot.replace(waker), - None => waker_slot.take(), - } + pub(crate) unsafe fn set_waker(&self, waker: NonNull) -> NonNull { + mem::replace(&mut (*self.extra.get()).waker, waker) } pub(crate) unsafe fn make_userdata(&self, data: UserDataCell) -> Result diff --git a/src/thread.rs b/src/thread.rs index f708095..0894065 100644 --- a/src/thread.rs +++ b/src/thread.rs @@ -24,6 +24,7 @@ use { std::{ marker::PhantomData, pin::Pin, + ptr::NonNull, task::{Context, Poll, Waker}, }, }; @@ -390,7 +391,7 @@ where _ => return Poll::Ready(None), }; - let _wg = WakerGuard::new(lua, cx.waker().clone()); + let _wg = WakerGuard::new(lua, cx.waker()); // This is safe as we are not moving the whole struct let this = unsafe { self.get_unchecked_mut() }; @@ -424,7 +425,7 @@ where _ => return Poll::Ready(Err(Error::CoroutineInactive)), }; - let _wg = WakerGuard::new(lua, cx.waker().clone()); + let _wg = WakerGuard::new(lua, cx.waker()); // This is safe as we are not moving the whole struct let this = unsafe { self.get_unchecked_mut() }; @@ -462,15 +463,15 @@ fn is_poll_pending(val: &MultiValue) -> bool { #[cfg(feature = "async")] struct WakerGuard<'lua> { lua: &'lua Lua, - prev: Option, + prev: NonNull, } #[cfg(feature = "async")] impl<'lua> WakerGuard<'lua> { #[inline] - pub fn new(lua: &Lua, waker: Waker) -> Result { + pub fn new(lua: &'lua Lua, waker: &Waker) -> Result> { unsafe { - let prev = lua.set_waker(Some(waker)); + let prev = lua.set_waker(NonNull::from(waker)); Ok(WakerGuard { lua, prev }) } } @@ -480,7 +481,7 @@ impl<'lua> WakerGuard<'lua> { impl<'lua> Drop for WakerGuard<'lua> { fn drop(&mut self) { unsafe { - self.lua.set_waker(self.prev.take()); + self.lua.set_waker(self.prev); } } }