Skip to content

Commit

Permalink
Add feature controlling the global reference pool to enable avoiding …
Browse files Browse the repository at this point in the history
…its overhead.
  • Loading branch information
adamreichold committed Apr 18, 2024
1 parent b11174e commit 43d791e
Show file tree
Hide file tree
Showing 4 changed files with 43 additions and 14 deletions.
5 changes: 4 additions & 1 deletion Cargo.toml
Expand Up @@ -63,7 +63,7 @@ futures = "0.3.28"
pyo3-build-config = { path = "pyo3-build-config", version = "=0.21.2", features = ["resolve-config"] }

[features]
default = ["macros"]
default = ["macros", "reference-pool"]

# Enables support for `async fn` for `#[pyfunction]` and `#[pymethods]`.
experimental-async = ["macros", "pyo3-macros/experimental-async"]
Expand Down Expand Up @@ -106,6 +106,9 @@ auto-initialize = []
# Allows use of the deprecated "GIL Refs" APIs.
gil-refs = []

# Allows `Clone`ing references to Python objects without the GIL being held.
reference-pool = []

# Optimizes PyObject to Vec conversion and so on.
nightly = []

Expand Down
42 changes: 33 additions & 9 deletions src/gil.rs
Expand Up @@ -2,7 +2,9 @@

use crate::impl_::not_send::{NotSend, NOT_SEND};
use crate::{ffi, Python};
use parking_lot::{const_mutex, Mutex, Once};
use parking_lot::Once;
#[cfg(feature = "reference-pool")]
use parking_lot::{const_mutex, Mutex};
use std::cell::Cell;
#[cfg(debug_assertions)]
use std::cell::RefCell;
Expand Down Expand Up @@ -246,12 +248,14 @@ impl Drop for GILGuard {
// Vector of PyObject
type PyObjVec = Vec<NonNull<ffi::PyObject>>;

#[cfg(feature = "reference-pool")]
/// Thread-safe storage for objects which were inc_ref / dec_ref while the GIL was not held.
struct ReferencePool {
// .0 is INCREFs, .1 is DECREFs
pointer_ops: Mutex<(PyObjVec, PyObjVec)>,
}

#[cfg(feature = "reference-pool")]
impl ReferencePool {
const fn new() -> Self {
Self {
Expand Down Expand Up @@ -289,8 +293,10 @@ impl ReferencePool {
}
}

#[cfg(feature = "reference-pool")]
unsafe impl Sync for ReferencePool {}

#[cfg(feature = "reference-pool")]
static POOL: ReferencePool = ReferencePool::new();

/// A guard which can be used to temporarily release the GIL and restore on `Drop`.
Expand All @@ -315,6 +321,7 @@ impl Drop for SuspendGIL {
ffi::PyEval_RestoreThread(self.tstate);

// Update counts of PyObjects / Py that were cloned or dropped while the GIL was released.
#[cfg(feature = "reference-pool")]
POOL.update_counts(Python::assume_gil_acquired());
}
}
Expand Down Expand Up @@ -389,6 +396,7 @@ impl GILPool {
pub unsafe fn new() -> GILPool {
increment_gil_count();
// Update counts of PyObjects / Py that have been cloned or dropped since last acquisition
#[cfg(feature = "reference-pool")]
POOL.update_counts(Python::assume_gil_acquired());
GILPool {
start: OWNED_OBJECTS
Expand Down Expand Up @@ -451,7 +459,10 @@ pub unsafe fn register_incref(obj: NonNull<ffi::PyObject>) {
if gil_is_acquired() {
ffi::Py_INCREF(obj.as_ptr())
} else {
#[cfg(feature = "reference-pool")]
POOL.register_incref(obj);
#[cfg(not(feature = "reference-pool"))]
panic!("Cannot clone pointer into Python heap without the GIL being held.");
}
}

Expand All @@ -467,7 +478,10 @@ pub unsafe fn register_decref(obj: NonNull<ffi::PyObject>) {
if gil_is_acquired() {
ffi::Py_DECREF(obj.as_ptr())
} else {
#[cfg(feature = "reference-pool")]
POOL.register_decref(obj);
#[cfg(all(not(feature = "reference-pool"), debug_assertions))]
eprintln!("Leaking pointer into Python heap because the GIL is not held.");
}
}

Expand Down Expand Up @@ -520,10 +534,12 @@ fn decrement_gil_count() {
mod tests {
#[allow(deprecated)]
use super::GILPool;
use super::{gil_is_acquired, GIL_COUNT, OWNED_OBJECTS, POOL};
use super::{gil_is_acquired, GIL_COUNT, OWNED_OBJECTS};
#[cfg(feature = "reference-pool")]
use super::POOL;
use crate::types::any::PyAnyMethods;
use crate::{ffi, gil, PyObject, Python};
#[cfg(not(target_arch = "wasm32"))]
#[cfg(all(feature = "reference-pool", not(target_arch = "wasm32")))]
use parking_lot::{const_mutex, Condvar, Mutex};
use std::ptr::NonNull;

Expand All @@ -539,6 +555,7 @@ mod tests {
len
}

#[cfg(feature = "reference-pool")]
fn pool_inc_refs_does_not_contain(obj: &PyObject) -> bool {
!POOL
.pointer_ops
Expand All @@ -547,6 +564,7 @@ mod tests {
.contains(&unsafe { NonNull::new_unchecked(obj.as_ptr()) })
}

#[cfg(feature = "reference-pool")]
fn pool_dec_refs_does_not_contain(obj: &PyObject) -> bool {
!POOL
.pointer_ops
Expand All @@ -555,7 +573,7 @@ mod tests {
.contains(&unsafe { NonNull::new_unchecked(obj.as_ptr()) })
}

#[cfg(not(target_arch = "wasm32"))]
#[cfg(all(feature = "reference-pool", not(target_arch = "wasm32")))]
fn pool_dec_refs_contains(obj: &PyObject) -> bool {
POOL.pointer_ops
.lock()
Expand Down Expand Up @@ -632,20 +650,24 @@ mod tests {
let reference = obj.clone_ref(py);

assert_eq!(obj.get_refcnt(py), 2);
#[cfg(feature = "reference-pool")]
assert!(pool_inc_refs_does_not_contain(&obj));
#[cfg(feature = "reference-pool")]
assert!(pool_dec_refs_does_not_contain(&obj));

// With the GIL held, reference count will be decreased immediately.
drop(reference);

assert_eq!(obj.get_refcnt(py), 1);
#[cfg(feature = "reference-pool")]
assert!(pool_inc_refs_does_not_contain(&obj));
#[cfg(feature = "reference-pool")]
assert!(pool_dec_refs_does_not_contain(&obj));
});
}

#[test]
#[cfg(not(target_arch = "wasm32"))] // We are building wasm Python with pthreads disabled
#[cfg(all(feature = "reference-pool", not(target_arch = "wasm32")))] // We are building wasm Python with pthreads disabled
fn test_pyobject_drop_without_gil_doesnt_decrease_refcnt() {
let obj = Python::with_gil(|py| {
let obj = get_object(py);
Expand Down Expand Up @@ -729,6 +751,7 @@ mod tests {
}

#[test]
#[cfg_attr(not(feature = "reference-pool"), should_panic)]
fn test_allow_threads_updates_refcounts() {
Python::with_gil(|py| {
// Make a simple object with 1 reference
Expand Down Expand Up @@ -768,13 +791,13 @@ mod tests {
})
}

#[cfg(not(target_arch = "wasm32"))]
#[cfg(all(feature = "reference-pool", not(target_arch = "wasm32")))]
struct Event {
set: Mutex<bool>,
wait: Condvar,
}

#[cfg(not(target_arch = "wasm32"))]
#[cfg(all(feature = "reference-pool", not(target_arch = "wasm32")))]
impl Event {
const fn new() -> Self {
Self {
Expand All @@ -797,7 +820,7 @@ mod tests {
}

#[test]
#[cfg(not(target_arch = "wasm32"))] // We are building wasm Python with pthreads disabled
#[cfg(all(feature = "reference-pool", not(target_arch = "wasm32")))] // We are building wasm Python with pthreads disabled
fn test_clone_without_gil() {
use crate::{Py, PyAny};
use std::{sync::Arc, thread};
Expand Down Expand Up @@ -862,7 +885,7 @@ mod tests {
}

#[test]
#[cfg(not(target_arch = "wasm32"))] // We are building wasm Python with pthreads disabled
#[cfg(all(feature = "reference-pool", not(target_arch = "wasm32")))] // We are building wasm Python with pthreads disabled
fn test_clone_in_other_thread() {
use crate::Py;
use std::{sync::Arc, thread};
Expand Down Expand Up @@ -908,6 +931,7 @@ mod tests {
}

#[test]
#[cfg(feature = "reference-pool")]
fn test_update_counts_does_not_deadlock() {
// update_counts can run arbitrary Python code during Py_DECREF.
// if the locking is implemented incorrectly, it will deadlock.
Expand Down
8 changes: 4 additions & 4 deletions tests/test_class_basics.rs
Expand Up @@ -229,7 +229,7 @@ impl UnsendableChild {
}

fn test_unsendable<T: PyClass + 'static>() -> PyResult<()> {
let obj = Python::with_gil(|py| -> PyResult<_> {
let (obj, keep_obj_here) = Python::with_gil(|py| -> PyResult<_> {
let obj: Py<T> = PyType::new_bound::<T>(py).call1((5,))?.extract()?;

// Accessing the value inside this thread should not panic
Expand All @@ -241,10 +241,10 @@ fn test_unsendable<T: PyClass + 'static>() -> PyResult<()> {
.is_err();

assert!(!caught_panic);
Ok(obj)
})?;

let keep_obj_here = obj.clone();
let keep_obj_here = obj.clone();
Ok((obj, keep_obj_here))
})?;

let caught_panic = std::thread::spawn(move || {
// This access must panic
Expand Down
2 changes: 2 additions & 0 deletions tests/test_gc.rs
Expand Up @@ -445,6 +445,7 @@ impl DropDuringTraversal {
}
}

#[cfg(feature = "reference-pool")]
#[test]
fn drop_during_traversal_with_gil() {
let drop_called = Arc::new(AtomicBool::new(false));
Expand Down Expand Up @@ -476,6 +477,7 @@ fn drop_during_traversal_with_gil() {
assert!(drop_called.load(Ordering::Relaxed));
}

#[cfg(feature = "reference-pool")]
#[test]
fn drop_during_traversal_without_gil() {
let drop_called = Arc::new(AtomicBool::new(false));
Expand Down

0 comments on commit 43d791e

Please sign in to comment.