fix multithreaded access to freelist pyclasses (#4902)

* make FreeList explicitly a wrapper for *mut PyObject

* fix multithreaded access to freelist pyclasses

* add changelog entry

* use a GILOnceCell to initialize the freelist

* respond to code review

* skip test on wasm

---------

Co-authored-by: David Hewitt <[email protected]>

Author: ngoldbaum

newsfragments/4902.fixed.md

@@ -0,0 +1 @@
+* Fixed thread-unsafe implementation of freelist pyclasses on the free-threaded build.

pyo3-macros-backend/src/pyclass.rs

@@ -2395,15 +2395,13 @@ impl<'a> PyClassImplsBuilder<'a> {
             quote! {
                 impl #pyo3_path::impl_::pyclass::PyClassWithFreeList for #cls {
                     #[inline]
-                    fn get_free_list(py: #pyo3_path::Python<'_>) -> &mut #pyo3_path::impl_::freelist::FreeList<*mut #pyo3_path::ffi::PyObject> {
-                        static mut FREELIST: *mut #pyo3_path::impl_::freelist::FreeList<*mut #pyo3_path::ffi::PyObject> = 0 as *mut _;
-                        unsafe {
-                            if FREELIST.is_null() {
-                                FREELIST = ::std::boxed::Box::into_raw(::std::boxed::Box::new(
-                                    #pyo3_path::impl_::freelist::FreeList::with_capacity(#freelist)));
-                            }
-                            &mut *FREELIST
-                        }
+                    fn get_free_list(py: #pyo3_path::Python<'_>) -> &'static ::std::sync::Mutex<#pyo3_path::impl_::freelist::PyObjectFreeList> {
+                        static FREELIST: #pyo3_path::sync::GILOnceCell<::std::sync::Mutex<#pyo3_path::impl_::freelist::PyObjectFreeList>> = #pyo3_path::sync::GILOnceCell::new();
+                        // If there's a race to fill the cell, the object created
+                        // by the losing thread will be deallocated via RAII
+                        &FREELIST.get_or_init(py, || {
+                            ::std::sync::Mutex::new(#pyo3_path::impl_::freelist::PyObjectFreeList::with_capacity(#freelist))
+                        })
                     }
                 }
             }

src/impl_/freelist.rs

@@ -8,58 +8,64 @@
 //!
 //! [1]: https://en.wikipedia.org/wiki/Free_list
 
+use crate::ffi;
 use std::mem;
 
-/// Represents a slot of a [`FreeList`].
-pub enum Slot<T> {
+/// Represents a slot of a [`PyObjectFreeList`].
+enum PyObjectSlot {
     /// A free slot.
     Empty,
     /// An allocated slot.
-    Filled(T),
+    Filled(*mut ffi::PyObject),
 }
 
-/// A free allocation list.
+// safety: access is guarded by a per-pyclass mutex
+unsafe impl Send for PyObjectSlot {}
+
+/// A free allocation list for PyObject ffi pointers.
 ///
 /// See [the parent module](crate::impl_::freelist) for more details.
-pub struct FreeList<T> {
-    entries: Vec<Slot<T>>,
+pub struct PyObjectFreeList {
+    entries: Box<[PyObjectSlot]>,
     split: usize,
     capacity: usize,
 }
 
-impl<T> FreeList<T> {
-    /// Creates a new `FreeList` instance with specified capacity.
-    pub fn with_capacity(capacity: usize) -> FreeList<T> {
-        let entries = (0..capacity).map(|_| Slot::Empty).collect::<Vec<_>>();
+impl PyObjectFreeList {
+    /// Creates a new `PyObjectFreeList` instance with specified capacity.
+    pub fn with_capacity(capacity: usize) -> PyObjectFreeList {
+        let entries = (0..capacity)
+            .map(|_| PyObjectSlot::Empty)
+            .collect::<Box<[_]>>();
 
-        FreeList {
+        PyObjectFreeList {
             entries,
             split: 0,
             capacity,
         }
     }
 
     /// Pops the first non empty item.
-    pub fn pop(&mut self) -> Option<T> {
+    pub fn pop(&mut self) -> Option<*mut ffi::PyObject> {
         let idx = self.split;
         if idx == 0 {
             None
         } else {
-            match mem::replace(&mut self.entries[idx - 1], Slot::Empty) {
-                Slot::Filled(v) => {
+            match mem::replace(&mut self.entries[idx - 1], PyObjectSlot::Empty) {
+                PyObjectSlot::Filled(v) => {
                     self.split = idx - 1;
                     Some(v)
                 }
-                _ => panic!("FreeList is corrupt"),
+                _ => panic!("PyObjectFreeList is corrupt"),
             }
         }
     }
 
-    /// Inserts a value into the list. Returns `Some(val)` if the `FreeList` is full.
-    pub fn insert(&mut self, val: T) -> Option<T> {
+    /// Inserts a value into the list. Returns `Some(val)` if the `PyObjectFreeList` is full.
+    pub fn insert(&mut self, val: *mut ffi::PyObject) -> Option<*mut ffi::PyObject> {
         let next = self.split + 1;
         if next < self.capacity {
-            self.entries[self.split] = Slot::Filled(val);
+            self.entries[self.split] = PyObjectSlot::Filled(val);
             self.split = next;
             None
         } else {

src/impl_/pyclass.rs

@@ -2,7 +2,7 @@ use crate::{
     exceptions::{PyAttributeError, PyNotImplementedError, PyRuntimeError, PyValueError},
     ffi,
     impl_::{
-        freelist::FreeList,
+        freelist::PyObjectFreeList,
         pycell::{GetBorrowChecker, PyClassMutability, PyClassObjectLayout},
         pyclass_init::PyObjectInit,
         pymethods::{PyGetterDef, PyMethodDefType},
@@ -20,6 +20,7 @@ use std::{
     marker::PhantomData,
     os::raw::{c_int, c_void},
     ptr::NonNull,
+    sync::Mutex,
     thread,
 };
 
@@ -910,7 +911,7 @@ use super::{pycell::PyClassObject, pymethods::BoundRef};
 /// Do not implement this trait manually. Instead, use `#[pyclass(freelist = N)]`
 /// on a Rust struct to implement it.
 pub trait PyClassWithFreeList: PyClass {
-    fn get_free_list(py: Python<'_>) -> &mut FreeList<*mut ffi::PyObject>;
+    fn get_free_list(py: Python<'_>) -> &'static Mutex<PyObjectFreeList>;
 }
 
 /// Implementation of tp_alloc for `freelist` classes.
@@ -931,7 +932,9 @@ pub unsafe extern "C" fn alloc_with_freelist<T: PyClassWithFreeList>(
     // If this type is a variable type or the subtype is not equal to this type, we cannot use the
     // freelist
     if nitems == 0 && subtype == self_type {
-        if let Some(obj) = T::get_free_list(py).pop() {
+        let mut free_list = T::get_free_list(py).lock().unwrap();
+        if let Some(obj) = free_list.pop() {
+            drop(free_list);
             ffi::PyObject_Init(obj, subtype);
             return obj as _;
         }
@@ -951,7 +954,11 @@ pub unsafe extern "C" fn free_with_freelist<T: PyClassWithFreeList>(obj: *mut c_
         T::type_object_raw(Python::assume_gil_acquired()),
         ffi::Py_TYPE(obj)
     );
-    if let Some(obj) = T::get_free_list(Python::assume_gil_acquired()).insert(obj) {
+    let mut free_list = T::get_free_list(Python::assume_gil_acquired())
+        .lock()
+        .unwrap();
+    if let Some(obj) = free_list.insert(obj) {
+        drop(free_list);
         let ty = ffi::Py_TYPE(obj);
 
         // Deduce appropriate inverse of PyType_GenericAlloc

tests/test_gc.rs

@@ -36,6 +36,35 @@ fn class_with_freelist() {
     });
 }
 
+#[pyclass(freelist = 2)]
+#[cfg(not(target_arch = "wasm32"))]
+struct ClassWithFreelistAndData {
+    data: Option<usize>,
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+fn spin_freelist(py: Python<'_>, data: usize) {
+    for _ in 0..500 {
+        let inst1 = Py::new(py, ClassWithFreelistAndData { data: Some(data) }).unwrap();
+        let inst2 = Py::new(py, ClassWithFreelistAndData { data: Some(data) }).unwrap();
+        assert_eq!(inst1.borrow(py).data, Some(data));
+        assert_eq!(inst2.borrow(py).data, Some(data));
+    }
+}
+
+#[test]
+#[cfg(not(target_arch = "wasm32"))]
+fn multithreaded_class_with_freelist() {
+    std::thread::scope(|s| {
+        s.spawn(|| {
+            Python::with_gil(|py| spin_freelist(py, 12));
+        });
+        s.spawn(|| {
+            Python::with_gil(|py| spin_freelist(py, 0x4d3d3d3));
+        });
+    });
+}
+
 /// Helper function to create a pair of objects that can be used to test drops;
 /// the first object is a guard that records when it has been dropped, the second
 /// object is a check that can be used to assert that the guard has been dropped.