From e12823579fd275479a04a457d6aed7258d4b5036 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 23 Jul 2026 13:15:43 +0200 Subject: [PATCH 01/24] Add safe context watcher API --- src/context.rs | 500 +++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 2 + src/types/context.rs | 21 ++ src/types/mod.rs | 4 + 4 files changed, 527 insertions(+) create mode 100644 src/context.rs create mode 100644 src/types/context.rs diff --git a/src/context.rs b/src/context.rs new file mode 100644 index 00000000000..343f6236d7e --- /dev/null +++ b/src/context.rs @@ -0,0 +1,500 @@ +#![deny(clippy::undocumented_unsafe_blocks)] + +//! Safe bindings for watching changes to Python's current [`contextvars.Context`][1]. +//! +//! Context watchers are registered for the current Python interpreter and are invoked whenever +//! the current context changes. +//! +//! [1]: https://docs.python.org/3/library/contextvars.html#contextvars.Context + +use crate::err::{error_on_minusone, error_on_minusone_with_result}; +use crate::ffi_ptr_ext::FfiPtrExt; +use crate::types::any::PyAnyMethods; +use crate::types::PyContext; +use crate::{ffi, Borrowed, PyAny, PyResult, Python}; +use core::ffi::c_int; + +/// An event passed to a context watcher. +/// +/// This enum is non-exhaustive because CPython may add context watcher events in future versions. +#[doc(alias = "PyContextEvent")] +#[non_exhaustive] +pub enum ContextEvent<'a, 'py> { + /// The current context changed. + /// + /// The value is the new current context, or `None` when there is no current context. + Switched(Option>), + + /// An event which is not known to this version of PyO3. + Unknown { + /// The raw CPython event value. + raw_event: ffi::PyContextEvent, + + /// The event-specific object, if one was provided. + object: Option>, + }, +} + +/// A guard which keeps a context watcher registered. +/// +/// The watcher is registered for the current Python interpreter and is removed when this guard is +/// dropped. Use [`clear`][Self::clear] to remove it explicitly and observe any error returned by +/// CPython. +/// +/// This guard is bound to the [`Python`] attachment used to create it. It therefore cannot be sent +/// to another thread, moved outside that attachment, or moved into [`Python::detach`]. +/// +/// If this guard is forgotten, the watcher remains registered. This does not create a dangling +/// function pointer because [`register_context_watcher!`](crate::register_context_watcher) creates +/// a static, monomorphized trampoline. +#[must_use = "dropping the guard immediately unregisters the context watcher"] +pub struct ContextWatcherGuard<'py> { + watcher_id: c_int, + py: Python<'py>, + active: bool, +} + +impl ContextWatcherGuard<'_> { + /// Removes this watcher from the current Python interpreter. + /// + /// Dropping the guard also removes the watcher, but cannot report a failure to the caller. + #[doc(alias = "PyContext_ClearWatcher")] + pub fn clear(mut self) -> PyResult<()> { + self.active = false; + + // SAFETY: + // - `self.py` proves that the thread is attached to the interpreter for which the watcher + // was registered + // - `watcher_id` was returned by `PyContext_AddWatcher` + error_on_minusone(self.py, unsafe { + ffi::PyContext_ClearWatcher(self.watcher_id) + }) + } +} + +impl Drop for ContextWatcherGuard<'_> { + fn drop(&mut self) { + if !self.active { + return; + } + + self.active = false; + + // A destructor must not replace an exception which was already pending. The Python token + // stored in the guard proves that this thread is still attached to the correct interpreter. + // + // SAFETY: + // - the thread is attached, as guaranteed by `self.py` + // - `PyErr_GetRaisedException` returns an owned reference or NULL + // - `watcher_id` was returned by `PyContext_AddWatcher` + // - `PyErr_SetRaisedException` steals the owned reference returned above + unsafe { + let pending_exception = ffi::PyErr_GetRaisedException(); + let result = ffi::PyContext_ClearWatcher(self.watcher_id); + + if result == -1 { + ffi::PyErr_WriteUnraisable(core::ptr::null_mut()); + } + + if !pending_exception.is_null() { + // Be defensive in case an unraisable hook itself left an exception set. + ffi::PyErr_Clear(); + ffi::PyErr_SetRaisedException(pending_exception); + } + } + } +} + +/// Registers a safe Rust function as a context watcher for the current interpreter. +/// +/// The callback must be a function path and must have this signature: +/// +/// ```rust +/// # #![cfg(all(Py_3_14, not(Py_LIMITED_API), not(any(PyPy, GraalPy, RustPython))))] +/// use pyo3::context::ContextEvent; +/// use pyo3::prelude::*; +/// +/// fn context_changed( +/// py: Python<'_>, +/// event: ContextEvent<'_, '_>, +/// ) -> PyResult<()> { +/// let _ = py; +/// let _ = event; +/// Ok(()) +/// } +/// +/// # fn main() -> PyResult<()> { +/// Python::attach(|py| { +/// let _watcher = pyo3::register_context_watcher!(py, context_changed)?; +/// Ok(()) +/// }) +/// # } +/// ``` +/// +/// A function path is required because CPython's context watcher callback has no user-data +/// pointer. The macro creates a unique static trampoline for the function, avoiding global callback +/// storage. State can still be shared through safe static synchronization primitives. +/// +/// The callback may run concurrently on free-threaded Python builds. Panics and returned +/// [`PyErr`][crate::PyErr] values are reported as unraisable exceptions and never unwind across the +/// C boundary. +#[doc(alias = "PyContext_AddWatcher")] +#[macro_export] +macro_rules! register_context_watcher { + ($py:expr, $callback:path) => {{ + struct Callback; + + impl $crate::context::impl_::ContextWatcherCallbackDef for Callback { + const CALLBACK: $crate::context::impl_::ContextWatcherCallback = $callback; + } + + $crate::context::impl_::register::($py) + }}; +} + +/// Implementation details used by [`register_context_watcher!`](crate::register_context_watcher). +#[doc(hidden)] +pub mod impl_ { + use super::*; + + /// The safe callback signature accepted by context watcher trampolines. + pub type ContextWatcherCallback = + for<'a, 'py> fn(Python<'py>, ContextEvent<'a, 'py>) -> PyResult<()>; + + /// Associates a generated trampoline type with its Rust callback. + pub trait ContextWatcherCallbackDef { + /// The Rust callback invoked by the generated C trampoline. + const CALLBACK: ContextWatcherCallback; + } + + /// Registers the trampoline specialized for `Callback`. + pub fn register( + py: Python<'_>, + ) -> PyResult> { + // SAFETY: + // - `py` proves that the thread is attached + // - `context_watcher::` is a static C-compatible function + let watcher_id = unsafe { ffi::PyContext_AddWatcher(context_watcher::) }; + let watcher_id = error_on_minusone_with_result(py, watcher_id)?; + + Ok(ContextWatcherGuard { + watcher_id, + py, + active: true, + }) + } + + /// C-compatible trampoline for a context watcher callback. + /// + /// # Safety + /// + /// - The thread must be attached to Python. + /// - `object` must follow the contract for the supplied `event`. + pub unsafe extern "C" fn context_watcher( + event: ffi::PyContextEvent, + object: *mut ffi::PyObject, + ) -> c_int { + // A context watcher may be called with an exception already set. Save it before invoking + // arbitrary Rust code so that safe PyO3 APIs can be used normally inside the callback. + // + // SAFETY: the caller guarantees that the thread is attached. + let pending_exception = unsafe { ffi::PyErr_GetRaisedException() }; + + // SAFETY: the caller guarantees that the thread is attached. `trampoline` catches all + // panics and converts callback errors into a Python exception with a -1 return value. + let result = unsafe { + crate::impl_::trampoline::trampoline(|py| { + let borrow_guard = (); + // SAFETY: + // - CPython guarantees that `object` follows the contract for `event` + // - `borrow_guard` limits the resulting borrow to this callback invocation + let event = event_from_raw(py, event, object, &borrow_guard); + + (Callback::CALLBACK)(py, event)?; + + // Although normal PyO3 APIs return errors as `PyResult`, `PyErr::restore` can be + // called directly. Do not allow an Ok return with an exception still set. + if crate::PyErr::occurred(py) { + return Err(crate::PyErr::fetch(py)); + } + + Ok(0) + }) + }; + + if pending_exception.is_null() { + return result; + } + + // When an exception was already pending on entry, CPython requires the callback to return + // 0 with that same exception still set. Report a new callback error ourselves before + // restoring the original exception. + // + // SAFETY: + // - the thread is attached + // - `object` is valid for the duration of the callback or NULL + // - `pending_exception` is an owned reference from `PyErr_GetRaisedException` + // - `PyErr_SetRaisedException` steals that reference + unsafe { + if result == -1 { + ffi::PyErr_WriteUnraisable(object); + } + + // Be defensive in case an unraisable hook itself left an exception set. + ffi::PyErr_Clear(); + ffi::PyErr_SetRaisedException(pending_exception); + } + + 0 + } + + unsafe fn event_from_raw<'a, 'py>( + py: Python<'py>, + event: ffi::PyContextEvent, + object: *mut ffi::PyObject, + _borrow_guard: &'a (), + ) -> ContextEvent<'a, 'py> { + match event { + ffi::Py_CONTEXT_SWITCHED => { + // SAFETY: CPython documents a non-null context object or `None` for this event. + let object = unsafe { object.assume_borrowed(py) }; + + if object.is_none() { + ContextEvent::Switched(None) + } else { + // SAFETY: CPython guarantees that a non-None object for this event is a + // `contextvars.Context`. + ContextEvent::Switched(Some(unsafe { object.cast_unchecked() })) + } + } + raw_event => { + // SAFETY: unknown events may have a NULL object; a non-null object is borrowed for + // at least the callback duration, which is bounded by `_borrow_guard`. + let object = unsafe { object.assume_borrowed_or_opt(py) }; + ContextEvent::Unknown { raw_event, object } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::impl_::{context_watcher, ContextWatcherCallback, ContextWatcherCallbackDef}; + use super::ContextEvent; + use crate::exceptions::{PyRuntimeError, PyValueError}; + use crate::test_utils::UnraisableCapture; + use crate::types::{PyAnyMethods, PyContext}; + use crate::{ffi, PyErr, PyResult, Python}; + use alloc::string::ToString; + use core::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; + use static_assertions::assert_not_impl_any; + + static SWITCH_COUNT: AtomicUsize = AtomicUsize::new(0); + static SAW_CONTEXT: AtomicBool = AtomicBool::new(false); + static SAW_NONE: AtomicBool = AtomicBool::new(false); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_switch(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if let ContextEvent::Switched(context) = event { + SWITCH_COUNT.fetch_add(1, Ordering::Relaxed); + match context { + Some(context) => { + assert!(context.is_exact_instance_of::()); + SAW_CONTEXT.store(true, Ordering::Relaxed); + } + None => SAW_NONE.store(true, Ordering::Relaxed), + } + } + Ok(()) + } + + #[test] + fn watcher_is_cleared_on_drop() { + Python::attach(|py| { + SWITCH_COUNT.store(0, Ordering::Relaxed); + SAW_CONTEXT.store(false, Ordering::Relaxed); + SAW_NONE.store(false, Ordering::Relaxed); + + let watcher = crate::register_context_watcher!(py, record_switch).unwrap(); + py.run( + c"import contextvars\ncontextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + + let count_after_first_run = SWITCH_COUNT.load(Ordering::Relaxed); + assert!(count_after_first_run >= 2); + assert!(SAW_CONTEXT.load(Ordering::Relaxed)); + assert!(SAW_NONE.load(Ordering::Relaxed)); + + drop(watcher); + + py.run(c"contextvars.Context().run(lambda: None)", None, None) + .unwrap(); + assert_eq!(SWITCH_COUNT.load(Ordering::Relaxed), count_after_first_run); + }); + } + + static EXPLICIT_CLEAR_COUNT: AtomicUsize = AtomicUsize::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_explicit_clear(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if matches!(event, ContextEvent::Switched(_)) { + EXPLICIT_CLEAR_COUNT.fetch_add(1, Ordering::Relaxed); + } + Ok(()) + } + + #[test] + fn watcher_can_be_cleared_explicitly() { + Python::attach(|py| { + EXPLICIT_CLEAR_COUNT.store(0, Ordering::Relaxed); + + let watcher = crate::register_context_watcher!(py, record_explicit_clear).unwrap(); + watcher.clear().unwrap(); + + py.run( + c"import contextvars\ncontextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + assert_eq!(EXPLICIT_CLEAR_COUNT.load(Ordering::Relaxed), 0); + }); + } + + fn fail_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { + Err(PyRuntimeError::new_err("watcher failed")) + } + + struct FailingCallback; + + impl ContextWatcherCallbackDef for FailingCallback { + const CALLBACK: ContextWatcherCallback = fail_callback; + } + + #[test] + fn callback_error_is_returned_without_a_pending_exception() { + Python::attach(|py| { + // SAFETY: the thread is attached and None is valid for Py_CONTEXT_SWITCHED. + let result = unsafe { + context_watcher::(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None()) + }; + + assert_eq!(result, -1); + let error = PyErr::fetch(py); + assert!(error.is_instance_of::(py)); + }); + } + + #[test] + fn callback_error_preserves_a_pending_exception() { + Python::attach(|py| { + UnraisableCapture::enter(py, |capture| { + PyValueError::new_err("original error").restore(py); + + // SAFETY: the thread is attached and None is valid for Py_CONTEXT_SWITCHED. + let result = unsafe { + context_watcher::(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None()) + }; + + assert_eq!(result, 0); + + let original_error = PyErr::fetch(py); + assert!(original_error.is_instance_of::(py)); + assert_eq!(original_error.to_string(), "ValueError: original error"); + + let (watcher_error, object) = + capture.take_capture().expect("missing unraisable error"); + assert!(watcher_error.is_instance_of::(py)); + assert!(object.is_none()); + }); + }); + } + + #[test] + fn registered_callback_errors_are_unraisable() { + Python::attach(|py| { + UnraisableCapture::enter(py, |capture| { + let watcher = crate::register_context_watcher!(py, fail_callback).unwrap(); + + py.run( + c"import contextvars\ncontextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + + let (watcher_error, _) = capture.take_capture().expect("missing unraisable error"); + assert!(watcher_error.is_instance_of::(py)); + + drop(watcher); + }); + }); + } + + fn panic_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { + panic!("context watcher panic") + } + + struct PanickingCallback; + + impl ContextWatcherCallbackDef for PanickingCallback { + const CALLBACK: ContextWatcherCallback = panic_callback; + } + + #[test] + fn callback_panic_does_not_cross_ffi_boundary() { + Python::attach(|py| { + // SAFETY: the thread is attached and None is valid for Py_CONTEXT_SWITCHED. + let result = unsafe { + context_watcher::(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None()) + }; + + assert_eq!(result, -1); + assert!(PyErr::occurred(py)); + + // SAFETY: the test has observed and intentionally discards the panic exception. + unsafe { ffi::PyErr_Clear() }; + }); + } + + static UNKNOWN_EVENT: AtomicU32 = AtomicU32::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_unknown(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if let ContextEvent::Unknown { raw_event, object } = event { + UNKNOWN_EVENT.store(raw_event, Ordering::Relaxed); + assert!(object.is_none()); + } + Ok(()) + } + + struct UnknownCallback; + + impl ContextWatcherCallbackDef for UnknownCallback { + const CALLBACK: ContextWatcherCallback = record_unknown; + } + + #[test] + fn unknown_events_are_forwarded() { + const FUTURE_EVENT: ffi::PyContextEvent = 123; + + Python::attach(|_py| { + UNKNOWN_EVENT.store(0, Ordering::Relaxed); + + // SAFETY: the thread is attached and unknown events accept a null object. + let result = + unsafe { context_watcher::(FUTURE_EVENT, core::ptr::null_mut()) }; + + assert_eq!(result, 0); + assert_eq!(UNKNOWN_EVENT.load(Ordering::Relaxed), FUTURE_EVENT); + }); + } + + #[test] + fn watcher_guard_is_not_send_or_sync() { + assert_not_impl_any!(super::ContextWatcherGuard<'_>: Send, Sync); + } +} diff --git a/src/lib.rs b/src/lib.rs index 8b0eeabc28e..6f6a91304c5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -422,6 +422,8 @@ pub mod impl_; pub mod buffer; pub mod call; +#[cfg(all(Py_3_14, not(Py_LIMITED_API), not(any(PyPy, GraalPy, RustPython))))] +pub mod context; pub mod conversion; mod conversions; #[cfg(feature = "experimental-async")] diff --git a/src/types/context.rs b/src/types/context.rs new file mode 100644 index 00000000000..43791f8b6a5 --- /dev/null +++ b/src/types/context.rs @@ -0,0 +1,21 @@ +#![deny(clippy::undocumented_unsafe_blocks)] + +use crate::{ffi, PyAny}; + +/// Represents a Python [`contextvars.Context`][1] object. +/// +/// Values of this type are accessed via PyO3's smart pointers, e.g. as +/// [`Py`][crate::Py] or [`Bound<'py, PyContext>`][crate::Bound]. +/// +/// [1]: https://docs.python.org/3/library/contextvars.html#contextvars.Context +#[repr(transparent)] +pub struct PyContext(PyAny); + +pyobject_native_type_core!( + PyContext, + pyobject_native_static_type_object!(ffi::PyContext_Type), + "contextvars", + "Context", + #module=Some("contextvars"), + #checkfunction=ffi::PyContext_CheckExact +); diff --git a/src/types/mod.rs b/src/types/mod.rs index cb71e626b28..d2b21d8219e 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -17,6 +17,8 @@ pub use self::capsule::{PyCapsule, PyCapsuleMethods}; pub use self::code::{PyCode, PyCodeMethods}; #[doc(inline)] pub use self::complex::{PyComplex, PyComplexMethods}; +#[cfg(all(Py_3_14, not(Py_LIMITED_API), not(any(PyPy, GraalPy, RustPython))))] +pub use self::context::PyContext; #[doc(inline)] pub use self::datetime::{PyDate, PyDateTime, PyDelta, PyTime, PyTzInfo, PyTzInfoAccess}; #[cfg(not(Py_LIMITED_API))] @@ -355,6 +357,8 @@ pub mod bytes; pub mod capsule; pub mod code; pub(crate) mod complex; +#[cfg(all(Py_3_14, not(Py_LIMITED_API), not(any(PyPy, GraalPy, RustPython))))] +mod context; pub mod datetime; pub mod dict; mod ellipsis; From fc02fb9d0357672c511254316dad63f319904716 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 23 Jul 2026 13:16:38 +0200 Subject: [PATCH 02/24] Add newsfragment --- newsfragments/6227.added.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 newsfragments/6227.added.md diff --git a/newsfragments/6227.added.md b/newsfragments/6227.added.md new file mode 100644 index 00000000000..a98c147f868 --- /dev/null +++ b/newsfragments/6227.added.md @@ -0,0 +1 @@ +Add safe bindings for registering context watchers on Python 3.14+. From c6f424baee646d5fb97ffafc2c2d83bb20d9ed77 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 23 Jul 2026 13:30:27 +0200 Subject: [PATCH 03/24] Refine context API availability --- newsfragments/6227.added.md | 3 ++- src/context.rs | 1 + src/types/mod.rs | 4 ++-- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/newsfragments/6227.added.md b/newsfragments/6227.added.md index a98c147f868..00e32d20147 100644 --- a/newsfragments/6227.added.md +++ b/newsfragments/6227.added.md @@ -1 +1,2 @@ -Add safe bindings for registering context watchers on Python 3.14+. +Add `PyContext` and the Python 3.14+ `register_context_watcher!`, `ContextEvent`, and +`ContextWatcherGuard` APIs. diff --git a/src/context.rs b/src/context.rs index 343f6236d7e..72044b27330 100644 --- a/src/context.rs +++ b/src/context.rs @@ -18,6 +18,7 @@ use core::ffi::c_int; /// /// This enum is non-exhaustive because CPython may add context watcher events in future versions. #[doc(alias = "PyContextEvent")] +#[derive(Debug)] #[non_exhaustive] pub enum ContextEvent<'a, 'py> { /// The current context changed. diff --git a/src/types/mod.rs b/src/types/mod.rs index d2b21d8219e..a6c956334ed 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -17,7 +17,7 @@ pub use self::capsule::{PyCapsule, PyCapsuleMethods}; pub use self::code::{PyCode, PyCodeMethods}; #[doc(inline)] pub use self::complex::{PyComplex, PyComplexMethods}; -#[cfg(all(Py_3_14, not(Py_LIMITED_API), not(any(PyPy, GraalPy, RustPython))))] +#[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy, RustPython)))] pub use self::context::PyContext; #[doc(inline)] pub use self::datetime::{PyDate, PyDateTime, PyDelta, PyTime, PyTzInfo, PyTzInfoAccess}; @@ -357,7 +357,7 @@ pub mod bytes; pub mod capsule; pub mod code; pub(crate) mod complex; -#[cfg(all(Py_3_14, not(Py_LIMITED_API), not(any(PyPy, GraalPy, RustPython))))] +#[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy, RustPython)))] mod context; pub mod datetime; pub mod dict; From 3f707771bc4994e1420af2b426cddc593624aa16 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 23 Jul 2026 13:45:30 +0200 Subject: [PATCH 04/24] Guard macro-dependent context tests --- src/context.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/context.rs b/src/context.rs index 72044b27330..9e43746df2b 100644 --- a/src/context.rs +++ b/src/context.rs @@ -282,10 +282,14 @@ pub mod impl_ { mod tests { use super::impl_::{context_watcher, ContextWatcherCallback, ContextWatcherCallbackDef}; use super::ContextEvent; - use crate::exceptions::{PyRuntimeError, PyValueError}; + use crate::exceptions::PyRuntimeError; + #[cfg(feature = "macros")] + use crate::exceptions::PyValueError; + #[cfg(feature = "macros")] use crate::test_utils::UnraisableCapture; use crate::types::{PyAnyMethods, PyContext}; use crate::{ffi, PyErr, PyResult, Python}; + #[cfg(feature = "macros")] use alloc::string::ToString; use core::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; use static_assertions::assert_not_impl_any; @@ -390,6 +394,7 @@ mod tests { } #[test] + #[cfg(feature = "macros")] fn callback_error_preserves_a_pending_exception() { Python::attach(|py| { UnraisableCapture::enter(py, |capture| { @@ -415,6 +420,7 @@ mod tests { } #[test] + #[cfg(feature = "macros")] fn registered_callback_errors_are_unraisable() { Python::attach(|py| { UnraisableCapture::enter(py, |capture| { From d238aee50ed8b017e6b31cc64c81d0ae5ee3ebad Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 23 Jul 2026 15:13:08 +0200 Subject: [PATCH 05/24] Refine context watcher safety and tests --- newsfragments/6227.added.md | 4 +- src/context.rs | 81 ++++++++++++++++++++++++++++++------- src/lib.rs | 7 +++- src/types/context.rs | 23 +++++++++++ 4 files changed, 98 insertions(+), 17 deletions(-) diff --git a/newsfragments/6227.added.md b/newsfragments/6227.added.md index 00e32d20147..17bd50f5d7a 100644 --- a/newsfragments/6227.added.md +++ b/newsfragments/6227.added.md @@ -1,2 +1,2 @@ -Add `PyContext` and the Python 3.14+ `register_context_watcher!`, `ContextEvent`, and -`ContextWatcherGuard` APIs. +Add `PyContext`, as well as `register_context_watcher!`, `ContextEvent`, and +`ContextWatcherGuard` on GIL-enabled CPython 3.14+. diff --git a/src/context.rs b/src/context.rs index 9e43746df2b..1b4f936d9e6 100644 --- a/src/context.rs +++ b/src/context.rs @@ -111,7 +111,12 @@ impl Drop for ContextWatcherGuard<'_> { /// The callback must be a function path and must have this signature: /// /// ```rust -/// # #![cfg(all(Py_3_14, not(Py_LIMITED_API), not(any(PyPy, GraalPy, RustPython))))] +/// # #![cfg(all( +/// # Py_3_14, +/// # not(Py_GIL_DISABLED), +/// # not(Py_LIMITED_API), +/// # not(any(PyPy, GraalPy, RustPython)) +/// # ))] /// use pyo3::context::ContextEvent; /// use pyo3::prelude::*; /// @@ -136,9 +141,8 @@ impl Drop for ContextWatcherGuard<'_> { /// pointer. The macro creates a unique static trampoline for the function, avoiding global callback /// storage. State can still be shared through safe static synchronization primitives. /// -/// The callback may run concurrently on free-threaded Python builds. Panics and returned -/// [`PyErr`][crate::PyErr] values are reported as unraisable exceptions and never unwind across the -/// C boundary. +/// Panics and returned [`PyErr`][crate::PyErr] values are reported as unraisable exceptions and +/// never unwind across the C boundary. #[doc(alias = "PyContext_AddWatcher")] #[macro_export] macro_rules! register_context_watcher { @@ -283,7 +287,6 @@ mod tests { use super::impl_::{context_watcher, ContextWatcherCallback, ContextWatcherCallbackDef}; use super::ContextEvent; use crate::exceptions::PyRuntimeError; - #[cfg(feature = "macros")] use crate::exceptions::PyValueError; #[cfg(feature = "macros")] use crate::test_utils::UnraisableCapture; @@ -296,18 +299,14 @@ mod tests { static SWITCH_COUNT: AtomicUsize = AtomicUsize::new(0); static SAW_CONTEXT: AtomicBool = AtomicBool::new(false); - static SAW_NONE: AtomicBool = AtomicBool::new(false); #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] fn record_switch(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { if let ContextEvent::Switched(context) = event { SWITCH_COUNT.fetch_add(1, Ordering::Relaxed); - match context { - Some(context) => { - assert!(context.is_exact_instance_of::()); - SAW_CONTEXT.store(true, Ordering::Relaxed); - } - None => SAW_NONE.store(true, Ordering::Relaxed), + if let Some(context) = context { + assert!(context.is_exact_instance_of::()); + SAW_CONTEXT.store(true, Ordering::Relaxed); } } Ok(()) @@ -318,7 +317,6 @@ mod tests { Python::attach(|py| { SWITCH_COUNT.store(0, Ordering::Relaxed); SAW_CONTEXT.store(false, Ordering::Relaxed); - SAW_NONE.store(false, Ordering::Relaxed); let watcher = crate::register_context_watcher!(py, record_switch).unwrap(); py.run( @@ -331,7 +329,6 @@ mod tests { let count_after_first_run = SWITCH_COUNT.load(Ordering::Relaxed); assert!(count_after_first_run >= 2); assert!(SAW_CONTEXT.load(Ordering::Relaxed)); - assert!(SAW_NONE.load(Ordering::Relaxed)); drop(watcher); @@ -369,6 +366,62 @@ mod tests { }); } + static DUPLICATE_CALLBACK_COUNT: AtomicUsize = AtomicUsize::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_duplicate_callback(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if matches!(event, ContextEvent::Switched(_)) { + DUPLICATE_CALLBACK_COUNT.fetch_add(1, Ordering::Relaxed); + } + Ok(()) + } + + #[test] + fn multiple_watchers_can_register_the_same_callback() { + Python::attach(|py| { + DUPLICATE_CALLBACK_COUNT.store(0, Ordering::Relaxed); + + let first = crate::register_context_watcher!(py, record_duplicate_callback).unwrap(); + let second = crate::register_context_watcher!(py, record_duplicate_callback).unwrap(); + + py.run( + c"import contextvars\ncontextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + assert!(DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed) >= 4); + + drop(first); + let count_with_both = DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed); + py.run(c"contextvars.Context().run(lambda: None)", None, None) + .unwrap(); + assert!(DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed) >= count_with_both + 2); + + drop(second); + let count_after_drop = DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed); + py.run(c"contextvars.Context().run(lambda: None)", None, None) + .unwrap(); + assert_eq!( + DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed), + count_after_drop + ); + }); + } + + #[test] + fn dropping_watcher_preserves_a_pending_exception() { + Python::attach(|py| { + let watcher = crate::register_context_watcher!(py, record_explicit_clear).unwrap(); + PyValueError::new_err("original error").restore(py); + + drop(watcher); + + let error = PyErr::fetch(py); + assert!(error.is_instance_of::(py)); + }); + } + fn fail_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { Err(PyRuntimeError::new_err("watcher failed")) } diff --git a/src/lib.rs b/src/lib.rs index 6f6a91304c5..a50666397d5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -422,7 +422,12 @@ pub mod impl_; pub mod buffer; pub mod call; -#[cfg(all(Py_3_14, not(Py_LIMITED_API), not(any(PyPy, GraalPy, RustPython))))] +#[cfg(all( + Py_3_14, + not(Py_GIL_DISABLED), + not(Py_LIMITED_API), + not(any(PyPy, GraalPy, RustPython)) +))] pub mod context; pub mod conversion; mod conversions; diff --git a/src/types/context.rs b/src/types/context.rs index 43791f8b6a5..cb5b3e96054 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -19,3 +19,26 @@ pyobject_native_type_core!( #module=Some("contextvars"), #checkfunction=ffi::PyContext_CheckExact ); + +#[cfg(test)] +mod tests { + use super::PyContext; + use crate::types::PyAnyMethods; + use crate::Python; + + #[test] + fn context_type() { + Python::attach(|py| { + let context = py + .import(c"contextvars") + .unwrap() + .getattr(c"Context") + .unwrap() + .call0() + .unwrap(); + + assert!(context.is_exact_instance_of::()); + context.cast::().unwrap(); + }); + } +} From 694a3a9b0879ae396937a21db6e00d66696397b0 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Mon, 27 Jul 2026 13:49:40 +0200 Subject: [PATCH 06/24] add forgotten guard test --- src/context.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/context.rs b/src/context.rs index 1b4f936d9e6..e224015bc11 100644 --- a/src/context.rs +++ b/src/context.rs @@ -557,4 +557,22 @@ mod tests { fn watcher_guard_is_not_send_or_sync() { assert_not_impl_any!(super::ContextWatcherGuard<'_>: Send, Sync); } + + #[test] + fn forgotten_guard() { + Python::attach(|py| { + SWITCH_COUNT.store(0, Ordering::Relaxed); + let watcher = crate::register_context_watcher!(py, record_switch).unwrap(); + core::mem::forget(watcher); + + py.run( + c"import contextvars; contextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + + assert!(SWITCH_COUNT.load(Ordering::Relaxed) > 0); + }); + } } From dafb6af68addd92b5a64efa8e978e0db76b2fba1 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Mon, 27 Jul 2026 15:24:00 +0200 Subject: [PATCH 07/24] Isolate forgotten context watcher test --- src/context.rs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/context.rs b/src/context.rs index e224015bc11..37b83420456 100644 --- a/src/context.rs +++ b/src/context.rs @@ -558,11 +558,21 @@ mod tests { assert_not_impl_any!(super::ContextWatcherGuard<'_>: Send, Sync); } + static FORGOTTEN_GUARD_COUNT: AtomicUsize = AtomicUsize::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_forgotten_guard(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if matches!(event, ContextEvent::Switched(_)) { + FORGOTTEN_GUARD_COUNT.fetch_add(1, Ordering::Relaxed); + } + Ok(()) + } + #[test] fn forgotten_guard() { Python::attach(|py| { - SWITCH_COUNT.store(0, Ordering::Relaxed); - let watcher = crate::register_context_watcher!(py, record_switch).unwrap(); + FORGOTTEN_GUARD_COUNT.store(0, Ordering::Relaxed); + let watcher = crate::register_context_watcher!(py, record_forgotten_guard).unwrap(); core::mem::forget(watcher); py.run( @@ -572,7 +582,7 @@ mod tests { ) .unwrap(); - assert!(SWITCH_COUNT.load(Ordering::Relaxed) > 0); + assert!(FORGOTTEN_GUARD_COUNT.load(Ordering::Relaxed) > 0); }); } } From b5d47c43cb4711c8a010b098a23bffb5ccd07dfb Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Tue, 28 Jul 2026 12:22:01 +0200 Subject: [PATCH 08/24] Simplify context watcher event conversion --- src/context.rs | 74 +++++++++++++++++++------------------------------- 1 file changed, 28 insertions(+), 46 deletions(-) diff --git a/src/context.rs b/src/context.rs index 37b83420456..bd0f2ed5c64 100644 --- a/src/context.rs +++ b/src/context.rs @@ -91,9 +91,7 @@ impl Drop for ContextWatcherGuard<'_> { // - `PyErr_SetRaisedException` steals the owned reference returned above unsafe { let pending_exception = ffi::PyErr_GetRaisedException(); - let result = ffi::PyContext_ClearWatcher(self.watcher_id); - - if result == -1 { + if ffi::PyContext_ClearWatcher(self.watcher_id) == -1 { ffi::PyErr_WriteUnraisable(core::ptr::null_mut()); } @@ -121,11 +119,9 @@ impl Drop for ContextWatcherGuard<'_> { /// use pyo3::prelude::*; /// /// fn context_changed( -/// py: Python<'_>, -/// event: ContextEvent<'_, '_>, +/// _py: Python<'_>, +/// _event: ContextEvent<'_, '_>, /// ) -> PyResult<()> { -/// let _ = py; -/// let _ = event; /// Ok(()) /// } /// @@ -179,8 +175,9 @@ pub mod impl_ { // SAFETY: // - `py` proves that the thread is attached // - `context_watcher::` is a static C-compatible function - let watcher_id = unsafe { ffi::PyContext_AddWatcher(context_watcher::) }; - let watcher_id = error_on_minusone_with_result(py, watcher_id)?; + let watcher_id = error_on_minusone_with_result(py, unsafe { + ffi::PyContext_AddWatcher(context_watcher::) + })?; Ok(ContextWatcherGuard { watcher_id, @@ -205,15 +202,29 @@ pub mod impl_ { // SAFETY: the caller guarantees that the thread is attached. let pending_exception = unsafe { ffi::PyErr_GetRaisedException() }; - // SAFETY: the caller guarantees that the thread is attached. `trampoline` catches all - // panics and converts callback errors into a Python exception with a -1 return value. + // SAFETY: + // - the caller guarantees that the thread is attached and `object` follows the contract + // for `event` + // - `trampoline` catches panics and converts callback errors into a Python exception + // - the callback's higher-ranked signature prevents borrowed event data from escaping let result = unsafe { crate::impl_::trampoline::trampoline(|py| { - let borrow_guard = (); - // SAFETY: - // - CPython guarantees that `object` follows the contract for `event` - // - `borrow_guard` limits the resulting borrow to this callback invocation - let event = event_from_raw(py, event, object, &borrow_guard); + let event = match event { + ffi::Py_CONTEXT_SWITCHED => { + let object = object.assume_borrowed(py); + + if object.is_none() { + ContextEvent::Switched(None) + } else { + ContextEvent::Switched(Some(object.cast_unchecked())) + } + } + + raw_event => { + let object = object.assume_borrowed_or_opt(py); + ContextEvent::Unknown { raw_event, object } + } + }; (Callback::CALLBACK)(py, event)?; @@ -252,42 +263,13 @@ pub mod impl_ { 0 } - - unsafe fn event_from_raw<'a, 'py>( - py: Python<'py>, - event: ffi::PyContextEvent, - object: *mut ffi::PyObject, - _borrow_guard: &'a (), - ) -> ContextEvent<'a, 'py> { - match event { - ffi::Py_CONTEXT_SWITCHED => { - // SAFETY: CPython documents a non-null context object or `None` for this event. - let object = unsafe { object.assume_borrowed(py) }; - - if object.is_none() { - ContextEvent::Switched(None) - } else { - // SAFETY: CPython guarantees that a non-None object for this event is a - // `contextvars.Context`. - ContextEvent::Switched(Some(unsafe { object.cast_unchecked() })) - } - } - raw_event => { - // SAFETY: unknown events may have a NULL object; a non-null object is borrowed for - // at least the callback duration, which is bounded by `_borrow_guard`. - let object = unsafe { object.assume_borrowed_or_opt(py) }; - ContextEvent::Unknown { raw_event, object } - } - } - } } #[cfg(test)] mod tests { use super::impl_::{context_watcher, ContextWatcherCallback, ContextWatcherCallbackDef}; use super::ContextEvent; - use crate::exceptions::PyRuntimeError; - use crate::exceptions::PyValueError; + use crate::exceptions::{PyRuntimeError, PyValueError}; #[cfg(feature = "macros")] use crate::test_utils::UnraisableCapture; use crate::types::{PyAnyMethods, PyContext}; From cc6106b697c83f96eea9a9ec7b8031b873d2d5fb Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 13 Aug 2026 16:27:17 +0200 Subject: [PATCH 09/24] Address context watcher review nits Fix the newsfragment's overclaim that PyContext requires GIL-enabled CPython 3.14+ (it's available on all supported versions; only the watcher API needs 3.14+), clarify the SAFETY comments on ContextWatcherGuard around single-interpreter attachment, and drop a redundant doctest cfg attribute already covered by the module-level gate in lib.rs. --- newsfragments/6227.added.md | 3 ++- src/context.rs | 15 ++++++--------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/newsfragments/6227.added.md b/newsfragments/6227.added.md index 17bd50f5d7a..e4b1d0635ec 100644 --- a/newsfragments/6227.added.md +++ b/newsfragments/6227.added.md @@ -1,2 +1,3 @@ Add `PyContext`, as well as `register_context_watcher!`, `ContextEvent`, and -`ContextWatcherGuard` on GIL-enabled CPython 3.14+. +`ContextWatcherGuard` for watching `contextvars.Context` on GIL-enabled CPython +3.14+. `PyContext` itself is available on all supported CPython versions. diff --git a/src/context.rs b/src/context.rs index bd0f2ed5c64..4e4f2632cd6 100644 --- a/src/context.rs +++ b/src/context.rs @@ -64,8 +64,9 @@ impl ContextWatcherGuard<'_> { self.active = false; // SAFETY: - // - `self.py` proves that the thread is attached to the interpreter for which the watcher - // was registered + // - `self.py` proves that the thread is attached to an interpreter; PyO3 does not + // currently support attaching to more than one interpreter, so this is the interpreter + // for which the watcher was registered // - `watcher_id` was returned by `PyContext_AddWatcher` error_on_minusone(self.py, unsafe { ffi::PyContext_ClearWatcher(self.watcher_id) @@ -82,7 +83,9 @@ impl Drop for ContextWatcherGuard<'_> { self.active = false; // A destructor must not replace an exception which was already pending. The Python token - // stored in the guard proves that this thread is still attached to the correct interpreter. + // stored in the guard proves that this thread is still attached to an interpreter; PyO3 + // does not currently support attaching to more than one interpreter, so this is the same + // interpreter the watcher was registered on. // // SAFETY: // - the thread is attached, as guaranteed by `self.py` @@ -109,12 +112,6 @@ impl Drop for ContextWatcherGuard<'_> { /// The callback must be a function path and must have this signature: /// /// ```rust -/// # #![cfg(all( -/// # Py_3_14, -/// # not(Py_GIL_DISABLED), -/// # not(Py_LIMITED_API), -/// # not(any(PyPy, GraalPy, RustPython)) -/// # ))] /// use pyo3::context::ContextEvent; /// use pyo3::prelude::*; /// From 2a9397ea10c90b0a15a9893a869c9894f5fe82a9 Mon Sep 17 00:00:00 2001 From: David Hewitt Date: Sat, 22 Aug 2026 19:42:17 +0100 Subject: [PATCH 10/24] add context watcher guard as associated method --- src/context.rs | 567 ----------------------------------------- src/lib.rs | 7 - src/types/context.rs | 585 ++++++++++++++++++++++++++++++++++++++++++- src/types/mod.rs | 2 +- 4 files changed, 583 insertions(+), 578 deletions(-) delete mode 100644 src/context.rs diff --git a/src/context.rs b/src/context.rs deleted file mode 100644 index 4e4f2632cd6..00000000000 --- a/src/context.rs +++ /dev/null @@ -1,567 +0,0 @@ -#![deny(clippy::undocumented_unsafe_blocks)] - -//! Safe bindings for watching changes to Python's current [`contextvars.Context`][1]. -//! -//! Context watchers are registered for the current Python interpreter and are invoked whenever -//! the current context changes. -//! -//! [1]: https://docs.python.org/3/library/contextvars.html#contextvars.Context - -use crate::err::{error_on_minusone, error_on_minusone_with_result}; -use crate::ffi_ptr_ext::FfiPtrExt; -use crate::types::any::PyAnyMethods; -use crate::types::PyContext; -use crate::{ffi, Borrowed, PyAny, PyResult, Python}; -use core::ffi::c_int; - -/// An event passed to a context watcher. -/// -/// This enum is non-exhaustive because CPython may add context watcher events in future versions. -#[doc(alias = "PyContextEvent")] -#[derive(Debug)] -#[non_exhaustive] -pub enum ContextEvent<'a, 'py> { - /// The current context changed. - /// - /// The value is the new current context, or `None` when there is no current context. - Switched(Option>), - - /// An event which is not known to this version of PyO3. - Unknown { - /// The raw CPython event value. - raw_event: ffi::PyContextEvent, - - /// The event-specific object, if one was provided. - object: Option>, - }, -} - -/// A guard which keeps a context watcher registered. -/// -/// The watcher is registered for the current Python interpreter and is removed when this guard is -/// dropped. Use [`clear`][Self::clear] to remove it explicitly and observe any error returned by -/// CPython. -/// -/// This guard is bound to the [`Python`] attachment used to create it. It therefore cannot be sent -/// to another thread, moved outside that attachment, or moved into [`Python::detach`]. -/// -/// If this guard is forgotten, the watcher remains registered. This does not create a dangling -/// function pointer because [`register_context_watcher!`](crate::register_context_watcher) creates -/// a static, monomorphized trampoline. -#[must_use = "dropping the guard immediately unregisters the context watcher"] -pub struct ContextWatcherGuard<'py> { - watcher_id: c_int, - py: Python<'py>, - active: bool, -} - -impl ContextWatcherGuard<'_> { - /// Removes this watcher from the current Python interpreter. - /// - /// Dropping the guard also removes the watcher, but cannot report a failure to the caller. - #[doc(alias = "PyContext_ClearWatcher")] - pub fn clear(mut self) -> PyResult<()> { - self.active = false; - - // SAFETY: - // - `self.py` proves that the thread is attached to an interpreter; PyO3 does not - // currently support attaching to more than one interpreter, so this is the interpreter - // for which the watcher was registered - // - `watcher_id` was returned by `PyContext_AddWatcher` - error_on_minusone(self.py, unsafe { - ffi::PyContext_ClearWatcher(self.watcher_id) - }) - } -} - -impl Drop for ContextWatcherGuard<'_> { - fn drop(&mut self) { - if !self.active { - return; - } - - self.active = false; - - // A destructor must not replace an exception which was already pending. The Python token - // stored in the guard proves that this thread is still attached to an interpreter; PyO3 - // does not currently support attaching to more than one interpreter, so this is the same - // interpreter the watcher was registered on. - // - // SAFETY: - // - the thread is attached, as guaranteed by `self.py` - // - `PyErr_GetRaisedException` returns an owned reference or NULL - // - `watcher_id` was returned by `PyContext_AddWatcher` - // - `PyErr_SetRaisedException` steals the owned reference returned above - unsafe { - let pending_exception = ffi::PyErr_GetRaisedException(); - if ffi::PyContext_ClearWatcher(self.watcher_id) == -1 { - ffi::PyErr_WriteUnraisable(core::ptr::null_mut()); - } - - if !pending_exception.is_null() { - // Be defensive in case an unraisable hook itself left an exception set. - ffi::PyErr_Clear(); - ffi::PyErr_SetRaisedException(pending_exception); - } - } - } -} - -/// Registers a safe Rust function as a context watcher for the current interpreter. -/// -/// The callback must be a function path and must have this signature: -/// -/// ```rust -/// use pyo3::context::ContextEvent; -/// use pyo3::prelude::*; -/// -/// fn context_changed( -/// _py: Python<'_>, -/// _event: ContextEvent<'_, '_>, -/// ) -> PyResult<()> { -/// Ok(()) -/// } -/// -/// # fn main() -> PyResult<()> { -/// Python::attach(|py| { -/// let _watcher = pyo3::register_context_watcher!(py, context_changed)?; -/// Ok(()) -/// }) -/// # } -/// ``` -/// -/// A function path is required because CPython's context watcher callback has no user-data -/// pointer. The macro creates a unique static trampoline for the function, avoiding global callback -/// storage. State can still be shared through safe static synchronization primitives. -/// -/// Panics and returned [`PyErr`][crate::PyErr] values are reported as unraisable exceptions and -/// never unwind across the C boundary. -#[doc(alias = "PyContext_AddWatcher")] -#[macro_export] -macro_rules! register_context_watcher { - ($py:expr, $callback:path) => {{ - struct Callback; - - impl $crate::context::impl_::ContextWatcherCallbackDef for Callback { - const CALLBACK: $crate::context::impl_::ContextWatcherCallback = $callback; - } - - $crate::context::impl_::register::($py) - }}; -} - -/// Implementation details used by [`register_context_watcher!`](crate::register_context_watcher). -#[doc(hidden)] -pub mod impl_ { - use super::*; - - /// The safe callback signature accepted by context watcher trampolines. - pub type ContextWatcherCallback = - for<'a, 'py> fn(Python<'py>, ContextEvent<'a, 'py>) -> PyResult<()>; - - /// Associates a generated trampoline type with its Rust callback. - pub trait ContextWatcherCallbackDef { - /// The Rust callback invoked by the generated C trampoline. - const CALLBACK: ContextWatcherCallback; - } - - /// Registers the trampoline specialized for `Callback`. - pub fn register( - py: Python<'_>, - ) -> PyResult> { - // SAFETY: - // - `py` proves that the thread is attached - // - `context_watcher::` is a static C-compatible function - let watcher_id = error_on_minusone_with_result(py, unsafe { - ffi::PyContext_AddWatcher(context_watcher::) - })?; - - Ok(ContextWatcherGuard { - watcher_id, - py, - active: true, - }) - } - - /// C-compatible trampoline for a context watcher callback. - /// - /// # Safety - /// - /// - The thread must be attached to Python. - /// - `object` must follow the contract for the supplied `event`. - pub unsafe extern "C" fn context_watcher( - event: ffi::PyContextEvent, - object: *mut ffi::PyObject, - ) -> c_int { - // A context watcher may be called with an exception already set. Save it before invoking - // arbitrary Rust code so that safe PyO3 APIs can be used normally inside the callback. - // - // SAFETY: the caller guarantees that the thread is attached. - let pending_exception = unsafe { ffi::PyErr_GetRaisedException() }; - - // SAFETY: - // - the caller guarantees that the thread is attached and `object` follows the contract - // for `event` - // - `trampoline` catches panics and converts callback errors into a Python exception - // - the callback's higher-ranked signature prevents borrowed event data from escaping - let result = unsafe { - crate::impl_::trampoline::trampoline(|py| { - let event = match event { - ffi::Py_CONTEXT_SWITCHED => { - let object = object.assume_borrowed(py); - - if object.is_none() { - ContextEvent::Switched(None) - } else { - ContextEvent::Switched(Some(object.cast_unchecked())) - } - } - - raw_event => { - let object = object.assume_borrowed_or_opt(py); - ContextEvent::Unknown { raw_event, object } - } - }; - - (Callback::CALLBACK)(py, event)?; - - // Although normal PyO3 APIs return errors as `PyResult`, `PyErr::restore` can be - // called directly. Do not allow an Ok return with an exception still set. - if crate::PyErr::occurred(py) { - return Err(crate::PyErr::fetch(py)); - } - - Ok(0) - }) - }; - - if pending_exception.is_null() { - return result; - } - - // When an exception was already pending on entry, CPython requires the callback to return - // 0 with that same exception still set. Report a new callback error ourselves before - // restoring the original exception. - // - // SAFETY: - // - the thread is attached - // - `object` is valid for the duration of the callback or NULL - // - `pending_exception` is an owned reference from `PyErr_GetRaisedException` - // - `PyErr_SetRaisedException` steals that reference - unsafe { - if result == -1 { - ffi::PyErr_WriteUnraisable(object); - } - - // Be defensive in case an unraisable hook itself left an exception set. - ffi::PyErr_Clear(); - ffi::PyErr_SetRaisedException(pending_exception); - } - - 0 - } -} - -#[cfg(test)] -mod tests { - use super::impl_::{context_watcher, ContextWatcherCallback, ContextWatcherCallbackDef}; - use super::ContextEvent; - use crate::exceptions::{PyRuntimeError, PyValueError}; - #[cfg(feature = "macros")] - use crate::test_utils::UnraisableCapture; - use crate::types::{PyAnyMethods, PyContext}; - use crate::{ffi, PyErr, PyResult, Python}; - #[cfg(feature = "macros")] - use alloc::string::ToString; - use core::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; - use static_assertions::assert_not_impl_any; - - static SWITCH_COUNT: AtomicUsize = AtomicUsize::new(0); - static SAW_CONTEXT: AtomicBool = AtomicBool::new(false); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_switch(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if let ContextEvent::Switched(context) = event { - SWITCH_COUNT.fetch_add(1, Ordering::Relaxed); - if let Some(context) = context { - assert!(context.is_exact_instance_of::()); - SAW_CONTEXT.store(true, Ordering::Relaxed); - } - } - Ok(()) - } - - #[test] - fn watcher_is_cleared_on_drop() { - Python::attach(|py| { - SWITCH_COUNT.store(0, Ordering::Relaxed); - SAW_CONTEXT.store(false, Ordering::Relaxed); - - let watcher = crate::register_context_watcher!(py, record_switch).unwrap(); - py.run( - c"import contextvars\ncontextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - - let count_after_first_run = SWITCH_COUNT.load(Ordering::Relaxed); - assert!(count_after_first_run >= 2); - assert!(SAW_CONTEXT.load(Ordering::Relaxed)); - - drop(watcher); - - py.run(c"contextvars.Context().run(lambda: None)", None, None) - .unwrap(); - assert_eq!(SWITCH_COUNT.load(Ordering::Relaxed), count_after_first_run); - }); - } - - static EXPLICIT_CLEAR_COUNT: AtomicUsize = AtomicUsize::new(0); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_explicit_clear(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if matches!(event, ContextEvent::Switched(_)) { - EXPLICIT_CLEAR_COUNT.fetch_add(1, Ordering::Relaxed); - } - Ok(()) - } - - #[test] - fn watcher_can_be_cleared_explicitly() { - Python::attach(|py| { - EXPLICIT_CLEAR_COUNT.store(0, Ordering::Relaxed); - - let watcher = crate::register_context_watcher!(py, record_explicit_clear).unwrap(); - watcher.clear().unwrap(); - - py.run( - c"import contextvars\ncontextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - assert_eq!(EXPLICIT_CLEAR_COUNT.load(Ordering::Relaxed), 0); - }); - } - - static DUPLICATE_CALLBACK_COUNT: AtomicUsize = AtomicUsize::new(0); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_duplicate_callback(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if matches!(event, ContextEvent::Switched(_)) { - DUPLICATE_CALLBACK_COUNT.fetch_add(1, Ordering::Relaxed); - } - Ok(()) - } - - #[test] - fn multiple_watchers_can_register_the_same_callback() { - Python::attach(|py| { - DUPLICATE_CALLBACK_COUNT.store(0, Ordering::Relaxed); - - let first = crate::register_context_watcher!(py, record_duplicate_callback).unwrap(); - let second = crate::register_context_watcher!(py, record_duplicate_callback).unwrap(); - - py.run( - c"import contextvars\ncontextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - assert!(DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed) >= 4); - - drop(first); - let count_with_both = DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed); - py.run(c"contextvars.Context().run(lambda: None)", None, None) - .unwrap(); - assert!(DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed) >= count_with_both + 2); - - drop(second); - let count_after_drop = DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed); - py.run(c"contextvars.Context().run(lambda: None)", None, None) - .unwrap(); - assert_eq!( - DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed), - count_after_drop - ); - }); - } - - #[test] - fn dropping_watcher_preserves_a_pending_exception() { - Python::attach(|py| { - let watcher = crate::register_context_watcher!(py, record_explicit_clear).unwrap(); - PyValueError::new_err("original error").restore(py); - - drop(watcher); - - let error = PyErr::fetch(py); - assert!(error.is_instance_of::(py)); - }); - } - - fn fail_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { - Err(PyRuntimeError::new_err("watcher failed")) - } - - struct FailingCallback; - - impl ContextWatcherCallbackDef for FailingCallback { - const CALLBACK: ContextWatcherCallback = fail_callback; - } - - #[test] - fn callback_error_is_returned_without_a_pending_exception() { - Python::attach(|py| { - // SAFETY: the thread is attached and None is valid for Py_CONTEXT_SWITCHED. - let result = unsafe { - context_watcher::(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None()) - }; - - assert_eq!(result, -1); - let error = PyErr::fetch(py); - assert!(error.is_instance_of::(py)); - }); - } - - #[test] - #[cfg(feature = "macros")] - fn callback_error_preserves_a_pending_exception() { - Python::attach(|py| { - UnraisableCapture::enter(py, |capture| { - PyValueError::new_err("original error").restore(py); - - // SAFETY: the thread is attached and None is valid for Py_CONTEXT_SWITCHED. - let result = unsafe { - context_watcher::(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None()) - }; - - assert_eq!(result, 0); - - let original_error = PyErr::fetch(py); - assert!(original_error.is_instance_of::(py)); - assert_eq!(original_error.to_string(), "ValueError: original error"); - - let (watcher_error, object) = - capture.take_capture().expect("missing unraisable error"); - assert!(watcher_error.is_instance_of::(py)); - assert!(object.is_none()); - }); - }); - } - - #[test] - #[cfg(feature = "macros")] - fn registered_callback_errors_are_unraisable() { - Python::attach(|py| { - UnraisableCapture::enter(py, |capture| { - let watcher = crate::register_context_watcher!(py, fail_callback).unwrap(); - - py.run( - c"import contextvars\ncontextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - - let (watcher_error, _) = capture.take_capture().expect("missing unraisable error"); - assert!(watcher_error.is_instance_of::(py)); - - drop(watcher); - }); - }); - } - - fn panic_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { - panic!("context watcher panic") - } - - struct PanickingCallback; - - impl ContextWatcherCallbackDef for PanickingCallback { - const CALLBACK: ContextWatcherCallback = panic_callback; - } - - #[test] - fn callback_panic_does_not_cross_ffi_boundary() { - Python::attach(|py| { - // SAFETY: the thread is attached and None is valid for Py_CONTEXT_SWITCHED. - let result = unsafe { - context_watcher::(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None()) - }; - - assert_eq!(result, -1); - assert!(PyErr::occurred(py)); - - // SAFETY: the test has observed and intentionally discards the panic exception. - unsafe { ffi::PyErr_Clear() }; - }); - } - - static UNKNOWN_EVENT: AtomicU32 = AtomicU32::new(0); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_unknown(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if let ContextEvent::Unknown { raw_event, object } = event { - UNKNOWN_EVENT.store(raw_event, Ordering::Relaxed); - assert!(object.is_none()); - } - Ok(()) - } - - struct UnknownCallback; - - impl ContextWatcherCallbackDef for UnknownCallback { - const CALLBACK: ContextWatcherCallback = record_unknown; - } - - #[test] - fn unknown_events_are_forwarded() { - const FUTURE_EVENT: ffi::PyContextEvent = 123; - - Python::attach(|_py| { - UNKNOWN_EVENT.store(0, Ordering::Relaxed); - - // SAFETY: the thread is attached and unknown events accept a null object. - let result = - unsafe { context_watcher::(FUTURE_EVENT, core::ptr::null_mut()) }; - - assert_eq!(result, 0); - assert_eq!(UNKNOWN_EVENT.load(Ordering::Relaxed), FUTURE_EVENT); - }); - } - - #[test] - fn watcher_guard_is_not_send_or_sync() { - assert_not_impl_any!(super::ContextWatcherGuard<'_>: Send, Sync); - } - - static FORGOTTEN_GUARD_COUNT: AtomicUsize = AtomicUsize::new(0); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_forgotten_guard(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if matches!(event, ContextEvent::Switched(_)) { - FORGOTTEN_GUARD_COUNT.fetch_add(1, Ordering::Relaxed); - } - Ok(()) - } - - #[test] - fn forgotten_guard() { - Python::attach(|py| { - FORGOTTEN_GUARD_COUNT.store(0, Ordering::Relaxed); - let watcher = crate::register_context_watcher!(py, record_forgotten_guard).unwrap(); - core::mem::forget(watcher); - - py.run( - c"import contextvars; contextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - - assert!(FORGOTTEN_GUARD_COUNT.load(Ordering::Relaxed) > 0); - }); - } -} diff --git a/src/lib.rs b/src/lib.rs index a50666397d5..8b0eeabc28e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -422,13 +422,6 @@ pub mod impl_; pub mod buffer; pub mod call; -#[cfg(all( - Py_3_14, - not(Py_GIL_DISABLED), - not(Py_LIMITED_API), - not(any(PyPy, GraalPy, RustPython)) -))] -pub mod context; pub mod conversion; mod conversions; #[cfg(feature = "experimental-async")] diff --git a/src/types/context.rs b/src/types/context.rs index cb5b3e96054..ce74bf5e623 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -1,7 +1,22 @@ #![deny(clippy::undocumented_unsafe_blocks)] +//! Types and APIs for Python [`contextvars.Context`][1] objects. +//! +//! On GIL-enabled Python 3.14 and newer, this module also provides safe bindings for watching +//! changes to the current context. +//! +//! [1]: https://docs.python.org/3/library/contextvars.html#contextvars.Context + use crate::{ffi, PyAny}; +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +use crate::{ + err::{error_on_minusone, error_on_minusone_with_result}, + Borrowed, PyResult, Python, +}; +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +use core::ffi::c_int; + /// Represents a Python [`contextvars.Context`][1] object. /// /// Values of this type are accessed via PyO3's smart pointers, e.g. as @@ -20,11 +35,281 @@ pyobject_native_type_core!( #checkfunction=ffi::PyContext_CheckExact ); -#[cfg(test)] +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +impl PyContext { + /// Registers a context watcher for the current interpreter. + /// + /// Use [`watch_callback!`] to create the callback passed to this method. The returned + /// [`ContextWatcherGuard`] removes the watcher when dropped. + /// + /// Panics and returned [`PyErr`][crate::PyErr] values are reported as unraisable exceptions and + /// never unwind across the C boundary. + #[doc(alias = "PyContext_AddWatcher")] + pub fn add_watcher( + py: Python<'_>, + callback: WatchCallback, + ) -> PyResult> { + // SAFETY: + // - `py` proves that the thread is attached + // - `callback` contains a static C-compatible function + let watcher_id = + error_on_minusone_with_result(py, unsafe { ffi::PyContext_AddWatcher(callback.0) })?; + + Ok(ContextWatcherGuard { + watcher_id, + py, + active: true, + }) + } +} + +/// An event passed to a context watcher. +/// +/// This enum is non-exhaustive because CPython may add context watcher events in future versions. +#[doc(alias = "PyContextEvent")] +#[derive(Debug)] +#[non_exhaustive] +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +pub enum ContextEvent<'a, 'py> { + /// The current context changed. + /// + /// The value is the new current context, or `None` when there is no current context. + Switched(Option>), + + /// An event which is not known to this version of PyO3. + Unknown { + /// The raw CPython event value. + raw_event: ffi::PyContextEvent, + + /// The event-specific object, if one was provided. + object: Option>, + }, +} + +/// A guard which keeps a context watcher registered. +/// +/// The watcher is registered for the current Python interpreter and is removed when this guard is +/// dropped. Use [`clear`][Self::clear] to remove it explicitly and observe any error returned by +/// CPython. +/// +/// This guard is bound to the [`Python`] attachment used to create it. It therefore cannot be sent +/// to another thread, moved outside that attachment, or moved into [`Python::detach`]. +/// +/// If this guard is forgotten, the watcher remains registered. This does not create a dangling +/// function pointer because [`watch_callback!`] creates a static, monomorphized trampoline. +#[must_use = "dropping the guard immediately unregisters the context watcher"] +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +pub struct ContextWatcherGuard<'py> { + watcher_id: c_int, + py: Python<'py>, + active: bool, +} + +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +impl ContextWatcherGuard<'_> { + /// Removes this watcher from the current Python interpreter. + /// + /// Dropping the guard also removes the watcher, but cannot report a failure to the caller. + #[doc(alias = "PyContext_ClearWatcher")] + pub fn clear(mut self) -> PyResult<()> { + self.active = false; + + // SAFETY: + // - `self.py` proves that the thread is attached to an interpreter; PyO3 does not + // currently support attaching to more than one interpreter, so this is the interpreter + // for which the watcher was registered + // - `watcher_id` was returned by `PyContext_AddWatcher` + error_on_minusone(self.py, unsafe { + ffi::PyContext_ClearWatcher(self.watcher_id) + }) + } +} + +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +impl Drop for ContextWatcherGuard<'_> { + fn drop(&mut self) { + if !self.active { + return; + } + + self.active = false; + + // A destructor must not replace an exception which was already pending. The Python token + // stored in the guard proves that this thread is still attached to an interpreter; PyO3 + // does not currently support attaching to more than one interpreter, so this is the same + // interpreter the watcher was registered on. + // + // SAFETY: + // - the thread is attached, as guaranteed by `self.py` + // - `PyErr_GetRaisedException` returns an owned reference or NULL + // - `watcher_id` was returned by `PyContext_AddWatcher` + // - `PyErr_SetRaisedException` steals the owned reference returned above + unsafe { + let pending_exception = ffi::PyErr_GetRaisedException(); + if ffi::PyContext_ClearWatcher(self.watcher_id) == -1 { + ffi::PyErr_WriteUnraisable(core::ptr::null_mut()); + } + + if !pending_exception.is_null() { + // Be defensive in case an unraisable hook itself left an exception set. + ffi::PyErr_Clear(); + ffi::PyErr_SetRaisedException(pending_exception); + } + } + } +} + +/// Callback type for context watchers. +/// +/// Values of this type are created by [`watch_callback!`]. +#[repr(transparent)] +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +pub struct WatchCallback(ffi::PyContext_WatchCallback); + +/// Creates a context watcher callback from a safe Rust function. +/// +/// A function path is required because CPython's context watcher callback has no user-data +/// pointer. The macro creates a unique static trampoline for the function, avoiding global callback +/// storage. State can still be shared through safe static synchronization primitives. +#[macro_export] +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +macro_rules! watch_callback { + ($callback:path) => {{ + struct Callback; + + impl $crate::types::context::impl_::ContextWatcherCallbackDef for Callback { + const CALLBACK: $crate::types::context::impl_::ContextWatcherCallback = $callback; + } + + $crate::types::context::impl_::new_watch_callback::() + }}; +} + +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +pub use crate::watch_callback; + +/// Implementation details used by [`watch_callback!`]. +#[doc(hidden)] +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +pub mod impl_ { + + use crate::{ffi_ptr_ext::FfiPtrExt, types::PyAnyMethods}; + + use super::*; + + /// The safe callback signature accepted by context watcher trampolines. + pub type ContextWatcherCallback = + for<'a, 'py> fn(Python<'py>, ContextEvent<'a, 'py>) -> PyResult<()>; + + /// Associates a generated trampoline type with its Rust callback. + pub trait ContextWatcherCallbackDef { + /// The Rust callback invoked by the generated C trampoline. + const CALLBACK: ContextWatcherCallback; + } + + pub fn new_watch_callback() -> WatchCallback { + WatchCallback(context_watcher::) + } + + /// C-compatible trampoline for a context watcher callback. + /// + /// # Safety + /// + /// - The thread must be attached to Python. + /// - `object` must follow the contract for the supplied `event`. + pub unsafe extern "C" fn context_watcher( + event: ffi::PyContextEvent, + object: *mut ffi::PyObject, + ) -> c_int { + // A context watcher may be called with an exception already set. Save it before invoking + // arbitrary Rust code so that safe PyO3 APIs can be used normally inside the callback. + // + // SAFETY: the caller guarantees that the thread is attached. + let pending_exception = unsafe { ffi::PyErr_GetRaisedException() }; + + // SAFETY: + // - the caller guarantees that the thread is attached and `object` follows the contract + // for `event` + // - `trampoline` catches panics and converts callback errors into a Python exception + // - the callback's higher-ranked signature prevents borrowed event data from escaping + let result = unsafe { + crate::impl_::trampoline::trampoline(|py| { + let event = match event { + ffi::Py_CONTEXT_SWITCHED => { + let object = object.assume_borrowed(py); + + if object.is_none() { + ContextEvent::Switched(None) + } else { + ContextEvent::Switched(Some(object.cast_unchecked())) + } + } + + raw_event => { + let object = object.assume_borrowed_or_opt(py); + ContextEvent::Unknown { raw_event, object } + } + }; + + (Callback::CALLBACK)(py, event)?; + + // Although normal PyO3 APIs return errors as `PyResult`, `PyErr::restore` can be + // called directly. Do not allow an Ok return with an exception still set. + if crate::PyErr::occurred(py) { + return Err(crate::PyErr::fetch(py)); + } + + Ok(0) + }) + }; + + if pending_exception.is_null() { + return result; + } + + // When an exception was already pending on entry, CPython requires the callback to return + // 0 with that same exception still set. Report a new callback error ourselves before + // restoring the original exception. + // + // SAFETY: + // - the thread is attached + // - `object` is valid for the duration of the callback or NULL + // - `pending_exception` is an owned reference from `PyErr_GetRaisedException` + // - `PyErr_SetRaisedException` steals that reference + unsafe { + if result == -1 { + ffi::PyErr_WriteUnraisable(object); + } + + // Be defensive in case an unraisable hook itself left an exception set. + ffi::PyErr_Clear(); + ffi::PyErr_SetRaisedException(pending_exception); + } + + 0 + } +} + +#[cfg(all(test, Py_3_14, not(Py_GIL_DISABLED)))] mod tests { - use super::PyContext; + use core::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; + + #[cfg(feature = "macros")] + use alloc::string::ToString; + use static_assertions::assert_not_impl_any; + + use crate::exceptions::{PyRuntimeError, PyValueError}; + use crate::types::context::impl_::{ + context_watcher, ContextWatcherCallback, ContextWatcherCallbackDef, + }; + use crate::types::context::ContextEvent; use crate::types::PyAnyMethods; - use crate::Python; + use crate::{PyErr, Python}; + + #[cfg(feature = "macros")] + use crate::test_utils::UnraisableCapture; + + use super::*; #[test] fn context_type() { @@ -41,4 +326,298 @@ mod tests { context.cast::().unwrap(); }); } + + static SWITCH_COUNT: AtomicUsize = AtomicUsize::new(0); + static SAW_CONTEXT: AtomicBool = AtomicBool::new(false); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_switch(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if let ContextEvent::Switched(context) = event { + SWITCH_COUNT.fetch_add(1, Ordering::Relaxed); + if let Some(context) = context { + assert!(context.is_exact_instance_of::()); + SAW_CONTEXT.store(true, Ordering::Relaxed); + } + } + Ok(()) + } + + #[test] + fn watcher_is_cleared_on_drop() { + Python::attach(|py| { + SWITCH_COUNT.store(0, Ordering::Relaxed); + SAW_CONTEXT.store(false, Ordering::Relaxed); + + let watcher = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap(); + py.run( + c"import contextvars\ncontextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + + let count_after_first_run = SWITCH_COUNT.load(Ordering::Relaxed); + assert!(count_after_first_run >= 2); + assert!(SAW_CONTEXT.load(Ordering::Relaxed)); + + drop(watcher); + + py.run(c"contextvars.Context().run(lambda: None)", None, None) + .unwrap(); + assert_eq!(SWITCH_COUNT.load(Ordering::Relaxed), count_after_first_run); + }); + } + + static EXPLICIT_CLEAR_COUNT: AtomicUsize = AtomicUsize::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_explicit_clear(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if matches!(event, ContextEvent::Switched(_)) { + EXPLICIT_CLEAR_COUNT.fetch_add(1, Ordering::Relaxed); + } + Ok(()) + } + + #[test] + fn watcher_can_be_cleared_explicitly() { + Python::attach(|py| { + EXPLICIT_CLEAR_COUNT.store(0, Ordering::Relaxed); + + let watcher = + PyContext::add_watcher(py, watch_callback!(record_explicit_clear)).unwrap(); + watcher.clear().unwrap(); + + py.run( + c"import contextvars\ncontextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + assert_eq!(EXPLICIT_CLEAR_COUNT.load(Ordering::Relaxed), 0); + }); + } + + static DUPLICATE_CALLBACK_COUNT: AtomicUsize = AtomicUsize::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_duplicate_callback(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if matches!(event, ContextEvent::Switched(_)) { + DUPLICATE_CALLBACK_COUNT.fetch_add(1, Ordering::Relaxed); + } + Ok(()) + } + + #[test] + fn multiple_watchers_can_register_the_same_callback() { + Python::attach(|py| { + DUPLICATE_CALLBACK_COUNT.store(0, Ordering::Relaxed); + + let first = + PyContext::add_watcher(py, watch_callback!(record_duplicate_callback)).unwrap(); + let second = + PyContext::add_watcher(py, watch_callback!(record_duplicate_callback)).unwrap(); + + py.run( + c"import contextvars\ncontextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + assert!(DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed) >= 4); + + drop(first); + let count_with_both = DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed); + py.run(c"contextvars.Context().run(lambda: None)", None, None) + .unwrap(); + assert!(DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed) >= count_with_both + 2); + + drop(second); + let count_after_drop = DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed); + py.run(c"contextvars.Context().run(lambda: None)", None, None) + .unwrap(); + assert_eq!( + DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed), + count_after_drop + ); + }); + } + + #[test] + fn dropping_watcher_preserves_a_pending_exception() { + Python::attach(|py| { + let watcher = + PyContext::add_watcher(py, watch_callback!(record_explicit_clear)).unwrap(); + PyValueError::new_err("original error").restore(py); + + drop(watcher); + + let error = PyErr::fetch(py); + assert!(error.is_instance_of::(py)); + }); + } + + fn fail_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { + Err(PyRuntimeError::new_err("watcher failed")) + } + + struct FailingCallback; + + impl ContextWatcherCallbackDef for FailingCallback { + const CALLBACK: ContextWatcherCallback = fail_callback; + } + + #[test] + fn callback_error_is_returned_without_a_pending_exception() { + Python::attach(|py| { + // SAFETY: the thread is attached and None is valid for Py_CONTEXT_SWITCHED. + let result = unsafe { + context_watcher::(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None()) + }; + + assert_eq!(result, -1); + let error = PyErr::fetch(py); + assert!(error.is_instance_of::(py)); + }); + } + + #[test] + #[cfg(feature = "macros")] + fn callback_error_preserves_a_pending_exception() { + Python::attach(|py| { + UnraisableCapture::enter(py, |capture| { + PyValueError::new_err("original error").restore(py); + + // SAFETY: the thread is attached and None is valid for Py_CONTEXT_SWITCHED. + let result = unsafe { + context_watcher::(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None()) + }; + + assert_eq!(result, 0); + + let original_error = PyErr::fetch(py); + assert!(original_error.is_instance_of::(py)); + assert_eq!(original_error.to_string(), "ValueError: original error"); + + let (watcher_error, object) = + capture.take_capture().expect("missing unraisable error"); + assert!(watcher_error.is_instance_of::(py)); + assert!(object.is_none()); + }); + }); + } + + #[test] + #[cfg(feature = "macros")] + fn registered_callback_errors_are_unraisable() { + Python::attach(|py| { + UnraisableCapture::enter(py, |capture| { + let watcher = PyContext::add_watcher(py, watch_callback!(fail_callback)).unwrap(); + + py.run( + c"import contextvars\ncontextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + + let (watcher_error, _) = capture.take_capture().expect("missing unraisable error"); + assert!(watcher_error.is_instance_of::(py)); + + drop(watcher); + }); + }); + } + + fn panic_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { + panic!("context watcher panic") + } + + struct PanickingCallback; + + impl ContextWatcherCallbackDef for PanickingCallback { + const CALLBACK: ContextWatcherCallback = panic_callback; + } + + #[test] + fn callback_panic_does_not_cross_ffi_boundary() { + Python::attach(|py| { + // SAFETY: the thread is attached and None is valid for Py_CONTEXT_SWITCHED. + let result = unsafe { + context_watcher::(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None()) + }; + + assert_eq!(result, -1); + assert!(PyErr::occurred(py)); + + // SAFETY: the test has observed and intentionally discards the panic exception. + unsafe { ffi::PyErr_Clear() }; + }); + } + + static UNKNOWN_EVENT: AtomicU32 = AtomicU32::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_unknown(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if let ContextEvent::Unknown { raw_event, object } = event { + UNKNOWN_EVENT.store(raw_event, Ordering::Relaxed); + assert!(object.is_none()); + } + Ok(()) + } + + struct UnknownCallback; + + impl ContextWatcherCallbackDef for UnknownCallback { + const CALLBACK: ContextWatcherCallback = record_unknown; + } + + #[test] + fn unknown_events_are_forwarded() { + const FUTURE_EVENT: ffi::PyContextEvent = 123; + + Python::attach(|_py| { + UNKNOWN_EVENT.store(0, Ordering::Relaxed); + + // SAFETY: the thread is attached and unknown events accept a null object. + let result = + unsafe { context_watcher::(FUTURE_EVENT, core::ptr::null_mut()) }; + + assert_eq!(result, 0); + assert_eq!(UNKNOWN_EVENT.load(Ordering::Relaxed), FUTURE_EVENT); + }); + } + + #[test] + fn watcher_guard_is_not_send_or_sync() { + assert_not_impl_any!(super::ContextWatcherGuard<'_>: Send, Sync); + } + + static FORGOTTEN_GUARD_COUNT: AtomicUsize = AtomicUsize::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_forgotten_guard(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if matches!(event, ContextEvent::Switched(_)) { + FORGOTTEN_GUARD_COUNT.fetch_add(1, Ordering::Relaxed); + } + Ok(()) + } + + #[test] + fn forgotten_guard() { + Python::attach(|py| { + FORGOTTEN_GUARD_COUNT.store(0, Ordering::Relaxed); + let watcher = + PyContext::add_watcher(py, watch_callback!(record_forgotten_guard)).unwrap(); + core::mem::forget(watcher); + + py.run( + c"import contextvars; contextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + + assert!(FORGOTTEN_GUARD_COUNT.load(Ordering::Relaxed) > 0); + }); + } } diff --git a/src/types/mod.rs b/src/types/mod.rs index a6c956334ed..4d5885e26ec 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -358,7 +358,7 @@ pub mod capsule; pub mod code; pub(crate) mod complex; #[cfg(not(any(Py_LIMITED_API, PyPy, GraalPy, RustPython)))] -mod context; +pub mod context; pub mod datetime; pub mod dict; mod ellipsis; From 7926e546086eb9e736dd9c2d1d7fe50f7827b556 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 27 Aug 2026 13:35:34 +0200 Subject: [PATCH 11/24] Use checked context watcher event conversion --- src/types/context.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/types/context.rs b/src/types/context.rs index ce74bf5e623..21873493b7d 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -222,7 +222,9 @@ pub mod impl_ { object: *mut ffi::PyObject, ) -> c_int { // A context watcher may be called with an exception already set. Save it before invoking - // arbitrary Rust code so that safe PyO3 APIs can be used normally inside the callback. + // arbitrary Rust code so that safe PyO3 APIs can be used normally inside the callback. The + // raw exception API is intentional because `PyErr::take` may resume a `PanicException`, + // which must not unwind across this C boundary. // // SAFETY: the caller guarantees that the thread is attached. let pending_exception = unsafe { ffi::PyErr_GetRaisedException() }; @@ -236,12 +238,12 @@ pub mod impl_ { crate::impl_::trampoline::trampoline(|py| { let event = match event { ffi::Py_CONTEXT_SWITCHED => { - let object = object.assume_borrowed(py); + let object = object.assume_borrowed_or_err(py)?; if object.is_none() { ContextEvent::Switched(None) } else { - ContextEvent::Switched(Some(object.cast_unchecked())) + ContextEvent::Switched(Some(object.cast()?)) } } From b09e9496ddd047a6a00f56c5be2c5519fce83e5f Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 27 Aug 2026 13:36:13 +0200 Subject: [PATCH 12/24] Remove watcher-specific trampoline error check --- src/types/context.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/types/context.rs b/src/types/context.rs index 21873493b7d..af78fc5e340 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -255,12 +255,6 @@ pub mod impl_ { (Callback::CALLBACK)(py, event)?; - // Although normal PyO3 APIs return errors as `PyResult`, `PyErr::restore` can be - // called directly. Do not allow an Ok return with an exception still set. - if crate::PyErr::occurred(py) { - return Err(crate::PyErr::fetch(py)); - } - Ok(0) }) }; From ac88a33343e15202b23b95023727703c8470dcea Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 27 Aug 2026 13:37:10 +0200 Subject: [PATCH 13/24] Support unbound context watcher guards --- newsfragments/6227.added.md | 6 +- src/types/context.rs | 388 +++++++++++++++++++++++++++++------- 2 files changed, 318 insertions(+), 76 deletions(-) diff --git a/newsfragments/6227.added.md b/newsfragments/6227.added.md index e4b1d0635ec..f32e9de12be 100644 --- a/newsfragments/6227.added.md +++ b/newsfragments/6227.added.md @@ -1,3 +1,3 @@ -Add `PyContext`, as well as `register_context_watcher!`, `ContextEvent`, and -`ContextWatcherGuard` for watching `contextvars.Context` on GIL-enabled CPython -3.14+. `PyContext` itself is available on all supported CPython versions. +Add `PyContext`, as well as `PyContext::add_watcher`, `watch_callback!`, `ContextEvent`, +`BoundContextWatcherGuard`, and `ContextWatcherGuard` for watching `contextvars.Context` on +GIL-enabled CPython 3.14+. `PyContext` itself is available on all supported CPython versions. diff --git a/src/types/context.rs b/src/types/context.rs index af78fc5e340..89d2f0f2733 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -39,8 +39,10 @@ pyobject_native_type_core!( impl PyContext { /// Registers a context watcher for the current interpreter. /// - /// Use [`watch_callback!`] to create the callback passed to this method. The returned - /// [`ContextWatcherGuard`] removes the watcher when dropped. + /// Use [`watch_callback!`] to create the callback passed to this method. + /// The returned [`BoundContextWatcherGuard`] removes the watcher when dropped; call + /// [`BoundContextWatcherGuard::unbind`] if the watcher needs to outlive the current Python + /// attachment. /// /// Panics and returned [`PyErr`][crate::PyErr] values are reported as unraisable exceptions and /// never unwind across the C boundary. @@ -48,14 +50,14 @@ impl PyContext { pub fn add_watcher( py: Python<'_>, callback: WatchCallback, - ) -> PyResult> { + ) -> PyResult> { // SAFETY: // - `py` proves that the thread is attached // - `callback` contains a static C-compatible function let watcher_id = error_on_minusone_with_result(py, unsafe { ffi::PyContext_AddWatcher(callback.0) })?; - Ok(ContextWatcherGuard { + Ok(BoundContextWatcherGuard { watcher_id, py, active: true, @@ -86,7 +88,7 @@ pub enum ContextEvent<'a, 'py> { }, } -/// A guard which keeps a context watcher registered. +/// A Python-bound guard which keeps a context watcher registered. /// /// The watcher is registered for the current Python interpreter and is removed when this guard is /// dropped. Use [`clear`][Self::clear] to remove it explicitly and observe any error returned by @@ -95,66 +97,167 @@ pub enum ContextEvent<'a, 'py> { /// This guard is bound to the [`Python`] attachment used to create it. It therefore cannot be sent /// to another thread, moved outside that attachment, or moved into [`Python::detach`]. /// -/// If this guard is forgotten, the watcher remains registered. This does not create a dangling -/// function pointer because [`watch_callback!`] creates a static, monomorphized trampoline. +/// Use [`unbind`][Self::unbind] to convert this guard into a [`ContextWatcherGuard`] which can be +/// stored outside the current attachment. #[must_use = "dropping the guard immediately unregisters the context watcher"] #[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] -pub struct ContextWatcherGuard<'py> { +pub struct BoundContextWatcherGuard<'py> { watcher_id: c_int, py: Python<'py>, active: bool, } #[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] -impl ContextWatcherGuard<'_> { +impl BoundContextWatcherGuard<'_> { /// Removes this watcher from the current Python interpreter. /// /// Dropping the guard also removes the watcher, but cannot report a failure to the caller. #[doc(alias = "PyContext_ClearWatcher")] pub fn clear(mut self) -> PyResult<()> { self.active = false; + clear_watcher(self.py, self.watcher_id) + } - // SAFETY: - // - `self.py` proves that the thread is attached to an interpreter; PyO3 does not - // currently support attaching to more than one interpreter, so this is the interpreter - // for which the watcher was registered - // - `watcher_id` was returned by `PyContext_AddWatcher` - error_on_minusone(self.py, unsafe { - ffi::PyContext_ClearWatcher(self.watcher_id) - }) + /// Removes the connection to the current Python attachment, allowing the guard to be stored + /// outside it or sent to another thread. + /// + /// Dropping the returned guard automatically attaches to Python to remove the watcher. To avoid + /// that attachment, convert it back with [`ContextWatcherGuard::into_bound`] before dropping it, + /// or call [`ContextWatcherGuard::clear`] while attached. + pub fn unbind(mut self) -> ContextWatcherGuard { + self.active = false; + ContextWatcherGuard { + watcher_id: self.watcher_id, + active: true, + } } } #[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] -impl Drop for ContextWatcherGuard<'_> { +impl Drop for BoundContextWatcherGuard<'_> { fn drop(&mut self) { if !self.active { return; } self.active = false; + clear_watcher_on_drop(self.py, self.watcher_id); + } +} - // A destructor must not replace an exception which was already pending. The Python token - // stored in the guard proves that this thread is still attached to an interpreter; PyO3 - // does not currently support attaching to more than one interpreter, so this is the same - // interpreter the watcher was registered on. - // - // SAFETY: - // - the thread is attached, as guaranteed by `self.py` - // - `PyErr_GetRaisedException` returns an owned reference or NULL - // - `watcher_id` was returned by `PyContext_AddWatcher` - // - `PyErr_SetRaisedException` steals the owned reference returned above - unsafe { - let pending_exception = ffi::PyErr_GetRaisedException(); - if ffi::PyContext_ClearWatcher(self.watcher_id) == -1 { - ffi::PyErr_WriteUnraisable(core::ptr::null_mut()); - } +/// An unbound guard which keeps a context watcher registered. +/// +/// Unlike [`BoundContextWatcherGuard`], this guard is not tied to a particular [`Python`] +/// attachment, so it can be stored outside that attachment or sent to another thread. +/// +/// Dropping this guard automatically attaches to Python to remove the watcher. Use +/// [`clear`][Self::clear] to remove it with an existing attachment and observe any error returned by +/// CPython, or [`into_bound`][Self::into_bound] to recover a bound guard. +/// +/// If Python cannot be attached during drop, the watcher remains registered. This does not create a +/// dangling function pointer because [`watch_callback!`] creates a static, monomorphized trampoline. +/// +/// # Example +/// +/// ```rust +/// use pyo3::prelude::*; +/// use pyo3::types::context::{watch_callback, ContextEvent}; +/// use pyo3::types::PyContext; +/// +/// fn context_changed( +/// _py: Python<'_>, +/// _event: ContextEvent<'_, '_>, +/// ) -> PyResult<()> { +/// Ok(()) +/// } +/// +/// # fn main() -> PyResult<()> { +/// let watcher = Python::attach(|py| -> PyResult<_> { +/// Ok(PyContext::add_watcher(py, watch_callback!(context_changed))?.unbind()) +/// })?; +/// +/// // The guard can be stored until an attachment is available for explicit cleanup. +/// Python::attach(|py| watcher.clear(py)) +/// # } +/// ``` +#[must_use = "dropping the guard unregisters the context watcher"] +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +pub struct ContextWatcherGuard { + watcher_id: c_int, + active: bool, +} - if !pending_exception.is_null() { - // Be defensive in case an unraisable hook itself left an exception set. - ffi::PyErr_Clear(); - ffi::PyErr_SetRaisedException(pending_exception); - } +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +impl ContextWatcherGuard { + /// Removes this watcher from the current Python interpreter. + /// + /// Dropping the guard also removes the watcher, but cannot report a failure to the caller. + #[doc(alias = "PyContext_ClearWatcher")] + pub fn clear(mut self, py: Python<'_>) -> PyResult<()> { + self.active = false; + clear_watcher(py, self.watcher_id) + } + + /// Connects this guard to the given Python attachment. + /// + /// PyO3 does not currently support using a module from multiple interpreters, so `py` is the + /// attachment for the interpreter in which this watcher was registered. + pub fn into_bound<'py>(mut self, py: Python<'py>) -> BoundContextWatcherGuard<'py> { + self.active = false; + BoundContextWatcherGuard { + watcher_id: self.watcher_id, + py, + active: true, + } + } +} + +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +impl Drop for ContextWatcherGuard { + fn drop(&mut self) { + if !self.active { + return; + } + + self.active = false; + let watcher_id = self.watcher_id; + let _ = Python::try_attach(|py| clear_watcher_on_drop(py, watcher_id)); + } +} + +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +fn clear_watcher(py: Python<'_>, watcher_id: c_int) -> PyResult<()> { + // SAFETY: + // - `py` proves that the thread is attached to an interpreter; PyO3 does not currently support + // attaching to more than one interpreter, so this is the interpreter for which the watcher + // was registered + // - `watcher_id` was returned by `PyContext_AddWatcher` + error_on_minusone(py, unsafe { ffi::PyContext_ClearWatcher(watcher_id) }) +} + +#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] +fn clear_watcher_on_drop(_py: Python<'_>, watcher_id: c_int) { + // A destructor must not replace an exception which was already pending. The Python token proves + // that this thread is attached to an interpreter; PyO3 does not currently support attaching to + // more than one interpreter, so this is the same interpreter the watcher was registered on. The + // raw exception API is intentional because `PyErr::take` may resume a `PanicException`, while + // `Drop` must preserve it without unwinding. + // + // SAFETY: + // - the thread is attached, as guaranteed by `_py` + // - `PyErr_GetRaisedException` returns an owned reference or NULL + // - `watcher_id` was returned by `PyContext_AddWatcher` + // - `PyErr_SetRaisedException` steals the owned reference returned above + unsafe { + let pending_exception = ffi::PyErr_GetRaisedException(); + if ffi::PyContext_ClearWatcher(watcher_id) == -1 { + ffi::PyErr_WriteUnraisable(core::ptr::null_mut()); + } + + if !pending_exception.is_null() { + // Be defensive in case an unraisable hook itself left an exception set. + ffi::PyErr_Clear(); + ffi::PyErr_SetRaisedException(pending_exception); } } } @@ -168,9 +271,32 @@ pub struct WatchCallback(ffi::PyContext_WatchCallback); /// Creates a context watcher callback from a safe Rust function. /// +/// The function must be a path with this signature: +/// +/// ```rust +/// use pyo3::prelude::*; +/// use pyo3::types::context::{watch_callback, ContextEvent}; +/// use pyo3::types::PyContext; +/// +/// fn context_changed( +/// _py: Python<'_>, +/// _event: ContextEvent<'_, '_>, +/// ) -> PyResult<()> { +/// Ok(()) +/// } +/// +/// # fn main() -> PyResult<()> { +/// Python::attach(|py| { +/// let _watcher = PyContext::add_watcher(py, watch_callback!(context_changed))?; +/// Ok(()) +/// }) +/// # } +/// ``` +/// /// A function path is required because CPython's context watcher callback has no user-data /// pointer. The macro creates a unique static trampoline for the function, avoiding global callback /// storage. State can still be shared through safe static synchronization primitives. +/// #[macro_export] #[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] macro_rules! watch_callback { @@ -192,7 +318,6 @@ pub use crate::watch_callback; #[doc(hidden)] #[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] pub mod impl_ { - use crate::{ffi_ptr_ext::FfiPtrExt, types::PyAnyMethods}; use super::*; @@ -287,41 +412,18 @@ pub mod impl_ { } #[cfg(all(test, Py_3_14, not(Py_GIL_DISABLED)))] -mod tests { - use core::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; - - #[cfg(feature = "macros")] - use alloc::string::ToString; - use static_assertions::assert_not_impl_any; - +mod watcher_tests { + use super::impl_::{context_watcher, ContextWatcherCallback, ContextWatcherCallbackDef}; + use super::{ContextEvent, PyContext}; use crate::exceptions::{PyRuntimeError, PyValueError}; - use crate::types::context::impl_::{ - context_watcher, ContextWatcherCallback, ContextWatcherCallbackDef, - }; - use crate::types::context::ContextEvent; - use crate::types::PyAnyMethods; - use crate::{PyErr, Python}; - #[cfg(feature = "macros")] use crate::test_utils::UnraisableCapture; - - use super::*; - - #[test] - fn context_type() { - Python::attach(|py| { - let context = py - .import(c"contextvars") - .unwrap() - .getattr(c"Context") - .unwrap() - .call0() - .unwrap(); - - assert!(context.is_exact_instance_of::()); - context.cast::().unwrap(); - }); - } + use crate::types::PyAnyMethods; + use crate::{ffi, PyErr, PyResult, Python}; + #[cfg(feature = "macros")] + use alloc::string::ToString; + use core::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; + use static_assertions::{assert_impl_all, assert_not_impl_any}; static SWITCH_COUNT: AtomicUsize = AtomicUsize::new(0); static SAW_CONTEXT: AtomicBool = AtomicBool::new(false); @@ -584,8 +686,125 @@ mod tests { } #[test] - fn watcher_guard_is_not_send_or_sync() { - assert_not_impl_any!(super::ContextWatcherGuard<'_>: Send, Sync); + fn bound_watcher_guard_is_not_send_or_sync() { + assert_not_impl_any!(super::BoundContextWatcherGuard<'_>: Send, Sync); + } + + #[test] + fn unbound_watcher_guard_is_send_and_sync() { + assert_impl_all!(super::ContextWatcherGuard: Send, Sync); + } + + static UNBOUND_DROP_COUNT: AtomicUsize = AtomicUsize::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_unbound_drop(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if matches!(event, ContextEvent::Switched(_)) { + UNBOUND_DROP_COUNT.fetch_add(1, Ordering::Relaxed); + } + Ok(()) + } + + #[test] + fn unbound_watcher_outlives_attachment_and_attaches_on_drop() { + UNBOUND_DROP_COUNT.store(0, Ordering::Relaxed); + let watcher = Python::attach(|py| { + PyContext::add_watcher(py, watch_callback!(record_unbound_drop)) + .unwrap() + .unbind() + }); + + Python::attach(|py| { + py.run( + c"import contextvars; contextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + }); + let count_before_drop = UNBOUND_DROP_COUNT.load(Ordering::Relaxed); + assert!(count_before_drop >= 2); + + drop(watcher); + + Python::attach(|py| { + py.run(c"contextvars.Context().run(lambda: None)", None, None) + .unwrap(); + }); + assert_eq!( + UNBOUND_DROP_COUNT.load(Ordering::Relaxed), + count_before_drop + ); + } + + static UNBOUND_CLEAR_COUNT: AtomicUsize = AtomicUsize::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_unbound_clear(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if matches!(event, ContextEvent::Switched(_)) { + UNBOUND_CLEAR_COUNT.fetch_add(1, Ordering::Relaxed); + } + Ok(()) + } + + #[test] + fn unbound_watcher_can_be_cleared_with_an_attachment() { + UNBOUND_CLEAR_COUNT.store(0, Ordering::Relaxed); + let watcher = Python::attach(|py| { + PyContext::add_watcher(py, watch_callback!(record_unbound_clear)) + .unwrap() + .unbind() + }); + + Python::attach(|py| watcher.clear(py).unwrap()); + + Python::attach(|py| { + py.run( + c"import contextvars; contextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + }); + assert_eq!(UNBOUND_CLEAR_COUNT.load(Ordering::Relaxed), 0); + } + + static REBOUND_COUNT: AtomicUsize = AtomicUsize::new(0); + + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn record_rebound(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { + if matches!(event, ContextEvent::Switched(_)) { + REBOUND_COUNT.fetch_add(1, Ordering::Relaxed); + } + Ok(()) + } + + #[test] + fn unbound_watcher_can_be_rebound() { + REBOUND_COUNT.store(0, Ordering::Relaxed); + let watcher = Python::attach(|py| { + PyContext::add_watcher(py, watch_callback!(record_rebound)) + .unwrap() + .unbind() + }); + + Python::attach(|py| { + let watcher = watcher.into_bound(py); + py.run( + c"import contextvars; contextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + + let count_before_drop = REBOUND_COUNT.load(Ordering::Relaxed); + assert!(count_before_drop >= 2); + drop(watcher); + + py.run(c"contextvars.Context().run(lambda: None)", None, None) + .unwrap(); + assert_eq!(REBOUND_COUNT.load(Ordering::Relaxed), count_before_drop); + }); } static FORGOTTEN_GUARD_COUNT: AtomicUsize = AtomicUsize::new(0); @@ -617,3 +836,26 @@ mod tests { }); } } + +#[cfg(test)] +mod tests { + use super::PyContext; + use crate::types::PyAnyMethods; + use crate::Python; + + #[test] + fn context_type() { + Python::attach(|py| { + let context = py + .import(c"contextvars") + .unwrap() + .getattr(c"Context") + .unwrap() + .call0() + .unwrap(); + + assert!(context.is_exact_instance_of::()); + context.cast::().unwrap(); + }); + } +} From c5636e50107b91a51c1de166992ea6c5d6c92ad2 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 27 Aug 2026 14:41:19 +0200 Subject: [PATCH 14/24] Fix race in unbound context watcher test --- src/types/context.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/types/context.rs b/src/types/context.rs index 89d2f0f2733..aca58e369d0 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -722,19 +722,16 @@ mod watcher_tests { ) .unwrap(); }); - let count_before_drop = UNBOUND_DROP_COUNT.load(Ordering::Relaxed); - assert!(count_before_drop >= 2); + assert!(UNBOUND_DROP_COUNT.load(Ordering::Relaxed) >= 2); drop(watcher); + let count_after_drop = UNBOUND_DROP_COUNT.load(Ordering::Relaxed); Python::attach(|py| { py.run(c"contextvars.Context().run(lambda: None)", None, None) .unwrap(); }); - assert_eq!( - UNBOUND_DROP_COUNT.load(Ordering::Relaxed), - count_before_drop - ); + assert_eq!(UNBOUND_DROP_COUNT.load(Ordering::Relaxed), count_after_drop); } static UNBOUND_CLEAR_COUNT: AtomicUsize = AtomicUsize::new(0); From c0e49d72b4297ed8cbdf04570a24d4d5bdf12d7e Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 27 Aug 2026 15:05:24 +0200 Subject: [PATCH 15/24] Fix race in unbound context watcher clear test --- src/types/context.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/types/context.rs b/src/types/context.rs index aca58e369d0..4be5053560b 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -753,7 +753,10 @@ mod watcher_tests { .unbind() }); - Python::attach(|py| watcher.clear(py).unwrap()); + let count_after_clear = Python::attach(|py| { + watcher.clear(py).unwrap(); + UNBOUND_CLEAR_COUNT.load(Ordering::Relaxed) + }); Python::attach(|py| { py.run( @@ -763,7 +766,10 @@ mod watcher_tests { ) .unwrap(); }); - assert_eq!(UNBOUND_CLEAR_COUNT.load(Ordering::Relaxed), 0); + assert_eq!( + UNBOUND_CLEAR_COUNT.load(Ordering::Relaxed), + count_after_clear + ); } static REBOUND_COUNT: AtomicUsize = AtomicUsize::new(0); From 8edd88a289fd115e63018ef4f038affe8cd77952 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Thu, 27 Aug 2026 16:18:00 +0200 Subject: [PATCH 16/24] Refactor context watcher tests --- src/types/context.rs | 249 ++++++++++++------------------------------- 1 file changed, 68 insertions(+), 181 deletions(-) diff --git a/src/types/context.rs b/src/types/context.rs index 4be5053560b..90e2e6d6ade 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -424,10 +424,34 @@ mod watcher_tests { use alloc::string::ToString; use core::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; use static_assertions::{assert_impl_all, assert_not_impl_any}; + use std::sync::{Mutex, MutexGuard, PoisonError}; + // Context watchers are interpreter-global and limited to eight slots, so tests which register + // watchers must not run concurrently. + static WATCHER_TEST_MUTEX: Mutex<()> = Mutex::new(()); static SWITCH_COUNT: AtomicUsize = AtomicUsize::new(0); static SAW_CONTEXT: AtomicBool = AtomicBool::new(false); + fn acquire_watcher_test_lock() -> MutexGuard<'static, ()> { + WATCHER_TEST_MUTEX + .lock() + .unwrap_or_else(PoisonError::into_inner) + } + + fn run_context_switch(py: Python<'_>) { + py.run( + c"import contextvars; contextvars.Context().run(lambda: None)", + None, + None, + ) + .unwrap(); + } + + fn assert_no_context_switches(py: Python<'_>, count_before: usize) { + run_context_switch(py); + assert_eq!(SWITCH_COUNT.load(Ordering::Relaxed), count_before); + } + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] fn record_switch(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { if let ContextEvent::Switched(context) = event { @@ -442,17 +466,13 @@ mod watcher_tests { #[test] fn watcher_is_cleared_on_drop() { + let _guard = acquire_watcher_test_lock(); Python::attach(|py| { SWITCH_COUNT.store(0, Ordering::Relaxed); SAW_CONTEXT.store(false, Ordering::Relaxed); let watcher = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap(); - py.run( - c"import contextvars\ncontextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); + run_context_switch(py); let count_after_first_run = SWITCH_COUNT.load(Ordering::Relaxed); assert!(count_after_first_run >= 2); @@ -460,91 +480,51 @@ mod watcher_tests { drop(watcher); - py.run(c"contextvars.Context().run(lambda: None)", None, None) - .unwrap(); - assert_eq!(SWITCH_COUNT.load(Ordering::Relaxed), count_after_first_run); + assert_no_context_switches(py, count_after_first_run); }); } - static EXPLICIT_CLEAR_COUNT: AtomicUsize = AtomicUsize::new(0); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_explicit_clear(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if matches!(event, ContextEvent::Switched(_)) { - EXPLICIT_CLEAR_COUNT.fetch_add(1, Ordering::Relaxed); - } - Ok(()) - } - #[test] fn watcher_can_be_cleared_explicitly() { + let _guard = acquire_watcher_test_lock(); Python::attach(|py| { - EXPLICIT_CLEAR_COUNT.store(0, Ordering::Relaxed); + SWITCH_COUNT.store(0, Ordering::Relaxed); - let watcher = - PyContext::add_watcher(py, watch_callback!(record_explicit_clear)).unwrap(); + let watcher = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap(); watcher.clear().unwrap(); - py.run( - c"import contextvars\ncontextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - assert_eq!(EXPLICIT_CLEAR_COUNT.load(Ordering::Relaxed), 0); + assert_no_context_switches(py, 0); }); } - static DUPLICATE_CALLBACK_COUNT: AtomicUsize = AtomicUsize::new(0); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_duplicate_callback(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if matches!(event, ContextEvent::Switched(_)) { - DUPLICATE_CALLBACK_COUNT.fetch_add(1, Ordering::Relaxed); - } - Ok(()) - } - #[test] fn multiple_watchers_can_register_the_same_callback() { + let _guard = acquire_watcher_test_lock(); Python::attach(|py| { - DUPLICATE_CALLBACK_COUNT.store(0, Ordering::Relaxed); + SWITCH_COUNT.store(0, Ordering::Relaxed); - let first = - PyContext::add_watcher(py, watch_callback!(record_duplicate_callback)).unwrap(); - let second = - PyContext::add_watcher(py, watch_callback!(record_duplicate_callback)).unwrap(); + let first = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap(); + let second = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap(); - py.run( - c"import contextvars\ncontextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - assert!(DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed) >= 4); + run_context_switch(py); + assert!(SWITCH_COUNT.load(Ordering::Relaxed) >= 4); drop(first); - let count_with_both = DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed); - py.run(c"contextvars.Context().run(lambda: None)", None, None) - .unwrap(); - assert!(DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed) >= count_with_both + 2); + let count_with_both = SWITCH_COUNT.load(Ordering::Relaxed); + run_context_switch(py); + assert!(SWITCH_COUNT.load(Ordering::Relaxed) >= count_with_both + 2); drop(second); - let count_after_drop = DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed); - py.run(c"contextvars.Context().run(lambda: None)", None, None) - .unwrap(); - assert_eq!( - DUPLICATE_CALLBACK_COUNT.load(Ordering::Relaxed), - count_after_drop - ); + let count_after_drop = SWITCH_COUNT.load(Ordering::Relaxed); + assert_no_context_switches(py, count_after_drop); }); } #[test] fn dropping_watcher_preserves_a_pending_exception() { + let _guard = acquire_watcher_test_lock(); Python::attach(|py| { - let watcher = - PyContext::add_watcher(py, watch_callback!(record_explicit_clear)).unwrap(); + let watcher = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap(); PyValueError::new_err("original error").restore(py); drop(watcher); @@ -607,16 +587,12 @@ mod watcher_tests { #[test] #[cfg(feature = "macros")] fn registered_callback_errors_are_unraisable() { + let _guard = acquire_watcher_test_lock(); Python::attach(|py| { UnraisableCapture::enter(py, |capture| { let watcher = PyContext::add_watcher(py, watch_callback!(fail_callback)).unwrap(); - py.run( - c"import contextvars\ncontextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); + run_context_switch(py); let (watcher_error, _) = capture.take_capture().expect("missing unraisable error"); assert!(watcher_error.is_instance_of::(py)); @@ -686,156 +662,67 @@ mod watcher_tests { } #[test] - fn bound_watcher_guard_is_not_send_or_sync() { + fn context_watcher_guard_traits() { assert_not_impl_any!(super::BoundContextWatcherGuard<'_>: Send, Sync); - } - - #[test] - fn unbound_watcher_guard_is_send_and_sync() { assert_impl_all!(super::ContextWatcherGuard: Send, Sync); } - static UNBOUND_DROP_COUNT: AtomicUsize = AtomicUsize::new(0); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_unbound_drop(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if matches!(event, ContextEvent::Switched(_)) { - UNBOUND_DROP_COUNT.fetch_add(1, Ordering::Relaxed); - } - Ok(()) - } - #[test] - fn unbound_watcher_outlives_attachment_and_attaches_on_drop() { - UNBOUND_DROP_COUNT.store(0, Ordering::Relaxed); + fn unbound_watcher_attaches_on_drop_from_another_thread() { + let _guard = acquire_watcher_test_lock(); + SWITCH_COUNT.store(0, Ordering::Relaxed); let watcher = Python::attach(|py| { - PyContext::add_watcher(py, watch_callback!(record_unbound_drop)) + PyContext::add_watcher(py, watch_callback!(record_switch)) .unwrap() .unbind() }); - Python::attach(|py| { - py.run( - c"import contextvars; contextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - }); - assert!(UNBOUND_DROP_COUNT.load(Ordering::Relaxed) >= 2); + Python::attach(run_context_switch); + assert!(SWITCH_COUNT.load(Ordering::Relaxed) >= 2); - drop(watcher); - let count_after_drop = UNBOUND_DROP_COUNT.load(Ordering::Relaxed); + std::thread::spawn(move || drop(watcher)).join().unwrap(); + let count_after_drop = SWITCH_COUNT.load(Ordering::Relaxed); - Python::attach(|py| { - py.run(c"contextvars.Context().run(lambda: None)", None, None) - .unwrap(); - }); - assert_eq!(UNBOUND_DROP_COUNT.load(Ordering::Relaxed), count_after_drop); - } - - static UNBOUND_CLEAR_COUNT: AtomicUsize = AtomicUsize::new(0); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_unbound_clear(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if matches!(event, ContextEvent::Switched(_)) { - UNBOUND_CLEAR_COUNT.fetch_add(1, Ordering::Relaxed); - } - Ok(()) + Python::attach(|py| assert_no_context_switches(py, count_after_drop)); } #[test] fn unbound_watcher_can_be_cleared_with_an_attachment() { - UNBOUND_CLEAR_COUNT.store(0, Ordering::Relaxed); + let _guard = acquire_watcher_test_lock(); + SWITCH_COUNT.store(0, Ordering::Relaxed); let watcher = Python::attach(|py| { - PyContext::add_watcher(py, watch_callback!(record_unbound_clear)) + PyContext::add_watcher(py, watch_callback!(record_switch)) .unwrap() .unbind() }); let count_after_clear = Python::attach(|py| { watcher.clear(py).unwrap(); - UNBOUND_CLEAR_COUNT.load(Ordering::Relaxed) + SWITCH_COUNT.load(Ordering::Relaxed) }); - Python::attach(|py| { - py.run( - c"import contextvars; contextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - }); - assert_eq!( - UNBOUND_CLEAR_COUNT.load(Ordering::Relaxed), - count_after_clear - ); - } - - static REBOUND_COUNT: AtomicUsize = AtomicUsize::new(0); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_rebound(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if matches!(event, ContextEvent::Switched(_)) { - REBOUND_COUNT.fetch_add(1, Ordering::Relaxed); - } - Ok(()) + Python::attach(|py| assert_no_context_switches(py, count_after_clear)); } #[test] fn unbound_watcher_can_be_rebound() { - REBOUND_COUNT.store(0, Ordering::Relaxed); + let _guard = acquire_watcher_test_lock(); + SWITCH_COUNT.store(0, Ordering::Relaxed); let watcher = Python::attach(|py| { - PyContext::add_watcher(py, watch_callback!(record_rebound)) + PyContext::add_watcher(py, watch_callback!(record_switch)) .unwrap() .unbind() }); Python::attach(|py| { let watcher = watcher.into_bound(py); - py.run( - c"import contextvars; contextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - - let count_before_drop = REBOUND_COUNT.load(Ordering::Relaxed); + run_context_switch(py); + + let count_before_drop = SWITCH_COUNT.load(Ordering::Relaxed); assert!(count_before_drop >= 2); drop(watcher); - py.run(c"contextvars.Context().run(lambda: None)", None, None) - .unwrap(); - assert_eq!(REBOUND_COUNT.load(Ordering::Relaxed), count_before_drop); - }); - } - - static FORGOTTEN_GUARD_COUNT: AtomicUsize = AtomicUsize::new(0); - - #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] - fn record_forgotten_guard(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> { - if matches!(event, ContextEvent::Switched(_)) { - FORGOTTEN_GUARD_COUNT.fetch_add(1, Ordering::Relaxed); - } - Ok(()) - } - - #[test] - fn forgotten_guard() { - Python::attach(|py| { - FORGOTTEN_GUARD_COUNT.store(0, Ordering::Relaxed); - let watcher = - PyContext::add_watcher(py, watch_callback!(record_forgotten_guard)).unwrap(); - core::mem::forget(watcher); - - py.run( - c"import contextvars; contextvars.Context().run(lambda: None)", - None, - None, - ) - .unwrap(); - - assert!(FORGOTTEN_GUARD_COUNT.load(Ordering::Relaxed) > 0); + assert_no_context_switches(py, count_before_drop); }); } } From 5e5eb07792a4dac39b80775de1cafb89baac8b4a Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Mon, 31 Aug 2026 14:27:21 +0200 Subject: [PATCH 17/24] Apply suggestions from code review Co-authored-by: David Hewitt <1939362+davidhewitt@users.noreply.github.com> --- src/types/context.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/types/context.rs b/src/types/context.rs index 90e2e6d6ade..d76346e8424 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -35,6 +35,8 @@ pyobject_native_type_core!( #checkfunction=ffi::PyContext_CheckExact ); +// TODO: enable support on free-threaded builds once +// https://github.com/python/cpython/issues/155619 is fixed #[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))] impl PyContext { /// Registers a context watcher for the current interpreter. @@ -363,7 +365,8 @@ pub mod impl_ { crate::impl_::trampoline::trampoline(|py| { let event = match event { ffi::Py_CONTEXT_SWITCHED => { - let object = object.assume_borrowed_or_err(py)?; + // SAFETY: `Py_CONTEXT_SWITCHED` is documented to always have None or a context object passed + let object = unsafe { object.assume_borrowed_unchecked(py) }; if object.is_none() { ContextEvent::Switched(None) From f623224cbce64b54431a6cfcd9d5261d8a942df8 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Mon, 31 Aug 2026 14:31:17 +0200 Subject: [PATCH 18/24] Reduce context watcher trampoline code size --- src/types/context.rs | 43 ++++++++++++++++++++++++++----------------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/src/types/context.rs b/src/types/context.rs index d76346e8424..6868f54913e 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -338,6 +338,31 @@ pub mod impl_ { WatchCallback(context_watcher::) } + unsafe fn event_from_raw<'a, 'py>( + py: Python<'py>, + event: ffi::PyContextEvent, + object: *mut ffi::PyObject, + ) -> PyResult> { + match event { + ffi::Py_CONTEXT_SWITCHED => { + // SAFETY: `Py_CONTEXT_SWITCHED` is documented to always have None or a context object passed + let object = unsafe { object.assume_borrowed_unchecked(py) }; + + if object.is_none() { + Ok(ContextEvent::Switched(None)) + } else { + Ok(ContextEvent::Switched(Some(object.cast()?))) + } + } + + raw_event => { + // SAFETY: the caller guarantees that `object` follows the contract for `event`. + let object = unsafe { object.assume_borrowed_or_opt(py) }; + Ok(ContextEvent::Unknown { raw_event, object }) + } + } + } + /// C-compatible trampoline for a context watcher callback. /// /// # Safety @@ -363,23 +388,7 @@ pub mod impl_ { // - the callback's higher-ranked signature prevents borrowed event data from escaping let result = unsafe { crate::impl_::trampoline::trampoline(|py| { - let event = match event { - ffi::Py_CONTEXT_SWITCHED => { - // SAFETY: `Py_CONTEXT_SWITCHED` is documented to always have None or a context object passed - let object = unsafe { object.assume_borrowed_unchecked(py) }; - - if object.is_none() { - ContextEvent::Switched(None) - } else { - ContextEvent::Switched(Some(object.cast()?)) - } - } - - raw_event => { - let object = object.assume_borrowed_or_opt(py); - ContextEvent::Unknown { raw_event, object } - } - }; + let event = event_from_raw(py, event, object)?; (Callback::CALLBACK)(py, event)?; From e7eca0852064f25980360c5a732d9c2de30304f1 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Mon, 31 Aug 2026 15:03:54 +0200 Subject: [PATCH 19/24] Handle watcher callbacks which leave exceptions set --- src/types/context.rs | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/src/types/context.rs b/src/types/context.rs index 6868f54913e..09b4f145e9f 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -392,6 +392,10 @@ pub mod impl_ { (Callback::CALLBACK)(py, event)?; + if crate::PyErr::occurred(py) { + return Err(crate::PyErr::fetch(py)); + } + Ok(0) }) }; @@ -570,6 +574,33 @@ mod watcher_tests { }); } + #[allow(clippy::unnecessary_wraps, reason = "context watcher callback")] + fn restore_error_callback(py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { + PyRuntimeError::new_err("watcher restored error").restore(py); + Ok(()) + } + + struct RestoringErrorCallback; + + impl ContextWatcherCallbackDef for RestoringErrorCallback { + const CALLBACK: ContextWatcherCallback = restore_error_callback; + } + + #[test] + fn callback_cannot_return_success_with_an_exception_set() { + Python::attach(|py| { + // SAFETY: the thread is attached and None is valid for Py_CONTEXT_SWITCHED. + let result = unsafe { + context_watcher::(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None()) + }; + + assert_eq!(result, -1); + let error = PyErr::fetch(py); + assert!(error.is_instance_of::(py)); + assert_eq!(error.to_string(), "RuntimeError: watcher restored error"); + }); + } + #[test] #[cfg(feature = "macros")] fn callback_error_preserves_a_pending_exception() { From 349d904ce691ee7e1be3f7ca3017c41d01a47f3c Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Mon, 31 Aug 2026 15:04:48 +0200 Subject: [PATCH 20/24] Use no-std mutex helper in context watcher tests --- src/types/context.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/types/context.rs b/src/types/context.rs index 09b4f145e9f..77d209bf07e 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -432,15 +432,14 @@ mod watcher_tests { use super::impl_::{context_watcher, ContextWatcherCallback, ContextWatcherCallbackDef}; use super::{ContextEvent, PyContext}; use crate::exceptions::{PyRuntimeError, PyValueError}; + use crate::platform::sync::non_poison::{Mutex, MutexGuard}; #[cfg(feature = "macros")] use crate::test_utils::UnraisableCapture; use crate::types::PyAnyMethods; use crate::{ffi, PyErr, PyResult, Python}; - #[cfg(feature = "macros")] use alloc::string::ToString; use core::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; use static_assertions::{assert_impl_all, assert_not_impl_any}; - use std::sync::{Mutex, MutexGuard, PoisonError}; // Context watchers are interpreter-global and limited to eight slots, so tests which register // watchers must not run concurrently. @@ -449,9 +448,7 @@ mod watcher_tests { static SAW_CONTEXT: AtomicBool = AtomicBool::new(false); fn acquire_watcher_test_lock() -> MutexGuard<'static, ()> { - WATCHER_TEST_MUTEX - .lock() - .unwrap_or_else(PoisonError::into_inner) + WATCHER_TEST_MUTEX.lock() } fn run_context_switch(py: Python<'_>) { From 9def05278d7298aef5623a529113f4ec40c82b19 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Wed, 2 Sep 2026 13:36:46 +0200 Subject: [PATCH 21/24] Skip threaded context watcher test on wasm --- src/types/context.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/types/context.rs b/src/types/context.rs index 77d209bf07e..58d0c35053a 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -707,6 +707,7 @@ mod watcher_tests { assert_impl_all!(super::ContextWatcherGuard: Send, Sync); } + #[cfg(not(target_arch = "wasm32"))] // We are building wasm Python with pthreads disabled #[test] fn unbound_watcher_attaches_on_drop_from_another_thread() { let _guard = acquire_watcher_test_lock(); From 1f0fe5ac4cc1fd104374390fe095b0f054008b1e Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Wed, 2 Sep 2026 13:53:47 +0200 Subject: [PATCH 22/24] Skip watcher panic test without std --- src/types/context.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/types/context.rs b/src/types/context.rs index 58d0c35053a..a402f31518b 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -642,16 +642,20 @@ mod watcher_tests { }); } + #[cfg(wip_feature_std)] fn panic_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { panic!("context watcher panic") } + #[cfg(wip_feature_std)] struct PanickingCallback; + #[cfg(wip_feature_std)] impl ContextWatcherCallbackDef for PanickingCallback { const CALLBACK: ContextWatcherCallback = panic_callback; } + #[cfg(wip_feature_std)] #[test] fn callback_panic_does_not_cross_ffi_boundary() { Python::attach(|py| { From 43e16b70ae8db757e9d5c03f8c86d4c6999b2880 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Mon, 7 Sep 2026 11:20:44 +0200 Subject: [PATCH 23/24] Skip watcher panic test on wasm --- src/types/context.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/types/context.rs b/src/types/context.rs index a402f31518b..19fa6738732 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -642,20 +642,20 @@ mod watcher_tests { }); } - #[cfg(wip_feature_std)] + #[cfg(all(wip_feature_std, not(target_arch = "wasm32")))] fn panic_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { panic!("context watcher panic") } - #[cfg(wip_feature_std)] + #[cfg(all(wip_feature_std, not(target_arch = "wasm32")))] struct PanickingCallback; - #[cfg(wip_feature_std)] + #[cfg(all(wip_feature_std, not(target_arch = "wasm32")))] impl ContextWatcherCallbackDef for PanickingCallback { const CALLBACK: ContextWatcherCallback = panic_callback; } - #[cfg(wip_feature_std)] + #[cfg(all(wip_feature_std, not(target_arch = "wasm32")))] #[test] fn callback_panic_does_not_cross_ffi_boundary() { Python::attach(|py| { From 1f61a3451ab00db9b1a790f4f0505eeaf1b63db5 Mon Sep 17 00:00:00 2001 From: Florentin Labelle Date: Mon, 7 Sep 2026 13:02:22 +0200 Subject: [PATCH 24/24] Gate watcher panic test on unwind support --- src/types/context.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/types/context.rs b/src/types/context.rs index 19fa6738732..815e96ea0cb 100644 --- a/src/types/context.rs +++ b/src/types/context.rs @@ -642,20 +642,20 @@ mod watcher_tests { }); } - #[cfg(all(wip_feature_std, not(target_arch = "wasm32")))] + #[cfg(all(wip_feature_std, panic = "unwind"))] fn panic_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> { panic!("context watcher panic") } - #[cfg(all(wip_feature_std, not(target_arch = "wasm32")))] + #[cfg(all(wip_feature_std, panic = "unwind"))] struct PanickingCallback; - #[cfg(all(wip_feature_std, not(target_arch = "wasm32")))] + #[cfg(all(wip_feature_std, panic = "unwind"))] impl ContextWatcherCallbackDef for PanickingCallback { const CALLBACK: ContextWatcherCallback = panic_callback; } - #[cfg(all(wip_feature_std, not(target_arch = "wasm32")))] + #[cfg(all(wip_feature_std, panic = "unwind"))] #[test] fn callback_panic_does_not_cross_ffi_boundary() { Python::attach(|py| {