Preserve public Rust timer API after cancellation lifetime fix

src/bindings/loop_api/methods.rs

@@ -102,15 +102,15 @@ 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,
             args.clone().unbind(),
             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"))]

src/engine/callbacks.rs

@@ -348,6 +348,10 @@ pub struct TimerCallback {
     state: Mutex<TimerState>,
     cancelled: AtomicBool,
     cancelled_timers: Arc<AtomicUsize>,
+    // The public Rust API historically returns an immutable Arc<ReadyCallback>.
+    // Preserve that explicit shared ownership for Rust embedders; Python timers
+    // use the detachable payload above instead.
+    legacy: Option<Arc<ReadyCallback>>,
 }
 
 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<ReadyCallback>,
+        cancelled_timers: Arc<AtomicUsize>,
+    ) -> 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,14 +426,17 @@ 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,
     /// cross-thread removal and loop shutdown count each tombstone once.
     #[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<TimerCallback>,
+    callback: TimerHandleCallback,
     cancelled: AtomicBool,
 }
 
+enum TimerHandleCallback {
+    Owned(Weak<TimerCallback>),
+    Shared(Weak<ReadyCallback>),
+}
+
 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<TimerCallback>) -> Self {
+    pub fn new(callback_id: CallbackId, when: f64, callback: &Arc<ReadyCallback>) -> 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<TimerCallback>,
+    ) -> 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(())
     }

src/engine/loop_core.rs

@@ -620,6 +620,35 @@ impl LoopCore {
         callback: Py<PyAny>,
         args: Py<PyTuple>,
         context: Option<Py<PyAny>>,
+    ) -> PyResult<(Arc<ReadyCallback>, 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<ReadyCallback>.
+    #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "LoopCore"))]
+    pub(crate) fn schedule_cancellable_timer(
+        self: &Arc<Self>,
+        py: Python<'_>,
+        delay: Duration,
+        callback: Py<PyAny>,
+        args: Py<PyTuple>,
+        context: Option<Py<PyAny>>,
     ) -> PyResult<(Arc<TimerCallback>, 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<Self>,
+        delay: Duration,
+        ready: Arc<TimerCallback>,
+    ) -> PyResult<f64> {
         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<ReadyCallback>, 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::<i32>(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(),