1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188
//! A write-once cell mediated by the Python GIL.
use crate::Python;
use std::cell::UnsafeCell;
/// A write-once cell similar to [`once_cell::OnceCell`](https://docs.rs/once_cell/1.4.0/once_cell/).
///
/// 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_sync` or
/// `lazy_static`'s synchronization strategy can lead to deadlocks when interacting with the Python
/// GIL. For an example, see [the FAQ section](https://pyo3.rs/latest/faq.html) of the guide.
///
/// # Examples
///
/// The following example shows how to use `GILOnceCell` to share a reference to a Python list
/// between threads:
///
/// ```
/// use pyo3::once_cell::GILOnceCell;
/// use pyo3::prelude::*;
/// use pyo3::types::PyList;
///
/// static LIST_CELL: GILOnceCell<Py<PyList>> = GILOnceCell::new();
///
/// pub fn get_shared_list(py: Python<'_>) -> &PyList {
/// LIST_CELL
/// .get_or_init(py, || PyList::empty(py).into())
/// .as_ref(py)
/// }
/// # Python::with_gil(|py| assert_eq!(get_shared_list(py).len(), 0));
/// ```
pub struct GILOnceCell<T>(UnsafeCell<Option<T>>);
// T: Send is needed for Sync because the thread which drops the GILOnceCell can be different
// to the thread which fills it.
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))
}
/// 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()
}
/// Get a reference to the contained value, initializing it if needed using the provided
/// closure.
///
/// Note that:
/// 1) reentrant initialization can cause a stack overflow.
/// 2) if f() temporarily releases the GIL (e.g. by calling `Python::import`) then it is
/// possible (and well-defined) that a second thread may also call get_or_init and begin
/// calling `f()`. Even when this happens `GILOnceCell` guarantees that only **one** write
/// to the cell ever occurs - other threads will simply discard the value they compute and
/// return the result of the first complete computation.
/// 3) if f() does not release the GIL and does not panic, it is guaranteed to be called
/// exactly once, even if multiple threads attempt to call `get_or_init`
/// 4) if f() can panic but still does not release the GIL, it may be called multiple times,
/// but it is guaranteed that f() will never be called concurrently
#[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, f)
}
#[cold]
fn init<F>(&self, py: Python<'_>, f: F) -> &T
where
F: FnOnce() -> 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.
let value = f();
let _ = self.set(py, value);
self.get(py).unwrap()
}
/// 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> {
// Safe because we have &mut self
unsafe { &mut *self.0.get() }.as_mut()
}
/// 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);
}
*inner = Some(value);
Ok(())
}
}
/// 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.
///
/// # Example: Using `intern!` to avoid needlessly recreating the same Python string
///
/// ```
/// use pyo3::intern;
/// # use pyo3::{pyfunction, types::PyDict, wrap_pyfunction, PyResult, Python};
///
/// #[pyfunction]
/// fn create_dict(py: Python<'_>) -> PyResult<&PyDict> {
/// let dict = PyDict::new(py);
/// // 👇 A new `PyString` is created
/// // for every call of this function.
/// dict.set_item("foo", 42)?;
/// Ok(dict)
/// }
///
/// #[pyfunction]
/// fn create_dict_faster(py: Python<'_>) -> PyResult<&PyDict> {
/// let dict = PyDict::new(py);
/// // 👇 A `PyString` is created once and reused
/// // for the lifetime of the program.
/// dict.set_item(intern!(py, "foo"), 42)?;
/// Ok(dict)
/// }
/// #
/// # Python::with_gil(|py| {
/// # let fun = wrap_pyfunction!(create_dict_faster, py).unwrap();
/// # let dict = fun.call0().unwrap();
/// # assert!(dict.contains("foo").unwrap());
/// # });
/// ```
#[macro_export]
macro_rules! intern {
($py: expr, $text: expr) => {{
fn isolate_from_dyn_env(py: $crate::Python<'_>) -> &$crate::types::PyString {
static INTERNED: $crate::once_cell::GILOnceCell<$crate::Py<$crate::types::PyString>> =
$crate::once_cell::GILOnceCell::new();
INTERNED
.get_or_init(py, || {
$crate::conversion::IntoPy::into_py(
$crate::types::PyString::intern(py, $text),
py,
)
})
.as_ref(py)
}
isolate_from_dyn_env($py)
}};
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::PyDict;
#[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().extract::<usize>().unwrap(), 42);
});
}
}