Async: store pointer to Waker in extra data rather than in Lua ref thread

This commit is contained in:
Alex Orlenko 2023-03-03 15:19:54 +00:00
parent 5394faca16
commit 8c18fa1764
No known key found for this signature in database
GPG Key ID: 4C150C250863B96D
2 changed files with 19 additions and 45 deletions

View File

@ -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<Waker>` userdata on the ref thread
// Waker for polling futures
#[cfg(feature = "async")]
ref_waker_idx: c_int,
waker: NonNull<Waker>,
#[cfg(not(feature = "luau"))]
hook_callback: Option<HookCallback>,
@ -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::<Option<Waker>>(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<Waker> {
let extra = &*self.extra.get();
(*get_userdata::<Option<Waker>>(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<Waker>) -> Option<Waker> {
let extra = &*self.extra.get();
let waker_slot = &mut *get_userdata::<Option<Waker>>(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<Waker>) -> NonNull<Waker> {
mem::replace(&mut (*self.extra.get()).waker, waker)
}
pub(crate) unsafe fn make_userdata<T>(&self, data: UserDataCell<T>) -> Result<AnyUserData>

View File

@ -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<Waker>,
prev: NonNull<Waker>,
}
#[cfg(feature = "async")]
impl<'lua> WakerGuard<'lua> {
#[inline]
pub fn new(lua: &Lua, waker: Waker) -> Result<WakerGuard> {
pub fn new(lua: &'lua Lua, waker: &Waker) -> Result<WakerGuard<'lua>> {
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);
}
}
}