From 4771721718effbaad028c8fc2edfdc0e663d16e3 Mon Sep 17 00:00:00 2001 From: Yehor Smoliakov Date: Sun, 11 Oct 2026 02:25:49 -0700 Subject: [PATCH] fix: retain public Rust timer signatures after Python cancellation fix --- src/bindings/loop_api/methods.rs | 4 +- src/engine/callbacks.rs | 75 ++++++++++++++++++--- src/engine/loop_core.rs | 110 +++++++++++++++++++++++++++++-- 3 files changed, 173 insertions(+), 16 deletions(-) diff --git a/src/bindings/loop_api/methods.rs b/src/bindings/loop_api/methods.rs index 1e8b7fa..a844623 100644 --- a/src/bindings/loop_api/methods.rs +++ b/src/bindings/loop_api/methods.rs @@ -102,7 +102,7 @@ impl PyLoop { return Err(PyValueError::new_err("delay cannot be NaN")); } let delay = delay.clamp(0.0, MAX_TIMER_DELAY_SECS); - let (ready, when) = self.core.schedule_timer( + let (ready, when) = self.core.schedule_cancellable_timer( py, Duration::from_secs_f64(delay), callback, @@ -110,7 +110,7 @@ impl PyLoop { context, )?; - Py::new(py, PyTimerHandle::new(ready.id(), when, &ready)) + Py::new(py, PyTimerHandle::new_timer(ready.id(), when, &ready)) } #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "PyLoop"))] diff --git a/src/engine/callbacks.rs b/src/engine/callbacks.rs index 43df01d..f2c4a3d 100644 --- a/src/engine/callbacks.rs +++ b/src/engine/callbacks.rs @@ -348,6 +348,10 @@ pub struct TimerCallback { state: Mutex, cancelled: AtomicBool, cancelled_timers: Arc, + // The public Rust API historically returns an immutable Arc. + // Preserve that explicit shared ownership for Rust embedders; Python timers + // use the detachable payload above instead. + legacy: Option>, } struct TimerState { @@ -366,9 +370,32 @@ impl TimerCallback { }), cancelled: AtomicBool::new(false), cancelled_timers, + legacy: None, } } + #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "TimerCallback"))] + pub(super) fn from_shared( + ready: Arc, + cancelled_timers: Arc, + ) -> Self { + Self { + id: ready.id(), + state: Mutex::new(TimerState { + ready: None, + queued: true, + }), + cancelled: AtomicBool::new(false), + cancelled_timers, + legacy: Some(ready), + } + } + + #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "TimerCallback"))] + pub(super) fn shared_ready(&self) -> Option<&ReadyCallback> { + self.legacy.as_deref() + } + #[inline] #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "TimerCallback"))] pub fn id(&self) -> CallbackId { @@ -379,6 +406,7 @@ impl TimerCallback { #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "TimerCallback"))] pub fn cancelled(&self) -> bool { self.cancelled.load(Ordering::Relaxed) + || self.legacy.as_ref().is_some_and(|ready| ready.cancelled()) } /// Linearize cancellation and detach Python ownership under the same lock @@ -398,6 +426,9 @@ impl TimerCallback { state.ready.take() }; drop(ready); + if let Some(ready) = &self.legacy { + ready.cancel(); + } } /// Queue removal and cancellation share the lock so repeated cancellation, @@ -405,7 +436,7 @@ impl TimerCallback { #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "TimerCallback"))] pub(super) fn dequeue(&self) { let mut state = self.state.lock().expect("poisoned timer state"); - if state.queued && self.cancelled() { + if state.queued && self.cancelled.load(Ordering::Relaxed) { self.cancelled_timers.fetch_sub(1, Ordering::Relaxed); } state.queued = false; @@ -452,19 +483,38 @@ impl Drop for TimerCallback { pub struct PyTimerHandle { callback_id: CallbackId, when: f64, - callback: Weak, + callback: TimerHandleCallback, cancelled: AtomicBool, } +enum TimerHandleCallback { + Owned(Weak), + Shared(Weak), +} + impl PyTimerHandle { #[inline] /// Creates a timer handle for a callback scheduled at loop time `when`. #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "PyTimerHandle"))] - pub fn new(callback_id: CallbackId, when: f64, callback: &Arc) -> Self { + pub fn new(callback_id: CallbackId, when: f64, callback: &Arc) -> Self { Self { callback_id, when, - callback: Arc::downgrade(callback), + callback: TimerHandleCallback::Shared(Arc::downgrade(callback)), + cancelled: AtomicBool::new(false), + } + } + + #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "PyTimerHandle"))] + pub(crate) fn new_timer( + callback_id: CallbackId, + when: f64, + callback: &Arc, + ) -> Self { + Self { + callback_id, + when, + callback: TimerHandleCallback::Owned(Arc::downgrade(callback)), cancelled: AtomicBool::new(false), } } @@ -701,7 +751,7 @@ mod tests { callback_with_args(py, 12, PyTuple::empty(py).unbind()), Arc::new(AtomicUsize::new(0)), )); - let timer = PyTimerHandle::new(12, 1.25, &ready); + let timer = PyTimerHandle::new_timer(12, 1.25, &ready); assert!((timer.when() - 1.25).abs() < f64::EPSILON); assert!(!timer.cancelled()); @@ -717,7 +767,7 @@ mod tests { callback_with_args(py, 13, PyTuple::empty(py).unbind()), Arc::new(AtomicUsize::new(0)), )); - let dropped_timer = PyTimerHandle::new(13, 2.0, &dropped); + let dropped_timer = PyTimerHandle::new_timer(13, 2.0, &dropped); drop(dropped); dropped_timer .cancel() @@ -732,8 +782,17 @@ impl PyTimerHandle { #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "PyTimerHandle"))] fn cancel(&self) -> PyResult<()> { self.cancelled.store(true, Ordering::Relaxed); - if let Some(callback) = self.callback.upgrade() { - callback.cancel(); + match &self.callback { + TimerHandleCallback::Owned(callback) => { + if let Some(callback) = callback.upgrade() { + callback.cancel(); + } + } + TimerHandleCallback::Shared(callback) => { + if let Some(callback) = callback.upgrade() { + callback.cancel(); + } + } } Ok(()) } diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs index eef4fe0..0378123 100644 --- a/src/engine/loop_core.rs +++ b/src/engine/loop_core.rs @@ -620,6 +620,35 @@ impl LoopCore { callback: Py, args: Py, context: Option>, + ) -> PyResult<(Arc, f64)> { + let (captured, context_needs_run) = capture_context(py, context)?; + let ready = Arc::new(ReadyCallback::new( + py, + self.next_callback_id(), + CallbackKind::Timer, + callback, + args, + captured, + context_needs_run, + )); + let timer = Arc::new(TimerCallback::from_shared( + Arc::clone(&ready), + Arc::clone(&self.cancelled_timers), + )); + let when = self.enqueue_timer(delay, timer)?; + Ok((ready, when)) + } + + /// Python timers own a detachable payload instead of the Rust API's + /// externally borrowable immutable Arc. + #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "LoopCore"))] + pub(crate) fn schedule_cancellable_timer( + self: &Arc, + py: Python<'_>, + delay: Duration, + callback: Py, + args: Py, + context: Option>, ) -> PyResult<(Arc, f64)> { let (captured, context_needs_run) = capture_context(py, context)?; let ready = Arc::new(TimerCallback::new( @@ -635,14 +664,25 @@ impl LoopCore { Arc::clone(&self.cancelled_timers), )); + let when = self.enqueue_timer(delay, Arc::clone(&ready))?; + Ok((ready, when)) + } + + #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "LoopCore"))] + fn enqueue_timer( + self: &Arc, + delay: Duration, + ready: Arc, + ) -> PyResult { let deadline = Instant::now() + delay; // Derive the public timestamp from the actual heap deadline. This also // avoids reading the monotonic clock twice for every timer. let when = deadline.duration_since(self.start).as_secs_f64(); + let seq = ready.id(); let entry = TimerEntry { - callback: Arc::clone(&ready), + callback: ready, when: deadline, - seq: ready.id(), + seq, }; let remote = ACTIVE_LOOP_TLS.with(|tls| { if std::ptr::eq(tls.core.get(), Arc::as_ptr(self)) && !tls.timers.get().is_null() { @@ -669,7 +709,7 @@ impl LoopCore { drop(state); self.wake.signal(); } - Ok((ready, when)) + Ok(when) } /// Runs callbacks and I/O until [`LoopCore::schedule_stop`] is processed. @@ -847,8 +887,9 @@ impl LoopCore { } } ReadyItem::TimerCallback(timer) => { - if let Some(callback) = timer.take_ready() { - if let Some(err) = self.execute_ready(py, Some(&loop_obj), &callback)? { + let owned = timer.take_ready(); + if let Some(callback) = owned.as_ref().or_else(|| timer.shared_ready()) { + if let Some(err) = self.execute_ready(py, Some(&loop_obj), callback)? { ready_error = Some(err); break; } @@ -1782,6 +1823,63 @@ mod wake_tests { }); } + #[test] + fn public_rust_timer_api_keeps_shared_ready_callback_and_handle_contract() { + crate::initialize_python_for_tests(); + Python::attach(|py| { + let core = LoopCore::new(); + let callback = py + .eval(pyo3::ffi::c_str!("lambda: 42"), None, None) + .expect("callable") + .unbind(); + let (ready, when): (Arc, f64) = core + .schedule_timer( + py, + Duration::ZERO, + callback, + PyTuple::empty(py).unbind(), + None, + ) + .expect("shared Rust timer"); + let handle = Py::new( + py, + super::super::callbacks::PyTimerHandle::new(ready.id(), when, &ready), + ) + .expect("public Rust handle constructor"); + assert_eq!( + ready + .invoke(py) + .expect("invoke shared payload") + .extract::(py) + .expect("result"), + 42 + ); + handle + .bind(py) + .call_method0("cancel") + .expect("cancel shared timer"); + assert!(ready.cancelled()); + let mut local = LocalTimers::new(&core); + let mut batch = VecDeque::new(); + local.collect(&mut batch); + let ReadyItem::TimerCallback(timer) = batch.pop_front().expect("timer") else { + panic!("expected timer callback"); + }; + assert!(timer.cancelled()); + assert!(timer.take_ready().is_none()); + let shared = timer.shared_ready().expect("legacy shared ownership"); + assert!(std::ptr::eq(shared, ready.as_ref())); + assert!( + core.execute_ready(py, None, shared) + .expect("skip cancelled") + .is_none() + ); + assert_eq!(core.cancelled_timers.load(Ordering::Relaxed), 0); + drop(local); + core.close().expect("close loop"); + }); + } + #[test] fn cancelled_timer_burst_is_compacted_behind_an_earlier_live_timer() { crate::initialize_python_for_tests(); @@ -1797,7 +1895,7 @@ mod wake_tests { .expect("sentinel"); for _ in 0..256 { let (callback, _) = core - .schedule_timer( + .schedule_cancellable_timer( py, Duration::from_secs(3600), py.None(),