-
Notifications
You must be signed in to change notification settings - Fork 678
/
once_cell.rs
146 lines (134 loc) · 5.52 KB
/
once_cell.rs
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
//! 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.
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
pub fn get_or_init<F>(&self, py: Python<'_>, f: F) -> &T
where
F: FnOnce() -> T,
{
let inner = unsafe { &*self.0.get() }.as_ref();
if let Some(value) = inner {
return value;
}
// 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(())
}
}
/// Converts `value` into a Python object and stores it in static storage. The same Python object
/// is returned on each invocation.
///
/// Because it is stored in a static, this object's destructor will not run.
///
/// # Example: Using `intern!` to avoid needlessly recreating the same object
///
/// ```
/// use pyo3::intern;
/// # use pyo3::{pyfunction, types::PyDict, 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)
/// }
/// ```
#[macro_export]
macro_rules! intern {
($py: expr, $value: expr) => {{
static INTERNED: $crate::once_cell::GILOnceCell<$crate::PyObject> =
$crate::once_cell::GILOnceCell::new();
INTERNED
.get_or_init($py, || {
$crate::conversion::ToPyObject::to_object($value, $py)
})
.as_ref($py)
}};
}