use crate::{
ffi,
sealed::Sealed,
types::{any::PyAnyMethods, PyAny, PyString},
Bound, Py, PyResult, PyTypeCheck, Python,
};
use std::{
cell::UnsafeCell,
marker::PhantomData,
mem::MaybeUninit,
sync::{Once, OnceState},
};
#[cfg(not(Py_GIL_DISABLED))]
use crate::PyVisit;
#[cfg(not(Py_GIL_DISABLED))]
pub struct GILProtected<T> {
value: T,
}
#[cfg(not(Py_GIL_DISABLED))]
impl<T> GILProtected<T> {
pub const fn new(value: T) -> Self {
Self { value }
}
pub fn get<'py>(&'py self, _py: Python<'py>) -> &'py T {
&self.value
}
pub fn traverse<'py>(&'py self, _visit: PyVisit<'py>) -> &'py T {
&self.value
}
}
#[cfg(not(Py_GIL_DISABLED))]
unsafe impl<T> Sync for GILProtected<T> where T: Send {}
#[doc = concat!(
"[the FAQ section](https://pyo3.rs/v",
env!("CARGO_PKG_VERSION"),
"/faq.html#im-experiencing-deadlocks-using-pyo3-with-stdsynconcelock-stdsynclazylock-lazy_static-and-once_cell)"
)]
pub struct GILOnceCell<T> {
once: Once,
data: UnsafeCell<MaybeUninit<T>>,
_marker: PhantomData<T>,
}
impl<T> Default for GILOnceCell<T> {
fn default() -> Self {
Self::new()
}
}
unsafe impl<T: Send + Sync> Sync for GILOnceCell<T> {}
unsafe impl<T: Send> Send for GILOnceCell<T> {}
impl<T> GILOnceCell<T> {
pub const fn new() -> Self {
Self {
once: Once::new(),
data: UnsafeCell::new(MaybeUninit::uninit()),
_marker: PhantomData,
}
}
#[inline]
pub fn get(&self, _py: Python<'_>) -> Option<&T> {
if self.once.is_completed() {
Some(unsafe { (*self.data.get()).assume_init_ref() })
} else {
None
}
}
#[inline]
pub fn get_or_init<F>(&self, py: Python<'_>, f: F) -> &T
where
F: FnOnce() -> T,
{
if let Some(value) = self.get(py) {
return value;
}
self.init(py, || Ok::<T, std::convert::Infallible>(f()))
.unwrap()
}
#[inline]
pub fn get_or_try_init<F, E>(&self, py: Python<'_>, f: F) -> Result<&T, E>
where
F: FnOnce() -> Result<T, E>,
{
if let Some(value) = self.get(py) {
return Ok(value);
}
self.init(py, f)
}
#[cold]
fn init<F, E>(&self, py: Python<'_>, f: F) -> Result<&T, E>
where
F: FnOnce() -> Result<T, E>,
{
let value = f()?;
let _ = self.set(py, value);
Ok(self.get(py).unwrap())
}
pub fn get_mut(&mut self) -> Option<&mut T> {
if self.once.is_completed() {
Some(unsafe { (*self.data.get()).assume_init_mut() })
} else {
None
}
}
pub fn set(&self, _py: Python<'_>, value: T) -> Result<(), T> {
let mut value = Some(value);
self.once.call_once_force(|_| {
unsafe {
(*self.data.get()).write(value.take().unwrap());
}
});
match value {
Some(value) => Err(value),
None => Ok(()),
}
}
pub fn take(&mut self) -> Option<T> {
if self.once.is_completed() {
self.once = Once::new();
Some(unsafe { self.data.get_mut().assume_init_read() })
} else {
None
}
}
pub fn into_inner(mut self) -> Option<T> {
self.take()
}
}
impl<T> GILOnceCell<Py<T>> {
pub fn clone_ref(&self, py: Python<'_>) -> Self {
let cloned = Self {
once: Once::new(),
data: UnsafeCell::new(MaybeUninit::uninit()),
_marker: PhantomData,
};
if let Some(value) = self.get(py) {
let _ = cloned.set(py, value.clone_ref(py));
}
cloned
}
}
impl<T> GILOnceCell<Py<T>>
where
T: PyTypeCheck,
{
pub fn import<'py>(
&self,
py: Python<'py>,
module_name: &str,
attr_name: &str,
) -> PyResult<&Bound<'py, T>> {
self.get_or_try_init(py, || {
let type_object = py
.import(module_name)?
.getattr(attr_name)?
.downcast_into()?;
Ok(type_object.unbind())
})
.map(|ty| ty.bind(py))
}
}
impl<T> Drop for GILOnceCell<T> {
fn drop(&mut self) {
if self.once.is_completed() {
unsafe { MaybeUninit::assume_init_drop(self.data.get_mut()) }
}
}
}
#[macro_export]
macro_rules! intern {
($py: expr, $text: expr) => {{
static INTERNED: $crate::sync::Interned = $crate::sync::Interned::new($text);
INTERNED.get($py)
}};
}
#[doc(hidden)]
pub struct Interned(&'static str, GILOnceCell<Py<PyString>>);
impl Interned {
pub const fn new(value: &'static str) -> Self {
Interned(value, GILOnceCell::new())
}
#[inline]
pub fn get<'py>(&self, py: Python<'py>) -> &Bound<'py, PyString> {
self.1
.get_or_init(py, || PyString::intern(py, self.0).into())
.bind(py)
}
}
#[cfg_attr(not(Py_GIL_DISABLED), allow(unused_variables))]
pub fn with_critical_section<F, R>(object: &Bound<'_, PyAny>, f: F) -> R
where
F: FnOnce() -> R,
{
#[cfg(Py_GIL_DISABLED)]
{
struct Guard(crate::ffi::PyCriticalSection);
impl Drop for Guard {
fn drop(&mut self) {
unsafe {
crate::ffi::PyCriticalSection_End(&mut self.0);
}
}
}
let mut guard = Guard(unsafe { std::mem::zeroed() });
unsafe { crate::ffi::PyCriticalSection_Begin(&mut guard.0, object.as_ptr()) };
f()
}
#[cfg(not(Py_GIL_DISABLED))]
{
f()
}
}
#[cfg(rustc_has_once_lock)]
mod once_lock_ext_sealed {
pub trait Sealed {}
impl<T> Sealed for std::sync::OnceLock<T> {}
}
pub trait OnceExt: Sealed {
fn call_once_py_attached(&self, py: Python<'_>, f: impl FnOnce());
fn call_once_force_py_attached(&self, py: Python<'_>, f: impl FnOnce(&OnceState));
}
#[cfg(rustc_has_once_lock)]
pub trait OnceLockExt<T>: once_lock_ext_sealed::Sealed {
fn get_or_init_py_attached<F>(&self, py: Python<'_>, f: F) -> &T
where
F: FnOnce() -> T;
}
struct Guard(*mut crate::ffi::PyThreadState);
impl Drop for Guard {
fn drop(&mut self) {
unsafe { ffi::PyEval_RestoreThread(self.0) };
}
}
impl OnceExt for Once {
fn call_once_py_attached(&self, py: Python<'_>, f: impl FnOnce()) {
if self.is_completed() {
return;
}
init_once_py_attached(self, py, f)
}
fn call_once_force_py_attached(&self, py: Python<'_>, f: impl FnOnce(&OnceState)) {
if self.is_completed() {
return;
}
init_once_force_py_attached(self, py, f);
}
}
#[cfg(rustc_has_once_lock)]
impl<T> OnceLockExt<T> for std::sync::OnceLock<T> {
fn get_or_init_py_attached<F>(&self, py: Python<'_>, f: F) -> &T
where
F: FnOnce() -> T,
{
#[allow(clippy::incompatible_msrv)]
self.get()
.unwrap_or_else(|| init_once_lock_py_attached(self, py, f))
}
}
#[cold]
fn init_once_py_attached<F, T>(once: &Once, _py: Python<'_>, f: F)
where
F: FnOnce() -> T,
{
let ts_guard = Guard(unsafe { ffi::PyEval_SaveThread() });
once.call_once(move || {
drop(ts_guard);
f();
});
}
#[cold]
fn init_once_force_py_attached<F, T>(once: &Once, _py: Python<'_>, f: F)
where
F: FnOnce(&OnceState) -> T,
{
let ts_guard = Guard(unsafe { ffi::PyEval_SaveThread() });
once.call_once_force(move |state| {
drop(ts_guard);
f(state);
});
}
#[cfg(rustc_has_once_lock)]
#[cold]
fn init_once_lock_py_attached<'a, F, T>(
lock: &'a std::sync::OnceLock<T>,
_py: Python<'_>,
f: F,
) -> &'a T
where
F: FnOnce() -> T,
{
let ts_guard = Guard(unsafe { ffi::PyEval_SaveThread() });
#[allow(clippy::incompatible_msrv)]
let value = lock.get_or_init(move || {
drop(ts_guard);
f()
});
value
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::{PyDict, PyDictMethods};
#[test]
fn test_intern() {
Python::with_gil(|py| {
let foo1 = "foo";
let foo2 = intern!(py, "foo");
let foo3 = intern!(py, stringify!(foo));
let dict = PyDict::new(py);
dict.set_item(foo1, 42_usize).unwrap();
assert!(dict.contains(foo2).unwrap());
assert_eq!(
dict.get_item(foo3)
.unwrap()
.unwrap()
.extract::<usize>()
.unwrap(),
42
);
});
}
#[test]
fn test_once_cell() {
Python::with_gil(|py| {
let mut cell = GILOnceCell::new();
assert!(cell.get(py).is_none());
assert_eq!(cell.get_or_try_init(py, || Err(5)), Err(5));
assert!(cell.get(py).is_none());
assert_eq!(cell.get_or_try_init(py, || Ok::<_, ()>(2)), Ok(&2));
assert_eq!(cell.get(py), Some(&2));
assert_eq!(cell.get_or_try_init(py, || Err(5)), Ok(&2));
assert_eq!(cell.take(), Some(2));
assert_eq!(cell.into_inner(), None);
let cell_py = GILOnceCell::new();
assert!(cell_py.clone_ref(py).get(py).is_none());
cell_py.get_or_init(py, || py.None());
assert!(cell_py.clone_ref(py).get(py).unwrap().is_none(py));
})
}
#[test]
fn test_once_cell_drop() {
#[derive(Debug)]
struct RecordDrop<'a>(&'a mut bool);
impl Drop for RecordDrop<'_> {
fn drop(&mut self) {
*self.0 = true;
}
}
Python::with_gil(|py| {
let mut dropped = false;
let cell = GILOnceCell::new();
cell.set(py, RecordDrop(&mut dropped)).unwrap();
let drop_container = cell.get(py).unwrap();
assert!(!*drop_container.0);
drop(cell);
assert!(dropped);
});
}
#[cfg(feature = "macros")]
#[cfg(not(target_arch = "wasm32"))] #[test]
fn test_critical_section() {
use std::sync::{
atomic::{AtomicBool, Ordering},
Barrier,
};
let barrier = Barrier::new(2);
#[crate::pyclass(crate = "crate")]
struct BoolWrapper(AtomicBool);
let bool_wrapper = Python::with_gil(|py| -> Py<BoolWrapper> {
Py::new(py, BoolWrapper(AtomicBool::new(false))).unwrap()
});
std::thread::scope(|s| {
s.spawn(|| {
Python::with_gil(|py| {
let b = bool_wrapper.bind(py);
with_critical_section(b, || {
barrier.wait();
std::thread::sleep(std::time::Duration::from_millis(10));
b.borrow().0.store(true, Ordering::Release);
})
});
});
s.spawn(|| {
barrier.wait();
Python::with_gil(|py| {
let b = bool_wrapper.bind(py);
with_critical_section(b, || {
assert!(b.borrow().0.load(Ordering::Acquire));
});
});
});
});
}
#[test]
#[cfg(not(target_arch = "wasm32"))] fn test_once_ext() {
let init = Once::new();
std::thread::scope(|s| {
let handle = s.spawn(|| {
Python::with_gil(|py| {
init.call_once_py_attached(py, || panic!());
})
});
assert!(handle.join().is_err());
let handle = s.spawn(|| {
Python::with_gil(|py| {
init.call_once_py_attached(py, || {});
});
});
assert!(handle.join().is_err());
Python::with_gil(|py| {
init.call_once_force_py_attached(py, |state| {
assert!(state.is_poisoned());
});
init.call_once_py_attached(py, || {});
});
});
}
#[cfg(rustc_has_once_lock)]
#[cfg(not(target_arch = "wasm32"))] #[test]
fn test_once_lock_ext() {
let cell = std::sync::OnceLock::new();
std::thread::scope(|s| {
assert!(cell.get().is_none());
s.spawn(|| {
Python::with_gil(|py| {
assert_eq!(*cell.get_or_init_py_attached(py, || 12345), 12345);
});
});
});
assert_eq!(cell.get(), Some(&12345));
}
}