1//! Synchronization mechanisms based on the Python GIL.
2use crate::{types::PyString, types::PyType, Py, PyErr, PyVisit, Python};
3use std::cell::UnsafeCell;
4
5/// Value with concurrent access protected by the GIL.
6///
7/// This is a synchronization primitive based on Python's global interpreter lock (GIL).
8/// It ensures that only one thread at a time can access the inner value via shared references.
9/// It can be combined with interior mutability to obtain mutable references.
10///
11/// # Example
12///
13/// Combining `GILProtected` with `RefCell` enables mutable access to static data:
14///
15/// ```
16/// # use pyo3::prelude::*;
17/// use pyo3::sync::GILProtected;
18/// use std::cell::RefCell;
19///
20/// static NUMBERS: GILProtected<RefCell<Vec<i32>>> = GILProtected::new(RefCell::new(Vec::new()));
21///
22/// Python::with_gil(|py| {
23/// NUMBERS.get(py).borrow_mut().push(42);
24/// });
25/// ```
26pub struct GILProtected<T> {
27 value: T,
28}
29
30impl<T> GILProtected<T> {
31 /// Place the given value under the protection of the GIL.
32 pub const fn new(value: T) -> Self {
33 Self { value }
34 }
35
36 /// Gain access to the inner value by giving proof of having acquired the GIL.
37 pub fn get<'py>(&'py self, _py: Python<'py>) -> &'py T {
38 &self.value
39 }
40
41 /// Gain access to the inner value by giving proof that garbage collection is happening.
42 pub fn traverse<'py>(&'py self, _visit: PyVisit<'py>) -> &'py T {
43 &self.value
44 }
45}
46
47unsafe impl<T> Sync for GILProtected<T> where T: Send {}
48
49/// A write-once cell similar to [`once_cell::OnceCell`](https://docs.rs/once_cell/latest/once_cell/).
50///
51/// Unlike `once_cell::sync` which blocks threads to achieve thread safety, this implementation
52/// uses the Python GIL to mediate concurrent access. This helps in cases where `once_cell` or
53/// `lazy_static`'s synchronization strategy can lead to deadlocks when interacting with the Python
54/// GIL. For an example, see [the FAQ section](https://pyo3.rs/latest/faq.html) of the guide.
55///
56/// Note that:
57/// 1) `get_or_init` and `get_or_try_init` do not protect against infinite recursion
58/// from reentrant initialization.
59/// 2) If the initialization function `f` provided to `get_or_init` (or `get_or_try_init`)
60/// temporarily releases the GIL (e.g. by calling `Python::import`) then it is possible
61/// for a second thread to also begin initializing the `GITOnceCell`. Even when this
62/// happens `GILOnceCell` guarantees that only **one** write to the cell ever occurs
63/// - this is treated as a race, other threads will discard the value they compute and
64/// return the result of the first complete computation.
65///
66/// # Examples
67///
68/// The following example shows how to use `GILOnceCell` to share a reference to a Python list
69/// between threads:
70///
71/// ```
72/// use pyo3::sync::GILOnceCell;
73/// use pyo3::prelude::*;
74/// use pyo3::types::PyList;
75///
76/// static LIST_CELL: GILOnceCell<Py<PyList>> = GILOnceCell::new();
77///
78/// pub fn get_shared_list(py: Python<'_>) -> &PyList {
79/// LIST_CELL
80/// .get_or_init(py, || PyList::empty(py).into())
81/// .as_ref(py)
82/// }
83/// # Python::with_gil(|py| assert_eq!(get_shared_list(py).len(), 0));
84/// ```
85pub struct GILOnceCell<T>(UnsafeCell<Option<T>>);
86
87// T: Send is needed for Sync because the thread which drops the GILOnceCell can be different
88// to the thread which fills it.
89unsafe impl<T: Send + Sync> Sync for GILOnceCell<T> {}
90unsafe impl<T: Send> Send for GILOnceCell<T> {}
91
92impl<T> GILOnceCell<T> {
93 /// Create a `GILOnceCell` which does not yet contain a value.
94 pub const fn new() -> Self {
95 Self(UnsafeCell::new(None))
96 }
97
98 /// Get a reference to the contained value, or `None` if the cell has not yet been written.
99 #[inline]
100 pub fn get(&self, _py: Python<'_>) -> Option<&T> {
101 // Safe because if the cell has not yet been written, None is returned.
102 unsafe { &*self.0.get() }.as_ref()
103 }
104
105 /// Get a reference to the contained value, initializing it if needed using the provided
106 /// closure.
107 ///
108 /// See the type-level documentation for detail on re-entrancy and concurrent initialization.
109 #[inline]
110 pub fn get_or_init<F>(&self, py: Python<'_>, f: F) -> &T
111 where
112 F: FnOnce() -> T,
113 {
114 if let Some(value) = self.get(py) {
115 return value;
116 }
117
118 match self.init(py, || Ok::<T, std::convert::Infallible>(f())) {
119 Ok(value) => value,
120 Err(void) => match void {},
121 }
122 }
123
124 /// Like `get_or_init`, but accepts a fallible initialization function. If it fails, the cell
125 /// is left uninitialized.
126 ///
127 /// See the type-level documentation for detail on re-entrancy and concurrent initialization.
128 #[inline]
129 pub fn get_or_try_init<F, E>(&self, py: Python<'_>, f: F) -> Result<&T, E>
130 where
131 F: FnOnce() -> Result<T, E>,
132 {
133 if let Some(value) = self.get(py) {
134 return Ok(value);
135 }
136
137 self.init(py, f)
138 }
139
140 #[cold]
141 fn init<F, E>(&self, py: Python<'_>, f: F) -> Result<&T, E>
142 where
143 F: FnOnce() -> Result<T, E>,
144 {
145 // Note that f() could temporarily release the GIL, so it's possible that another thread
146 // writes to this GILOnceCell before f() finishes. That's fine; we'll just have to discard
147 // the value computed here and accept a bit of wasted computation.
148 let value = f()?;
149 let _ = self.set(py, value);
150
151 Ok(self.get(py).unwrap())
152 }
153
154 /// Get the contents of the cell mutably. This is only possible if the reference to the cell is
155 /// unique.
156 pub fn get_mut(&mut self) -> Option<&mut T> {
157 self.0.get_mut().as_mut()
158 }
159
160 /// Set the value in the cell.
161 ///
162 /// If the cell has already been written, `Err(value)` will be returned containing the new
163 /// value which was not written.
164 pub fn set(&self, _py: Python<'_>, value: T) -> Result<(), T> {
165 // Safe because GIL is held, so no other thread can be writing to this cell concurrently.
166 let inner = unsafe { &mut *self.0.get() };
167 if inner.is_some() {
168 return Err(value);
169 }
170
171 *inner = Some(value);
172 Ok(())
173 }
174
175 /// Takes the value out of the cell, moving it back to an uninitialized state.
176 ///
177 /// Has no effect and returns None if the cell has not yet been written.
178 pub fn take(&mut self) -> Option<T> {
179 self.0.get_mut().take()
180 }
181
182 /// Consumes the cell, returning the wrapped value.
183 ///
184 /// Returns None if the cell has not yet been written.
185 pub fn into_inner(self) -> Option<T> {
186 self.0.into_inner()
187 }
188}
189
190impl GILOnceCell<Py<PyType>> {
191 /// Get a reference to the contained Python type, initializing it if needed.
192 ///
193 /// This is a shorthand method for `get_or_init` which imports the type from Python on init.
194 pub(crate) fn get_or_try_init_type_ref<'py>(
195 &'py self,
196 py: Python<'py>,
197 module_name: &str,
198 attr_name: &str,
199 ) -> Result<&'py PyType, PyErr> {
200 self.get_or_try_init(py, || py.import(module_name)?.getattr(attr_name)?.extract())
201 .map(|ty: &Py| ty.as_ref(py))
202 }
203}
204
205/// Interns `text` as a Python string and stores a reference to it in static storage.
206///
207/// A reference to the same Python string is returned on each invocation.
208///
209/// # Example: Using `intern!` to avoid needlessly recreating the same Python string
210///
211/// ```
212/// use pyo3::intern;
213/// # use pyo3::{pyfunction, types::PyDict, wrap_pyfunction, PyResult, Python};
214///
215/// #[pyfunction]
216/// fn create_dict(py: Python<'_>) -> PyResult<&PyDict> {
217/// let dict = PyDict::new(py);
218/// // 👇 A new `PyString` is created
219/// // for every call of this function.
220/// dict.set_item("foo", 42)?;
221/// Ok(dict)
222/// }
223///
224/// #[pyfunction]
225/// fn create_dict_faster(py: Python<'_>) -> PyResult<&PyDict> {
226/// let dict = PyDict::new(py);
227/// // 👇 A `PyString` is created once and reused
228/// // for the lifetime of the program.
229/// dict.set_item(intern!(py, "foo"), 42)?;
230/// Ok(dict)
231/// }
232/// #
233/// # Python::with_gil(|py| {
234/// # let fun_slow = wrap_pyfunction!(create_dict, py).unwrap();
235/// # let dict = fun_slow.call0().unwrap();
236/// # assert!(dict.contains("foo").unwrap());
237/// # let fun = wrap_pyfunction!(create_dict_faster, py).unwrap();
238/// # let dict = fun.call0().unwrap();
239/// # assert!(dict.contains("foo").unwrap());
240/// # });
241/// ```
242#[macro_export]
243macro_rules! intern {
244 ($py: expr, $text: expr) => {{
245 static INTERNED: $crate::sync::Interned = $crate::sync::Interned::new($text);
246 INTERNED.get($py)
247 }};
248}
249
250/// Implementation detail for `intern!` macro.
251#[doc(hidden)]
252pub struct Interned(&'static str, GILOnceCell<Py<PyString>>);
253
254impl Interned {
255 /// Creates an empty holder for an interned `str`.
256 pub const fn new(value: &'static str) -> Self {
257 Interned(value, GILOnceCell::new())
258 }
259
260 /// Gets or creates the interned `str` value.
261 #[inline]
262 pub fn get<'py>(&'py self, py: Python<'py>) -> &'py PyString {
263 self.1
264 .get_or_init(py, || PyString::intern(py, self.0).into())
265 .as_ref(py)
266 }
267}
268
269#[cfg(test)]
270mod tests {
271 use super::*;
272
273 use crate::types::PyDict;
274
275 #[test]
276 fn test_intern() {
277 Python::with_gil(|py| {
278 let foo1 = "foo";
279 let foo2 = intern!(py, "foo");
280 let foo3 = intern!(py, stringify!(foo));
281
282 let dict = PyDict::new(py);
283 dict.set_item(foo1, 42_usize).unwrap();
284 assert!(dict.contains(foo2).unwrap());
285 assert_eq!(
286 dict.get_item(foo3)
287 .unwrap()
288 .unwrap()
289 .extract::<usize>()
290 .unwrap(),
291 42
292 );
293 });
294 }
295
296 #[test]
297 fn test_once_cell() {
298 Python::with_gil(|py| {
299 let mut cell = GILOnceCell::new();
300
301 assert!(cell.get(py).is_none());
302
303 assert_eq!(cell.get_or_try_init(py, || Err(5)), Err(5));
304 assert!(cell.get(py).is_none());
305
306 assert_eq!(cell.get_or_try_init(py, || Ok::<_, ()>(2)), Ok(&2));
307 assert_eq!(cell.get(py), Some(&2));
308
309 assert_eq!(cell.get_or_try_init(py, || Err(5)), Ok(&2));
310
311 assert_eq!(cell.take(), Some(2));
312 assert_eq!(cell.into_inner(), None)
313 })
314 }
315}
316