Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

make GILOnceCell threadsafe #4512

Merged
merged 10 commits into from
Oct 8, 2024
12 changes: 6 additions & 6 deletions guide/src/faq.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,18 @@

Sorry that you're having trouble using PyO3. If you can't find the answer to your problem in the list below, you can also reach out for help on [GitHub Discussions](https://github.com/PyO3/pyo3/discussions) and on [Discord](https://discord.gg/33kcChzH7f).

## I'm experiencing deadlocks using PyO3 with lazy_static or once_cell!
## I'm experiencing deadlocks using PyO3 with `std::sync::OnceLock`, `std::sync::LazyLock`, `lazy_static`, and `once_cell`!

`lazy_static` and `once_cell::sync` both use locks to ensure that initialization is performed only by a single thread. Because the Python GIL is an additional lock this can lead to deadlocks in the following way:
`OnceLock`, `LazyLock`, and their thirdparty predecessors use blocking to ensure only one thread ever initializes them. Because the Python GIL is an additional lock this can lead to deadlocks in the following way:

1. A thread (thread A) which has acquired the Python GIL starts initialization of a `lazy_static` value.
1. A thread (thread A) which has acquired the Python GIL starts initialization of a `OnceLock` value.
2. The initialization code calls some Python API which temporarily releases the GIL e.g. `Python::import`.
3. Another thread (thread B) acquires the Python GIL and attempts to access the same `lazy_static` value.
4. Thread B is blocked, because it waits for `lazy_static`'s initialization to lock to release.
3. Another thread (thread B) acquires the Python GIL and attempts to access the same `OnceLock` value.
4. Thread B is blocked, because it waits for `OnceLock`'s initialization to lock to release.
5. Thread A is blocked, because it waits to re-acquire the GIL which thread B still holds.
6. Deadlock.

PyO3 provides a struct [`GILOnceCell`] which works equivalently to `OnceCell` but relies solely on the Python GIL for thread safety. This means it can be used in place of `lazy_static` or `once_cell` where you are experiencing the deadlock described above. See the documentation for [`GILOnceCell`] for an example how to use it.
PyO3 provides a struct [`GILOnceCell`] which works similarly to these types but avoids risk of deadlocking with the Python GIL. This means it can be used in place of other choices when you are experiencing the deadlock described above. See the documentation for [`GILOnceCell`] for further details and an example how to use it.

[`GILOnceCell`]: {{#PYO3_DOCS_URL}}/pyo3/sync/struct.GILOnceCell.html

Expand Down
2 changes: 2 additions & 0 deletions guide/src/migration.md
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,8 @@ impl<'a, 'py> IntoPyObject<'py> for &'a MyPyObjectWrapper {
PyO3 0.23 introduces preliminary support for the new free-threaded build of
CPython 3.13. PyO3 features that implicitly assumed the existence of the GIL
are not exposed in the free-threaded build, since they are no longer safe.
Other features, such as `GILOnceCell`, have been internally rewritten to be
threadsafe without the GIL.

If you make use of these features then you will need to account for the
unavailability of this API in the free-threaded build. One way to handle it is
Expand Down
1 change: 1 addition & 0 deletions newsfragments/4512.changed.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
`GILOnceCell` is now thread-safe for the Python 3.13 freethreaded builds.
2 changes: 1 addition & 1 deletion src/conversions/chrono_tz.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ mod tests {
fn test_into_pyobject() {
Python::with_gil(|py| {
let assert_eq = |l: Bound<'_, PyAny>, r: Bound<'_, PyAny>| {
assert!(l.eq(r).unwrap());
assert!(l.eq(&r).unwrap(), "{:?} != {:?}", l, r);
};

assert_eq(
Expand Down
191 changes: 155 additions & 36 deletions src/sync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use crate::{
types::{any::PyAnyMethods, PyAny, PyString},
Bound, Py, PyResult, PyTypeCheck, Python,
};
use std::cell::UnsafeCell;
use std::{cell::UnsafeCell, marker::PhantomData, mem::MaybeUninit, sync::Once};

#[cfg(not(Py_GIL_DISABLED))]
use crate::PyVisit;
Expand Down Expand Up @@ -62,24 +62,36 @@ impl<T> GILProtected<T> {
#[cfg(not(Py_GIL_DISABLED))]
unsafe impl<T> Sync for GILProtected<T> where T: Send {}

/// A write-once cell similar to [`once_cell::OnceCell`](https://docs.rs/once_cell/latest/once_cell/).
/// A write-once primitive similar to [`std::sync::OnceLock<T>`].
///
/// Unlike `once_cell::sync` which blocks threads to achieve thread safety, this implementation
/// uses the Python GIL to mediate concurrent access. This helps in cases where `once_cell` or
/// `lazy_static`'s synchronization strategy can lead to deadlocks when interacting with the Python
/// GIL. For an example, see
#[doc = concat!("[the FAQ section](https://pyo3.rs/v", env!("CARGO_PKG_VERSION"), "/faq.html)")]
/// Unlike `OnceLock<T>` which blocks threads to achieve thread safety, `GilOnceCell<T>`
/// allows calls to [`get_or_init`][GILOnceCell::get_or_init] and
/// [`get_or_try_init`][GILOnceCell::get_or_try_init] to race to create an initialized value.
/// (It is still guaranteed that only one thread will ever write to the cell.)
///
/// On Python versions that run with the Global Interpreter Lock (GIL), this helps to avoid
/// deadlocks between initialization and the GIL. For an example of such a deadlock, see
#[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)"
)]
/// of the guide.
///
/// Note that:
/// 1) `get_or_init` and `get_or_try_init` do not protect against infinite recursion
/// from reentrant initialization.
/// 2) If the initialization function `f` provided to `get_or_init` (or `get_or_try_init`)
/// temporarily releases the GIL (e.g. by calling `Python::import`) then it is possible
/// for a second thread to also begin initializing the `GITOnceCell`. Even when this
/// happens `GILOnceCell` guarantees that only **one** write to the cell ever occurs -
/// this is treated as a race, other threads will discard the value they compute and
/// return the result of the first complete computation.
/// Note that because the GIL blocks concurrent execution, in practice the means that
/// [`get_or_init`][GILOnceCell::get_or_init] and
/// [`get_or_try_init`][GILOnceCell::get_or_try_init] may race if the initialization
/// function leads to the GIL being released and a thread context switch. This can
/// happen when importing or calling any Python code, as long as it releases the
/// GIL at some point. On free-threaded Python without any GIL, the race is
/// more likely since there is no GIL to prevent races. In the future, PyO3 may change
/// the semantics of GILOnceCell to behave more like the GIL build in the future.
///
/// # Re-entrant initialization
///
/// [`get_or_init`][GILOnceCell::get_or_init] and
/// [`get_or_try_init`][GILOnceCell::get_or_try_init] do not protect against infinite recursion
/// from reentrant initialization.
///
/// # Examples
///
Expand All @@ -100,25 +112,64 @@ unsafe impl<T> Sync for GILProtected<T> where T: Send {}
/// }
/// # Python::with_gil(|py| assert_eq!(get_shared_list(py).len(), 0));
/// ```
#[derive(Default)]
pub struct GILOnceCell<T>(UnsafeCell<Option<T>>);
pub struct GILOnceCell<T> {
once: Once,
data: UnsafeCell<MaybeUninit<T>>,
davidhewitt marked this conversation as resolved.
Show resolved Hide resolved

/// (Copied from std::sync::OnceLock)
///
/// `PhantomData` to make sure dropck understands we're dropping T in our Drop impl.
davidhewitt marked this conversation as resolved.
Show resolved Hide resolved
///
/// ```compile_error,E0597
/// use pyo3::Python;
/// use pyo3::sync::GILOnceCell;
///
/// struct A<'a>(#[allow(dead_code)] &'a str);
///
/// impl<'a> Drop for A<'a> {
/// fn drop(&mut self) {}
/// }
///
/// let cell = GILOnceCell::new();
/// {
/// let s = String::new();
/// let _ = Python::with_gil(|py| cell.set(py,A(&s)));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One of the builds is failing here:

error[E0597]: `s` does not live long enough
  --> src/sync.rs:142:50
   |
24 |     let s = String::new();
   |         - binding `s` declared here
25 |     let _ = Python::with_gil(|py| cell.set(py,A(&s)));
   |                              ----                ^ borrowed value does not live long enough
   |                              |
   |                              value captured here
26 | }
   | - `s` dropped here while still borrowed
27 | } _doctest_main_src_sync_rs_121_0() }
   | - borrow might be used here, when `cell` is dropped and runs the `Drop` code for type `GILOnceCell`

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah my mistake, this is meant to be marked compile_fail as a UI test that the type behaves properly (imported from std).

/// }
/// ```
_marker: PhantomData<T>,
}

impl<T> Default for GILOnceCell<T> {
fn default() -> Self {
Self::new()
}
}

// T: Send is needed for Sync because the thread which drops the GILOnceCell can be different
// to the thread which fills it.
// to the thread which fills it. (e.g. think scoped thread which fills the cell and then exits,
// leaving the cell to be dropped by the main thread).
unsafe impl<T: Send + Sync> Sync for GILOnceCell<T> {}
unsafe impl<T: Send> Send for GILOnceCell<T> {}

impl<T> GILOnceCell<T> {
/// Create a `GILOnceCell` which does not yet contain a value.
pub const fn new() -> Self {
Self(UnsafeCell::new(None))
Self {
once: Once::new(),
data: UnsafeCell::new(MaybeUninit::uninit()),
_marker: PhantomData,
}
}

/// Get a reference to the contained value, or `None` if the cell has not yet been written.
#[inline]
pub fn get(&self, _py: Python<'_>) -> Option<&T> {
// Safe because if the cell has not yet been written, None is returned.
unsafe { &*self.0.get() }.as_ref()
if self.once.is_completed() {
// SAFETY: the cell has been written.
Some(unsafe { (*self.data.get()).assume_init_ref() })
} else {
None
}
}

/// Get a reference to the contained value, initializing it if needed using the provided
Expand Down Expand Up @@ -163,6 +214,10 @@ impl<T> GILOnceCell<T> {
// Note that f() could temporarily release the GIL, so it's possible that another thread
// writes to this GILOnceCell before f() finishes. That's fine; we'll just have to discard
// the value computed here and accept a bit of wasted computation.

// TODO: on the freethreaded build, consider wrapping this pair of operations in a
// critical section (requires a critical section API which can use a PyMutex without
// an object.)
let value = f()?;
let _ = self.set(py, value);

Expand All @@ -172,45 +227,77 @@ impl<T> GILOnceCell<T> {
/// Get the contents of the cell mutably. This is only possible if the reference to the cell is
/// unique.
pub fn get_mut(&mut self) -> Option<&mut T> {
self.0.get_mut().as_mut()
if self.once.is_completed() {
// SAFETY: the cell has been written.
Some(unsafe { (*self.data.get()).assume_init_mut() })
} else {
None
}
}

/// Set the value in the cell.
///
/// If the cell has already been written, `Err(value)` will be returned containing the new
/// value which was not written.
pub fn set(&self, _py: Python<'_>, value: T) -> Result<(), T> {
// Safe because GIL is held, so no other thread can be writing to this cell concurrently.
let inner = unsafe { &mut *self.0.get() };
if inner.is_some() {
return Err(value);
}
let mut value = Some(value);
// NB this can block, but since this is only writing a single value and
// does not call arbitrary python code, we don't need to worry about
// deadlocks with the GIL.
self.once.call_once_force(|_| {
// SAFETY: no other threads can be writing this value, because we are
// inside the `call_once_force` closure.
unsafe {
// `.take().unwrap()` will never panic
(*self.data.get()).write(value.take().unwrap());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

neat, I'd never seen this pattern before, it took me a few minutes to puzzle out but it's a nice pattern only consuming the value if it's actually written.

}
});

*inner = Some(value);
Ok(())
match value {
// Some other thread wrote to the cell first
Some(value) => Err(value),
None => Ok(()),
}
}

/// Takes the value out of the cell, moving it back to an uninitialized state.
///
/// Has no effect and returns None if the cell has not yet been written.
pub fn take(&mut self) -> Option<T> {
self.0.get_mut().take()
if self.once.is_completed() {
// Reset the cell to its default state so that it won't try to
// drop the value again.
self.once = Once::new();
// SAFETY: the cell has been written. `self.once` has been reset,
// so when `self` is dropped the value won't be read again.
Some(unsafe { self.data.get_mut().assume_init_read() })
} else {
None
}
}

/// Consumes the cell, returning the wrapped value.
///
/// Returns None if the cell has not yet been written.
pub fn into_inner(self) -> Option<T> {
self.0.into_inner()
pub fn into_inner(mut self) -> Option<T> {
self.take()
}
}

impl<T> GILOnceCell<Py<T>> {
/// Create a new cell that contains a new Python reference to the same contained object.
/// Creates a new cell that contains a new Python reference to the same contained object.
///
/// Returns an uninitialised cell if `self` has not yet been initialised.
/// Returns an uninitialized cell if `self` has not yet been initialized.
pub fn clone_ref(&self, py: Python<'_>) -> Self {
Self(UnsafeCell::new(self.get(py).map(|ob| ob.clone_ref(py))))
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
}
}

Expand Down Expand Up @@ -266,6 +353,15 @@ where
}
}

impl<T> Drop for GILOnceCell<T> {
fn drop(&mut self) {
if self.once.is_completed() {
// SAFETY: the cell has been written.
unsafe { MaybeUninit::assume_init_drop(self.data.get_mut()) }
}
}
}

/// Interns `text` as a Python string and stores a reference to it in static storage.
///
/// A reference to the same Python string is returned on each invocation.
Expand Down Expand Up @@ -429,6 +525,29 @@ mod tests {
})
}

#[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"))] // We are building wasm Python with pthreads disabled
#[test]
Expand Down
Loading