diff --git a/src/atomic.rs b/src/atomic.rs new file mode 100644 index 00000000000..0093c6d2c23 --- /dev/null +++ b/src/atomic.rs @@ -0,0 +1,220 @@ +use std::{ + marker::PhantomData, + ptr::NonNull, + sync::atomic::{AtomicPtr, Ordering}, +}; + +use crate::{ffi, Bound, Py, Python}; + +/// Variant of [`Py`] that can be atomically swapped +#[repr(transparent)] +#[derive(Debug)] +pub struct AtomicPy { + inner: AtomicPtr, + _marker: PhantomData, +} + +// SAFETY: same as `Py` +unsafe impl Send for AtomicPy {} +unsafe impl Sync for AtomicPy {} + +impl AtomicPy { + /// Create an [`AtomicPy`] holding the [`Py`] + #[inline] + pub fn from_py(obj: Py) -> Self { + Self { + inner: AtomicPtr::new(obj.into_ptr()), + _marker: PhantomData, + } + } + + /// Create an [`AtomicPy`] holding the [`Bound`] + #[inline] + pub fn from_bound(obj: Bound<'_, T>) -> Self { + Self::from_py(obj.unbind()) + } + + /// Consumes the [`AtomicPy`] and turns it back into a [`Py`] + #[inline] + pub fn into_inner(self) -> Py { + let ptr = *std::mem::ManuallyDrop::new(self).inner.get_mut(); + // SAFETY: `ptr` is a non-null, owned pointer to `ffi::PyObject` of type `T` + unsafe { Py::from_owned_ptr_unchecked(ptr) } + } + + /// Stores `obj` in the [`AtomicPy`], returning the previous value. + /// + /// See [`swap`](AtomicPy::swap) + #[inline] + pub fn swap_unbound(&self, obj: Py, order: Ordering) -> Py { + let ptr = self.inner.swap(obj.into_ptr(), order); + // SAFETY: `ptr` is a non-null, owned pointer to `ffi::PyObject` of type `T` + unsafe { Py::from_owned_ptr_unchecked(ptr) } + } + + /// Stores `obj` in the [`AtomicPy`], returning the previous value. + /// + /// `swap` takes an [Ordering] argument which describes the memory ordering of this operation. + /// All ordering modes are possible. Note that using [Acquire](Ordering::Acquire) makes the + /// store part of this operation [Relaxed](Ordering::Relaxed), and using + /// [Release](Ordering::Release) makes the load part [Relaxed](Ordering::Relaxed). + #[inline] + pub fn swap<'py>(&self, obj: Bound<'py, T>, order: Ordering) -> Bound<'py, T> { + let py = obj.py(); + self.swap_unbound(obj.unbind(), order).into_bound(py) + } +} + +impl From> for AtomicPy { + /// Create an [`AtomicPy`] holding the [`Py`] + fn from(obj: Py) -> Self { + Self::from_py(obj) + } +} + +impl From> for AtomicPy { + /// Create an [`AtomicPy`] holding the [`Bound`] + fn from(obj: Bound<'_, T>) -> Self { + Self::from_bound(obj) + } +} + +impl Drop for AtomicPy { + /// Drop the inner [`Py`] + fn drop(&mut self) { + // SAFETY: `inner` is a non-null, owned pointer to `ffi::PyObject` of type `T` + unsafe { Py::::from_owned_ptr_unchecked(*self.inner.get_mut()) }; + } +} + +/// Variant of [`Option>`] that can be atomically swapped +#[repr(transparent)] +#[derive(Debug)] +pub struct AtomicOptionPy { + inner: AtomicPtr, + _marker: PhantomData, +} + +// SAFETY: same as `Py` +unsafe impl Send for AtomicOptionPy {} +unsafe impl Sync for AtomicOptionPy {} + +impl AtomicOptionPy { + /// Create an [`AtomicOptionPy`] holding the [`Py`] + #[inline] + pub fn from_py(obj: Py) -> Self { + Self { + inner: AtomicPtr::new(obj.into_ptr()), + _marker: PhantomData, + } + } + + /// Create an [`AtomicOptionPy`] holding the [`Bound`] + #[inline] + pub fn from_bound(obj: Bound<'_, T>) -> Self { + Self::from_py(obj.unbind()) + } + + /// Consumes the [`AtomicOptionPy`] and turns it back into a [`Py`] or [`None`] if empty + #[inline] + pub fn into_inner(self) -> Option> { + let ptr = *std::mem::ManuallyDrop::new(self).inner.get_mut(); + // SAFETY: `ptr` is a owned pointer to `ffi::PyObject` of type `T` + NonNull::new(ptr).map(|ptr| unsafe { Py::from_owned_non_null(ptr) }) + } + + /// Takes the object out of the [`AtomicOptionPy`], leaving [`None`] in its place + /// + /// Note: This uses [`swap`](Self::swap) under the hood. + #[inline] + pub fn take<'py>(&self, py: Python<'py>, order: Ordering) -> Option> { + self.swap(py, None, order) + } + + /// Stores `obj` in the [`AtomicPy`], returning the previous value. + /// + /// See [`swap`](AtomicPy::swap) + #[inline] + pub fn swap_unbound(&self, obj: Option>, order: Ordering) -> Option> { + let ptr = self + .inner + .swap(obj.map(Py::into_ptr).unwrap_or_default(), order); + // SAFETY: `ptr` is an owned pointer to `ffi::PyObject` of type `T` + NonNull::new(ptr).map(|ptr| unsafe { Py::from_owned_non_null(ptr) }) + } + + /// Stores `obj` in the [`AtomicPy`], returning the previous value. + /// + /// `swap` takes an [Ordering] argument which describes the memory ordering of this operation. + /// All ordering modes are possible. Note that using [Acquire](Ordering::Acquire) makes the + /// store part of this operation [Relaxed](Ordering::Relaxed), and using + /// [Release](Ordering::Release) makes the load part [Relaxed](Ordering::Relaxed). + #[inline] + pub fn swap<'py>( + &self, + py: Python<'py>, + obj: Option>, + order: Ordering, + ) -> Option> { + self.swap_unbound(obj.map(Bound::unbind), order) + .map(|obj| obj.into_bound(py)) + } +} + +impl Default for AtomicOptionPy { + fn default() -> Self { + Self { + inner: Default::default(), + _marker: Default::default(), + } + } +} + +impl From> for AtomicOptionPy { + /// Create an [`AtomicPy`] holding the [`Py`] + fn from(obj: Py) -> Self { + Self::from_py(obj) + } +} + +impl From> for AtomicOptionPy { + /// Create an [`AtomicPy`] holding the [`Bound`] + fn from(obj: Bound<'_, T>) -> Self { + Self::from_bound(obj) + } +} + +impl From>> for AtomicOptionPy { + /// Create an [`AtomicPy`] holding the [`Py`] + fn from(obj: Option>) -> Self { + if let Some(obj) = obj { + Self::from_py(obj) + } else { + Self::default() + } + } +} + +impl From>> for AtomicOptionPy { + /// Create an [`AtomicPy`] holding the [`Bound`] + fn from(obj: Option>) -> Self { + if let Some(obj) = obj { + Self::from_bound(obj) + } else { + Self::default() + } + } +} + +impl Drop for AtomicOptionPy { + /// Drop the inner [`Py`] + fn drop(&mut self) { + if let Some(ptr) = NonNull::new(*self.inner.get_mut()) { + // SAFETY: `ptr` is an owned pointer to `ffi::PyObject` of type `T` + unsafe { Py::::from_owned_non_null(ptr) }; + } + } +} + +#[cfg(test)] +mod tests {} diff --git a/src/instance.rs b/src/instance.rs index 3bbbb4729b5..cd2e65b57d8 100644 --- a/src/instance.rs +++ b/src/instance.rs @@ -1876,6 +1876,15 @@ impl Py { Py(unsafe { NonNull::new_unchecked(ptr) }, PhantomData) } + /// Constructs a new `Py` instance by taking ownership of the given FFI pointer. + /// + /// # Safety + /// + /// - `ptr` must be an owned Python reference or type `T`. + pub(crate) unsafe fn from_owned_non_null(ptr: NonNull) -> Self { + Py(ptr, PhantomData) + } + /// Create a `Py` instance by creating a new reference from the given FFI pointer. /// /// # Safety diff --git a/src/lib.rs b/src/lib.rs index 34b4a6ae156..5e4407f00cb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -422,6 +422,7 @@ pub mod impl_; mod internal_tricks; mod internal; +mod atomic; pub mod buffer; pub mod call; pub mod conversion; @@ -449,6 +450,7 @@ mod version; #[allow(unused_imports)] // with no features enabled this module has no public exports pub use crate::conversions::*; +pub use atomic::{AtomicOptionPy, AtomicPy}; #[cfg(feature = "macros")] pub use pyo3_macros::{