Skip to content

Commit

Permalink
Use a blocking approach to reduce allocation overhead of the lock-fre…
Browse files Browse the repository at this point in the history
…e reference pool
  • Loading branch information
adamreichold committed May 18, 2024
1 parent e2ce401 commit ddd9f10
Showing 1 changed file with 55 additions and 10 deletions.
65 changes: 55 additions & 10 deletions src/gil.rs
Original file line number Diff line number Diff line change
Expand Up @@ -282,7 +282,8 @@ impl Drop for GILGuard {
#[cfg(not(pyo3_disable_reference_pool))]
/// Thread-safe storage for objects which were dec_ref while the GIL was not held.
struct PendingDecRef {
obj: NonNull<ffi::PyObject>,
objs: [NonNull<ffi::PyObject>; 32],
cnt: usize,
next: *mut PendingDecRef,
}

Expand All @@ -291,17 +292,57 @@ static POOL: AtomicPtr<PendingDecRef> = AtomicPtr::new(null_mut());

#[cfg(not(pyo3_disable_reference_pool))]
fn enqueue_decref(obj: NonNull<ffi::PyObject>) {
let old_top = POOL.swap(null_mut(), Ordering::AcqRel);

let mut curr = old_top;

while !curr.is_null() {
// SAFETY: We have exclusive ownership of the full stack starting at `old_top`.
let val = unsafe { &mut *curr };

if val.cnt < val.objs.len() {
val.objs[val.cnt] = obj;
val.cnt += 1;

let mut last_next = &mut val.next;

while !last_next.is_null() {
last_next = unsafe { &mut (**last_next).next };
}

POOL.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |new_top| {
*last_next = new_top;
Some(old_top)
})
.unwrap();
return;
}

curr = val.next;
}

// There is no existing stack or it has no unused capacity remaining,
// hence we allocate a new block and prepend it locally before publishing.
let mut objs = [NonNull::dangling(); 32];
objs[0] = obj;

let val = PendingDecRef {
obj,
next: null_mut(),
objs,
cnt: 1,
next: old_top,
};

let top = Box::into_raw(Box::new(val));

let next = unsafe { &mut (*top).next };
// SAFETY: We just allocated `top` using `Box::new`.
let mut last_next = unsafe { &mut (*top).next };

while !last_next.is_null() {
last_next = unsafe { &mut (**last_next).next };
}

POOL.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |new_top| {
*next = new_top;
*last_next = new_top;
Some(top)
})
.unwrap();
Expand All @@ -314,7 +355,9 @@ fn update_counts(_py: Python<'_>) {
// SAFETY: Was enqueued using `Box::into_raw`.
let val = unsafe { Box::from_raw(top) };

unsafe { ffi::Py_DECREF(val.obj.as_ptr()) };
for obj in &val.objs[..val.cnt] {
unsafe { ffi::Py_DECREF(obj.as_ptr()) };
}

top = val.next;
}
Expand Down Expand Up @@ -581,11 +624,13 @@ mod tests {
// SAFETY: Was enqueued using `Box::into_raw`.
let val = unsafe { Box::from_raw(top) };

if val.obj.as_ptr() == obj.as_ptr() {
found = true;
}
for obj1 in &val.objs[..val.cnt] {
if obj1.as_ptr() == obj.as_ptr() {
found = true;
}

unsafe { ffi::Py_DECREF(val.obj.as_ptr()) };
unsafe { ffi::Py_DECREF(obj1.as_ptr()) };
}

top = val.next;
}
Expand Down

0 comments on commit ddd9f10

Please sign in to comment.