Something went wrong. Try again.
Personal fork of prime agent. Mirror of: https://mnemosyne.sigilzero.dev/sigilzero/prime-agent
Something went wrong. Try again.
35 kB · 789 lines
at atom/orphans
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790--- a/src/thread.rs+++ b/src/thread.rs@@ -22,6 +22,8 @@ marker::PhantomData, pin::Pin, ptr::NonNull,+ sync::{Arc, Weak, atomic::{AtomicBool, Ordering}},+ task::Wake, task::{Context, Poll, Waker}, }, };@@ -88,6 +90,9 @@ thread: Thread, ret: PhantomData<fn() -> R>, recycle: bool,+ resume_nargs: Option<c_int>,+ wake_debt: Arc<AtomicBool>,+ wake_route: Weak<FutureWake>, } impl Thread {@@ -435,6 +440,9 @@ thread: self, ret: PhantomData, recycle: false,+ resume_nargs: None,+ wake_debt: Arc::new(AtomicBool::new(false)),+ wake_route: Weak::new(), }) } }@@ -514,6 +522,13 @@ #[cfg(feature = "async")] impl<R> AsyncThread<R> {+ fn clear_wake_route(&self) {+ if let Some(route) = self.wake_route.upgrade() {+ let previous = route.caller.lock().take();+ drop(previous);+ }+ }+ #[inline(always)] pub(crate) fn set_recyclable(&mut self, recyclable: bool) { self.recycle = recyclable;@@ -523,6 +538,7 @@ #[cfg(feature = "async")] impl<R> Drop for AsyncThread<R> { fn drop(&mut self) {+ self.clear_wake_route(); if self.recycle { if let Some(lua) = self.thread.0.lua.try_lock() { unsafe {@@ -551,33 +567,63 @@ type Item = Result<R>; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {- let lua = self.thread.0.lua.lock();- let nargs = match self.thread.status_inner(&lua) {+ let this = self.get_mut();+ let lua = this.thread.0.lua.lock();+ let nargs = match this.thread.status_inner(&lua) { ThreadStatusInner::New(nargs) | ThreadStatusInner::Yielded(nargs) => nargs,- _ => return Poll::Ready(None),+ _ => { this.clear_wake_route(); return Poll::Ready(None); }, }; + // A retained hook-suspended frame is not a list of resume arguments.+ let nargs = this.resume_nargs.take().unwrap_or(nargs); let state = lua.state();- let thread_state = self.thread.state();+ let thread_state = this.thread.state(); unsafe { let _sg = StackGuard::new(state);- let _thread_sg = StackGuard::with_top(thread_state, 0);- let _wg = WakerGuard::new(&lua, cx.waker());-- let (status, nresults) = (self.thread).resume_inner(&lua, nargs)?;+ let mut thread_sg = StackGuard::with_top(thread_state, 0);+ let route = match this.wake_route.upgrade() {+ Some(route) => route,+ None => {+ let route = Arc::new(FutureWake {+ caller: parking_lot::Mutex::new(Some(Arc::new(cx.waker().clone()))),+ debt: this.wake_debt.clone(),+ });+ this.wake_route = Arc::downgrade(&route);+ route+ }+ };+ let caller = Arc::new(cx.waker().clone());+ let previous = std::mem::replace(&mut *route.caller.lock(), Some(caller));+ drop(previous);+ let proxy = Waker::from(route);+ let _wg = WakerGuard::new(&lua, &proxy, this.wake_debt.clone());++ let (status, nresults) = (this.thread).resume_inner(&lua, nargs).map_err(|error| { this.clear_wake_route(); error })?; if status.is_yielded() {+ this.resume_nargs = Some(0);+ if nresults == 0 {+ thread_sg.keep(ffi::lua_gettop(thread_state));+ } if nresults == 1 && is_poll_pending(thread_state) {+ if this.wake_debt.load(Ordering::SeqCst) {+ cx.waker().wake_by_ref();+ } return Poll::Pending; } // Continue polling cx.waker().wake_by_ref(); } - check_stack(state, nresults + 1)?;+ if !status.is_yielded() {+ this.clear_wake_route();+ }+ check_stack(state, nresults + 1).map_err(|error| { this.clear_wake_route(); error })?; ffi::lua_xmove(thread_state, state, nresults); - Poll::Ready(Some(R::from_stack_multi(nresults, &lua)))+ let result = R::from_stack_multi(nresults, &lua);+ if result.is_err() { this.clear_wake_route(); }+ Poll::Ready(Some(result)) } } }@@ -587,30 +633,61 @@ type Output = Result<R>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {- let lua = self.thread.0.lua.lock();- let nargs = match self.thread.status_inner(&lua) {+ let this = self.get_mut();+ let lua = this.thread.0.lua.lock();+ let nargs = match this.thread.status_inner(&lua) { ThreadStatusInner::New(nargs) | ThreadStatusInner::Yielded(nargs) => nargs,- _ => return Poll::Ready(Err(Error::CoroutineUnresumable)),+ _ => { this.clear_wake_route(); return Poll::Ready(Err(Error::CoroutineUnresumable)); }, }; + // A retained hook-suspended frame is not a list of resume arguments.+ let nargs = this.resume_nargs.take().unwrap_or(nargs); let state = lua.state();- let thread_state = self.thread.state();+ let thread_state = this.thread.state(); unsafe { let _sg = StackGuard::new(state);- let _thread_sg = StackGuard::with_top(thread_state, 0);- let _wg = WakerGuard::new(&lua, cx.waker());-- let (status, nresults) = self.thread.resume_inner(&lua, nargs)?;+ let mut thread_sg = StackGuard::with_top(thread_state, 0);+ let route = match this.wake_route.upgrade() {+ Some(route) => route,+ None => {+ let route = Arc::new(FutureWake {+ caller: parking_lot::Mutex::new(Some(Arc::new(cx.waker().clone()))),+ debt: this.wake_debt.clone(),+ });+ this.wake_route = Arc::downgrade(&route);+ route+ }+ };+ let caller = Arc::new(cx.waker().clone());+ let previous = std::mem::replace(&mut *route.caller.lock(), Some(caller));+ drop(previous);+ let proxy = Waker::from(route);+ let _wg = WakerGuard::new(&lua, &proxy, this.wake_debt.clone());++ let (status, nresults) = this.thread.resume_inner(&lua, nargs).map_err(|error| { this.clear_wake_route(); error })?; if status.is_yielded() {- if !(nresults == 1 && is_poll_pending(thread_state)) {+ this.resume_nargs = Some(0);+ if nresults == 0 {+ // A count-hook suspension exposes the current Lua frame, not+ // returned yield values. Keep those locals until resume.+ thread_sg.keep(ffi::lua_gettop(thread_state));+ }+ if nresults == 1 && is_poll_pending(thread_state) {+ if this.wake_debt.load(Ordering::SeqCst) {+ cx.waker().wake_by_ref();+ }+ } else { // Ignore values returned via yield() cx.waker().wake_by_ref(); } return Poll::Pending; } - check_stack(state, nresults + 1)?;+ if !status.is_yielded() {+ this.clear_wake_route();+ }+ check_stack(state, nresults + 1).map_err(|error| { this.clear_wake_route(); error })?; ffi::lua_xmove(thread_state, state, nresults); Poll::Ready(R::from_stack_multi(nresults, &lua))@@ -625,20 +702,44 @@ } #[cfg(feature = "async")]+struct FutureWake {+ caller: parking_lot::Mutex<Option<Arc<Waker>>>,+ debt: Arc<AtomicBool>,+}++#[cfg(feature = "async")]+impl Wake for FutureWake {+ fn wake(self: Arc<Self>) {+ self.wake_by_ref();+ }++ fn wake_by_ref(self: &Arc<Self>) {+ self.debt.store(true, Ordering::SeqCst);+ let caller = self.caller.lock().clone();+ if let Some(caller) = caller {+ caller.wake_by_ref();+ }+ }+}++#[cfg(feature = "async")] struct WakerGuard<'lua, 'a> { lua: &'lua RawLua, prev: NonNull<Waker>,+ previous_debt: Option<Arc<AtomicBool>>, _phantom: PhantomData<&'a ()>, } #[cfg(feature = "async")] impl<'lua, 'a> WakerGuard<'lua, 'a> { #[inline]- pub fn new(lua: &'lua RawLua, waker: &'a Waker) -> Result<WakerGuard<'lua, 'a>> {+ pub fn new(lua: &'lua RawLua, waker: &'a Waker, debt: Arc<AtomicBool>) -> Result<WakerGuard<'lua, 'a>> { let prev = lua.set_waker(NonNull::from(waker));+ let previous_debt = lua.set_wake_debt(Some(debt)); Ok(WakerGuard { lua, prev,+ previous_debt, _phantom: PhantomData, }) }@@ -648,6 +749,7 @@ impl Drop for WakerGuard<'_, '_> { fn drop(&mut self) { self.lua.set_waker(self.prev);+ self.lua.set_wake_debt(self.previous_debt.take()); } } @@ -664,3 +766,269 @@ #[cfg(all(feature = "async", feature = "send"))] static_assertions::assert_impl_all!(AsyncThread<()>: Send, Sync); }++#[cfg(all(test, feature = "async", any(feature = "lua54", feature = "lua55")))]+mod wake_debt_tests {+ use super::*;+ use crate::Lua;+ use std::sync::atomic::AtomicUsize;++ struct CountWake(AtomicUsize);+ impl Wake for CountWake {+ fn wake(self: Arc<Self>) { self.wake_by_ref(); }+ fn wake_by_ref(self: &Arc<Self>) { self.0.fetch_add(1, Ordering::SeqCst); }+ }++ #[test]+ fn route_migration_forwards_to_latest_parent_and_clears_on_drop() {+ let first = Arc::new(CountWake(AtomicUsize::new(0)));+ let second = Arc::new(CountWake(AtomicUsize::new(0)));+ let debt = Arc::new(AtomicBool::new(false));+ let route = Arc::new(FutureWake { caller: parking_lot::Mutex::new(Some(Arc::new(Waker::from(first.clone())))), debt: debt.clone() });+ let old_proxy = Waker::from(route.clone());+ let parent = Arc::new(Waker::from(second.clone()));+ let previous = std::mem::replace(&mut *route.caller.lock(), Some(parent));+ drop(previous);+ old_proxy.wake_by_ref();+ assert!(debt.load(Ordering::SeqCst));+ assert_eq!(first.0.load(Ordering::SeqCst), 0);+ assert_eq!(second.0.load(Ordering::SeqCst), 1);+ let previous = route.caller.lock().take();+ drop(previous);+ old_proxy.wake_by_ref();+ assert_eq!(second.0.load(Ordering::SeqCst), 1);+ }++ #[test]+ fn real_wake_debt_before_during_after_ack_is_not_consumed_by_forwarding() {+ let parent = Arc::new(CountWake(AtomicUsize::new(0)));+ let debt = Arc::new(AtomicBool::new(false));+ let route = Arc::new(FutureWake { caller: parking_lot::Mutex::new(Some(Arc::new(Waker::from(parent.clone())))), debt: debt.clone() });+ let wake = Waker::from(route);+ wake.wake_by_ref();+ assert!(debt.load(Ordering::SeqCst));+ debt.store(false, Ordering::SeqCst); // Actual Rust poll acknowledges prior wake.+ assert!(!debt.load(Ordering::SeqCst));+ let concurrent = wake.clone();+ std::thread::spawn(move || concurrent.wake()).join().unwrap();+ assert!(debt.load(Ordering::SeqCst));+ assert_eq!(parent.0.load(Ordering::SeqCst), 2);+ assert!(debt.load(Ordering::SeqCst)); // A conditional handoff does not acknowledge.+ debt.store(false, Ordering::SeqCst);+ wake.wake_by_ref();+ assert!(debt.load(Ordering::SeqCst));+ assert_eq!(parent.0.load(Ordering::SeqCst), 3);+ }++ #[test]+ fn nested_waker_guards_restore_both_bindings() -> crate::Result<()> {+ let lua = Lua::new();+ let raw = lua.lock();+ let initial = raw.waker() as *const Waker;+ let first = Waker::from(Arc::new(CountWake(AtomicUsize::new(0))));+ let second = Waker::from(Arc::new(CountWake(AtomicUsize::new(0))));+ let first_debt = Arc::new(AtomicBool::new(false));+ let second_debt = Arc::new(AtomicBool::new(false));+ {+ let _outer = WakerGuard::new(&raw, &first, first_debt.clone())?;+ assert!(std::ptr::eq(raw.waker(), &first));+ {+ let _inner = WakerGuard::new(&raw, &second, second_debt.clone())?;+ assert!(std::ptr::eq(raw.waker(), &second));+ let bound = raw.set_wake_debt(None);+ assert!(Arc::ptr_eq(bound.as_ref().unwrap(), &second_debt));+ raw.set_wake_debt(bound);+ }+ assert!(std::ptr::eq(raw.waker(), &first));+ let bound = raw.set_wake_debt(None);+ assert!(Arc::ptr_eq(bound.as_ref().unwrap(), &first_debt));+ raw.set_wake_debt(bound);+ }+ assert_eq!(raw.waker() as *const Waker, initial);+ assert!(raw.set_wake_debt(None).is_none());+ Ok(())+ }++ #[test]+ fn never_ready_no_wake_future_does_not_schedule_marker_spin() -> crate::Result<()> {+ let lua = Lua::new();+ lua.globals().set("wait_forever", lua.create_async_function(|_, ()| async {+ std::future::pending::<()>().await;+ Ok(())+ })?)?;+ let function = lua.load("wait_forever()").into_function()?;+ let mut execution = lua.create_thread(function)?.into_async::<()>(())?;+ let parent = Arc::new(CountWake(AtomicUsize::new(0)));+ let waker = Waker::from(parent.clone());+ let mut cx = Context::from_waker(&waker);+ assert!(Pin::new(&mut execution).poll(&mut cx).is_pending());+ assert_eq!(parent.0.load(Ordering::SeqCst), 0);+ assert!(!execution.wake_debt.load(Ordering::SeqCst));+ drop(execution);+ assert_eq!(parent.0.load(Ordering::SeqCst), 0);+ Ok(())+ }++ #[test]+ fn async_thread_migration_routes_registered_proxy_without_rust_repoll() -> crate::Result<()> {+ struct PendingProbe {+ polls: Arc<AtomicUsize>,+ registered: Arc<parking_lot::Mutex<Option<Waker>>>,+ }+ impl Future for PendingProbe {+ type Output = crate::Result<i64>;+ fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {+ self.polls.fetch_add(1, Ordering::SeqCst);+ *self.registered.lock() = Some(cx.waker().clone());+ Poll::Pending+ }+ }+ let lua = Lua::new();+ let polls = Arc::new(AtomicUsize::new(0));+ let registered = Arc::new(parking_lot::Mutex::new(None));+ let callback_polls = polls.clone();+ let callback_registered = registered.clone();+ lua.globals().set("await_tick", lua.create_async_function(move |_, ()| PendingProbe {+ polls: callback_polls.clone(), registered: callback_registered.clone(),+ })?)?;+ let instructions = Arc::new(AtomicUsize::new(0));+ let counted = instructions.clone();+ lua.set_global_hook(HookTriggers::new().every_nth_instruction(1), move |_, _| {+ let instruction = counted.fetch_add(1, Ordering::SeqCst) + 1;+ assert!(instruction < 100_000);+ Ok(if instruction == 49 { crate::VmState::Yield } else { crate::VmState::Continue })+ })?;+ let function = lua.load(r#"+ local run = {"alpha", "beta", "gamma", "delta"}+ local scalar = 271828+ local function add(term)+ assert(type(term) == "string", "concat result lost across hook yield")+ return term+ end+ local joined = add(table.concat(run))+ assert(joined == "alphabetagammadelta")+ assert(scalar == 271828 and run[3] == "gamma")+ assert(await_tick() == 17)+ assert(table.concat(run) == joined)+ return joined .. ":" .. scalar+ "#).into_function()?;+ let mut execution = lua.create_thread(function)?.into_async::<String>(())?;+ let first = Arc::new(CountWake(AtomicUsize::new(0)));+ let second = Arc::new(CountWake(AtomicUsize::new(0)));+ let first_waker = Waker::from(first.clone());+ let second_waker = Waker::from(second.clone());+ assert!(Pin::new(&mut execution).poll(&mut Context::from_waker(&first_waker)).is_pending());+ assert_eq!(polls.load(Ordering::SeqCst), 1, "hook must interrupt after Rust registration");+ let old_proxy = registered.lock().as_ref().unwrap().clone();+ assert!(Pin::new(&mut execution).poll(&mut Context::from_waker(&second_waker)).is_pending());+ assert_eq!(polls.load(Ordering::SeqCst), 1, "migration poll must only advance hook handoff");+ let before_first = first.0.load(Ordering::SeqCst);+ let before_second = second.0.load(Ordering::SeqCst);+ old_proxy.wake_by_ref();+ assert_eq!(first.0.load(Ordering::SeqCst), before_first);+ assert_eq!(second.0.load(Ordering::SeqCst), before_second + 1);+ // Drain only actual notifications, never manually spin a Pending future.+ let mut consumed = before_second;+ let mut drains = 0;+ while second.0.load(Ordering::SeqCst) > consumed {+ consumed = second.0.load(Ordering::SeqCst);+ drains += 1;+ assert!(drains < 16, "never-ready future scheduled indefinitely");+ assert!(Pin::new(&mut execution).poll(&mut Context::from_waker(&second_waker)).is_pending());+ }+ assert_eq!(polls.load(Ordering::SeqCst), 2);+ assert!(!execution.wake_debt.load(Ordering::SeqCst));+ drop(execution);+ let after_drop = second.0.load(Ordering::SeqCst);+ old_proxy.wake_by_ref();+ assert_eq!(second.0.load(Ordering::SeqCst), after_drop);+ Ok(())+ }++ #[test]+ fn concurrent_route_registration_and_real_wake_forward_once() {+ for _ in 0..64 {+ let first = Arc::new(CountWake(AtomicUsize::new(0)));+ let second = Arc::new(CountWake(AtomicUsize::new(0)));+ let debt = Arc::new(AtomicBool::new(false));+ let route = Arc::new(FutureWake { caller: parking_lot::Mutex::new(Some(Arc::new(Waker::from(first.clone())))), debt: debt.clone() });+ let gate = Arc::new(std::sync::Barrier::new(2));+ let notifying_route = route.clone();+ let notifying_gate = gate.clone();+ let notifier = std::thread::spawn(move || {+ notifying_gate.wait();+ Waker::from(notifying_route).wake();+ });+ let latest = Arc::new(Waker::from(second.clone()));+ gate.wait();+ let previous = std::mem::replace(&mut *route.caller.lock(), Some(latest));+ drop(previous);+ notifier.join().unwrap();+ assert!(debt.load(Ordering::SeqCst));+ assert_eq!(first.0.load(Ordering::SeqCst) + second.0.load(Ordering::SeqCst), 1);+ // Once registration finished, every old proxy must reach latest parent.+ let before = second.0.load(Ordering::SeqCst);+ Waker::from(route).wake();+ assert_eq!(second.0.load(Ordering::SeqCst), before + 1);+ }+ }++ #[test]+ fn actual_rust_poll_ack_preserves_concurrent_notification() -> crate::Result<()> {+ struct RacingPending {+ polls: Arc<AtomicUsize>,+ registered: Arc<parking_lot::Mutex<Option<Waker>>>,+ }+ impl Future for RacingPending {+ type Output = crate::Result<()>;+ fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {+ let poll = self.polls.fetch_add(1, Ordering::SeqCst) + 1;+ *self.registered.lock() = Some(cx.waker().clone());+ if poll == 1 {+ // raw.rs has acknowledged debt before entering this actual poll.+ let gate = Arc::new(std::sync::Barrier::new(2));+ let notify_gate = gate.clone();+ let notify = cx.waker().clone();+ let notifier = std::thread::spawn(move || {+ notify_gate.wait();+ notify.wake();+ });+ gate.wait();+ notifier.join().unwrap();+ }+ Poll::Pending+ }+ }+ let lua = Lua::new();+ let polls = Arc::new(AtomicUsize::new(0));+ let registered = Arc::new(parking_lot::Mutex::new(None));+ let callback_polls = polls.clone();+ let callback_registered = registered.clone();+ lua.globals().set("racing_pending", lua.create_async_function(move |_, ()| RacingPending {+ polls: callback_polls.clone(), registered: callback_registered.clone(),+ })?)?;+ let function = lua.load("racing_pending()").into_function()?;+ let mut execution = lua.create_thread(function)?.into_async::<()>(())?;+ let parent = Arc::new(CountWake(AtomicUsize::new(0)));+ let waker = Waker::from(parent.clone());+ let mut cx = Context::from_waker(&waker);+ assert!(Pin::new(&mut execution).poll(&mut cx).is_pending());+ assert_eq!(polls.load(Ordering::SeqCst), 1);+ assert!(execution.wake_debt.load(Ordering::SeqCst), "wake during actual poll must survive raw ack");+ assert!(parent.0.load(Ordering::SeqCst) >= 1);+ // Consume the real notification: actual Rust poll now acknowledges it.+ assert!(Pin::new(&mut execution).poll(&mut cx).is_pending());+ assert_eq!(polls.load(Ordering::SeqCst), 2);+ assert!(!execution.wake_debt.load(Ordering::SeqCst));+ let count = parent.0.load(Ordering::SeqCst);+ // An after-poll wake is retained until the next actual poll.+ let registered = registered.lock().as_ref().unwrap().clone();+ registered.wake_by_ref();+ assert_eq!(parent.0.load(Ordering::SeqCst), count + 1);+ assert!(execution.wake_debt.load(Ordering::SeqCst));+ assert!(Pin::new(&mut execution).poll(&mut cx).is_pending());+ assert_eq!(polls.load(Ordering::SeqCst), 3);+ assert!(!execution.wake_debt.load(Ordering::SeqCst));+ Ok(())+ }+}--- a/src/state/raw.rs+++ b/src/state/raw.rs@@ -1356,7 +1356,12 @@ let fut = &mut (*future).data; let mut ctx = Context::from_waker(rawlua.waker());- match fut.as_mut().map(|fut| fut.as_mut().poll(&mut ctx)) {+ match fut.as_mut().map(|fut| {+ if let Some(debt) = (*extra).wake_debt.as_ref() {+ debt.store(false, std::sync::atomic::Ordering::SeqCst);+ }+ fut.as_mut().poll(&mut ctx)+ }) { Some(Poll::Pending) => { let fut_nvals = ffi::lua_gettop(state) - 1; // Exclude the future itself if fut_nvals >= 3 && ffi::lua_tolightuserdata(state, -3) == Lua::poll_yield().0 {@@ -1483,6 +1488,14 @@ } #[cfg(feature = "async")]+ pub(crate) fn set_wake_debt(+ &self,+ debt: Option<Arc<std::sync::atomic::AtomicBool>>,+ ) -> Option<Arc<std::sync::atomic::AtomicBool>> {+ unsafe { mem::replace(&mut (*self.extra.get()).wake_debt, debt) }+ }++ #[cfg(feature = "async")] #[inline] pub(crate) fn set_waker(&self, waker: NonNull<Waker>) -> NonNull<Waker> { unsafe { mem::replace(&mut (*self.extra.get()).waker, waker) }--- a/src/state/extra.rs+++ b/src/state/extra.rs@@ -72,6 +72,8 @@ // Waker for polling futures #[cfg(feature = "async")] pub(super) waker: NonNull<Waker>,+ #[cfg(feature = "async")]+ pub(super) wake_debt: Option<Arc<std::sync::atomic::AtomicBool>>, #[cfg(not(feature = "luau"))] pub(super) hook_callback: Option<crate::types::HookCallback>,@@ -178,6 +180,8 @@ wrapped_failure_mt_ptr, #[cfg(feature = "async")] waker: NonNull::from(noop_waker_ref()),+ #[cfg(feature = "async")]+ wake_debt: None, #[cfg(not(feature = "luau"))] hook_callback: None, #[cfg(not(feature = "luau"))]--- a/tests/hooks.rs+++ b/tests/hooks.rs@@ -325,3 +325,212 @@ Ok(()) }++// Sweep each instruction boundary rather than relying on scheduler timing.+#[cfg(all(feature = "async", any(feature = "lua54", feature = "lua55")))]+#[tokio::test]+async fn zero_result_hook_yields_preserve_concat_arguments_and_frame_locals() -> Result<()> {+ use std::sync::atomic::AtomicUsize;+ use futures_util::StreamExt;+ const SCRIPT: &str = r#"+ local run = {"alpha", "beta", "gamma", "delta"}+ local scalar = 271828+ local function add(term)+ assert(type(term) == "string", "concat result lost across hook yield")+ return term+ end+ local joined = add(table.concat(run))+ assert(joined == "alphabetagammadelta")+ assert(scalar == 271828 and run[3] == "gamma")+ assert(await_tick() == 17)+ assert(table.concat(run) == joined)+ return joined .. ":" .. scalar+ "#;+ for stream in [false, true] {+ for boundary in 1..=256 {+ let started = std::time::Instant::now();+ eprintln!("hook-arity start stream={stream} boundary={boundary}");+ tokio::time::timeout(std::time::Duration::from_secs(3), async {+ let lua = Lua::new();+ lua.globals().set("await_tick", lua.create_async_function(|_, ()| async {+ eprintln!("hook-arity async-C enter");+ tokio::task::yield_now().await;+ eprintln!("hook-arity async-C return");+ Ok(17)+ })?)?;+ let instruction = Arc::new(AtomicUsize::new(0));+ let observed = instruction.clone();+ lua.set_global_hook(HookTriggers::new().every_nth_instruction(1), move |_, _| {+ let current = observed.fetch_add(1, Ordering::Relaxed) + 1;+ if current > 100_000 {+ return Err(Error::RuntimeError(format!("instruction non-progress stream={stream} boundary={boundary} count={current}")));+ }+ Ok(if current == boundary { VmState::Yield } else { VmState::Continue })+ })?;+ let function = lua.load(SCRIPT).into_function()?;+ let thread = lua.create_thread(function)?;+ if stream {+ let mut execution = thread.into_async::<mlua::MultiValue>(())?;+ let mut final_text = None;+ while let Some(row) = execution.next().await {+ let row = row?;+ if let Some(Value::String(text)) = row.front() {+ final_text = Some(text.to_str()?.to_string());+ } else {+ assert!(row.is_empty(), "unexpected hook/async yield payload");+ }+ }+ assert_eq!(final_text.as_deref(), Some("alphabetagammadelta:271828"), "stream boundary {boundary}");+ } else {+ let result: String = thread.into_async(())?.await?;+ assert_eq!(result, "alphabetagammadelta:271828", "future boundary {boundary}");+ }+ Ok::<(), Error>(())+ }).await.unwrap_or_else(|_| panic!("hook-arity deadline stream={stream} boundary={boundary}"))?;+ assert!(started.elapsed() < std::time::Duration::from_secs(1), "timer-rescued hook boundary stream={stream} boundary={boundary}");+ eprintln!("hook-arity complete stream={stream} boundary={boundary}");+ }+ }+ Ok(())+}++#[cfg(all(feature = "async", any(feature = "lua54", feature = "lua55")))]+#[tokio::test]+async fn async_stream_keeps_ordinary_yield_payloads_and_initial_arguments() -> Result<()> {+ use futures_util::StreamExt;+ let lua = Lua::new();+ let function = lua.load(r#"+ local initial = ...+ assert(initial == 42)+ local resumed = coroutine.yield("first", initial)+ assert(resumed == nil)+ coroutine.yield()+ return "last", initial + 1+ "#).into_function()?;+ let mut stream = lua.create_thread(function)?.into_async::<mlua::MultiValue>(42)?;+ let first = stream.next().await.unwrap()?;+ assert_eq!(first.len(), 2);+ assert!(matches!(first.get(1), Some(Value::Integer(42))));+ assert!(stream.next().await.unwrap()?.is_empty());+ let last = stream.next().await.unwrap()?;+ assert_eq!(last.len(), 2);+ assert!(matches!(last.get(1), Some(Value::Integer(43))));+ assert!(stream.next().await.is_none());+ Ok(())+}++#[cfg(all(feature = "async", any(feature = "lua54", feature = "lua55")))]+#[tokio::test]+async fn async_c_yield_preserves_initial_arguments_and_return_values() -> Result<()> {+ let lua = Lua::new();+ lua.globals().set("await_echo", lua.create_async_function(|_, value: i64| async move {+ tokio::task::yield_now().await;+ tokio::task::yield_now().await;+ Ok((value, "echo"))+ })?)?;+ let function = lua.load(r#"+ local initial = ...+ assert(initial == 42)+ local number, text = await_echo(initial)+ assert(number == initial and text == "echo")+ return number + 1+ "#).into_function()?;+ let answer: i64 = lua.create_thread(function)?.into_async(42)?.await?;+ assert_eq!(answer, 43);+ Ok(())+}++#[cfg(all(feature = "async", any(feature = "lua54", feature = "lua55")))]+#[tokio::test]+async fn counted_hook_async_handoff_timeline() -> Result<()> {+ use std::future::Future;+ use std::pin::Pin;+ use std::sync::atomic::AtomicUsize;+ use std::task::{Context, Poll, Waker};+ use futures_util::task::{ArcWake, waker_ref};+ struct TraceWake { target: Waker, generation: Arc<AtomicUsize> }+ impl ArcWake for TraceWake {+ fn wake_by_ref(this: &Arc<Self>) {+ let generation = this.generation.fetch_add(1, Ordering::SeqCst) + 1;+ eprintln!("timeline WAKE generation={generation}");+ this.target.wake_by_ref();+ }+ }+ struct OnePending { polled: bool, count: Arc<AtomicUsize> }+ impl Future for OnePending {+ type Output = Result<i64>;+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<i64>> {+ let count = self.count.fetch_add(1, Ordering::SeqCst) + 1;+ eprintln!("timeline C_POLL count={count}");+ if self.polled { return Poll::Ready(Ok(17)); }+ self.polled = true;+ cx.waker().wake_by_ref();+ Poll::Pending+ }+ }+ for boundary in [49usize, 50, 51] {+ let started = std::time::Instant::now();+ eprintln!("timeline START boundary={boundary}");+ let lua = Lua::new();+ let c_polls = Arc::new(AtomicUsize::new(0));+ let callback_polls = c_polls.clone();+ lua.globals().set("await_tick", lua.create_async_function(move |_, ()| {+ OnePending { polled: false, count: callback_polls.clone() }+ })?)?;+ let instructions = Arc::new(AtomicUsize::new(0));+ let counted = instructions.clone();+ lua.set_global_hook(HookTriggers::new().every_nth_instruction(1), move |_, _| {+ let instruction = counted.fetch_add(1, Ordering::SeqCst) + 1;+ if instruction > 100_000 { return Err(Error::RuntimeError("timeline instruction cap".into())); }+ if instruction == boundary {+ eprintln!("timeline HOOK_YIELD boundary={boundary}");+ return Ok(VmState::Yield);+ }+ Ok(VmState::Continue)+ })?;+ let function = lua.load(r#"+ local run = {"alpha", "beta", "gamma", "delta"}+ local scalar = 271828+ local function add(term)+ assert(type(term) == "string", "concat result lost across hook yield")+ return term+ end+ local joined = add(table.concat(run))+ assert(joined == "alphabetagammadelta")+ assert(scalar == 271828 and run[3] == "gamma")+ assert(await_tick() == 17)+ assert(table.concat(run) == joined)+ return joined .. ":" .. scalar+ "#).into_function()?;+ let execution = lua.create_thread(function)?.into_async::<String>(())?;+ let deadline = tokio::time::timeout(std::time::Duration::from_secs(3), execution);+ tokio::pin!(deadline);+ let wakes = Arc::new(AtomicUsize::new(0));+ let mut polls = 0usize;+ let result = futures_util::future::poll_fn(|cx| {+ polls += 1;+ let proxy = Arc::new(TraceWake { target: cx.waker().clone(), generation: wakes.clone() });+ let wake = waker_ref(&proxy);+ let mut traced = Context::from_waker(&wake);+ eprintln!("timeline OUTER_ENTER boundary={boundary} poll={polls} wakes={} c_polls={}", wakes.load(Ordering::SeqCst), c_polls.load(Ordering::SeqCst));+ let result = deadline.as_mut().poll(&mut traced);+ eprintln!("timeline OUTER_RETURN boundary={boundary} poll={polls} ready={} wakes={} c_polls={}", result.is_ready(), wakes.load(Ordering::SeqCst), c_polls.load(Ordering::SeqCst));+ result+ }).await;+ let answer = result.unwrap_or_else(|_| panic!("timeline DEADLINE boundary={boundary} polls={polls} c_polls={}", c_polls.load(Ordering::SeqCst)))?;+ assert_eq!(answer, "alphabetagammadelta:271828");+ assert_eq!(c_polls.load(Ordering::SeqCst), 2);+ assert!(started.elapsed() < std::time::Duration::from_secs(1), "timer-rescued counted boundary={boundary}");+ eprintln!("timeline COMPLETE boundary={boundary}");+ }+ Ok(())+}++#[cfg(all(feature = "async", any(feature = "lua54", feature = "lua55")))]+#[tokio::test]+async fn counted_hook_timer_control() {+ eprintln!("timeline TIMER_CONTROL enter");+ let deadline = tokio::time::timeout(std::time::Duration::from_millis(100), std::future::pending::<()>());+ assert!(deadline.await.is_err());+ eprintln!("timeline TIMER_CONTROL complete");+}