Skip to content

Commit 5cbebf4

Browse files
committed
Retain the original weakref stripe lock like CPython
Store a safe handle to the selected static weakref stripe instead of retaining and rehashing the referent address. Keep the same lock after referent clearing, matching CPython free-threaded weakrefs_lock. The nonthreaded handle is zero-sized and introduces no locking. Preserve lock selection, acquire/release ordering, fork reset, deferred GC publication, and callback-owner release outside the stripe. No new locks or unsafe code are needed. Assisted-by: OpenAI Codex:model version unavailable
1 parent 51e0ffe commit 5cbebf4

1 file changed

Lines changed: 55 additions & 23 deletions

File tree

‎crates/vm/src/object/core.rs‎

Lines changed: 55 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -593,29 +593,46 @@ mod weakref_lock {
593593

594594
static LOCKS: [AtomicU8; NUM_WEAKREF_LOCKS] = [const { AtomicU8::new(0) }; NUM_WEAKREF_LOCKS];
595595

596+
#[derive(Debug)]
597+
pub(super) struct WeakrefLock {
598+
lock: &'static AtomicU8,
599+
}
600+
596601
pub(super) struct WeakrefLockGuard {
597-
idx: usize,
602+
lock: &'static AtomicU8,
598603
}
599604

600605
impl Drop for WeakrefLockGuard {
601606
fn drop(&mut self) {
602-
LOCKS[self.idx].store(0, Ordering::Release);
607+
self.lock.store(0, Ordering::Release);
603608
}
604609
}
605610

606-
pub(super) fn lock(addr: usize) -> WeakrefLockGuard {
607-
let idx = (addr >> 4) % NUM_WEAKREF_LOCKS;
608-
loop {
609-
if LOCKS[idx]
610-
.compare_exchange_weak(0, 1, Ordering::Acquire, Ordering::Relaxed)
611-
.is_ok()
612-
{
613-
return WeakrefLockGuard { idx };
611+
impl WeakrefLock {
612+
pub(super) fn lock(&self) -> WeakrefLockGuard {
613+
loop {
614+
if self
615+
.lock
616+
.compare_exchange_weak(0, 1, Ordering::Acquire, Ordering::Relaxed)
617+
.is_ok()
618+
{
619+
return WeakrefLockGuard { lock: self.lock };
620+
}
621+
core::hint::spin_loop();
614622
}
615-
core::hint::spin_loop();
616623
}
617624
}
618625

626+
pub(super) fn select(addr: usize) -> WeakrefLock {
627+
WeakrefLock {
628+
lock: &LOCKS[(addr >> 4) % NUM_WEAKREF_LOCKS],
629+
}
630+
}
631+
632+
pub(super) fn lock(addr: usize) -> WeakrefLockGuard {
633+
select(addr).lock()
634+
}
635+
619636
/// Reset all weakref stripe locks after fork in child process.
620637
/// Locks held by parent threads would cause infinite spin in the child.
621638
#[cfg(all(unix, feature = "host_env"))]
@@ -628,14 +645,27 @@ mod weakref_lock {
628645

629646
#[cfg(not(feature = "threading"))]
630647
mod weakref_lock {
648+
#[derive(Debug)]
649+
pub(super) struct WeakrefLock;
650+
631651
pub(super) struct WeakrefLockGuard;
632652

633653
impl Drop for WeakrefLockGuard {
634654
fn drop(&mut self) {}
635655
}
636656

637-
pub(super) fn lock(_addr: usize) -> WeakrefLockGuard {
638-
WeakrefLockGuard
657+
impl WeakrefLock {
658+
pub(super) fn lock(&self) -> WeakrefLockGuard {
659+
WeakrefLockGuard
660+
}
661+
}
662+
663+
pub(super) fn select(_addr: usize) -> WeakrefLock {
664+
WeakrefLock
665+
}
666+
667+
pub(super) fn lock(addr: usize) -> WeakrefLockGuard {
668+
select(addr).lock()
639669
}
640670
}
641671

@@ -744,7 +774,7 @@ impl WeakRefList {
744774
let weak_payload = PyWeak {
745775
pointers: Pointers::new(),
746776
wr_object: Radium::new(ptr::null_mut()),
747-
referent_addr: obj as *const PyObject as usize,
777+
weakrefs_lock: weakref_lock::select(obj as *const PyObject as usize),
748778
callback: UnsafeCell::new(callback),
749779
hash: Radium::new(crate::common::hash::SENTINEL),
750780
};
@@ -1022,10 +1052,10 @@ pub struct PyWeak {
10221052
/// Direct pointer to the referent object, null when dead.
10231053
/// Equivalent to wr_object in PyWeakReference.
10241054
wr_object: PyAtomic<*mut PyObject>,
1025-
/// Immutable key for the original referent stripe, like CPython's
1026-
/// weakrefs_lock. Used only for locking, never dereferenced, so the same
1027-
/// stripe protects the callback even after wr_object becomes null.
1028-
referent_addr: usize,
1055+
/// Handle to the original static referent stripe, like CPython's
1056+
/// weakrefs_lock. The same stripe protects the callback even after
1057+
/// wr_object becomes null and the referent is freed.
1058+
weakrefs_lock: weakref_lock::WeakrefLock,
10291059
/// Owned callback edge, protected by the original referent stripe.
10301060
/// GC can retain this edge after unlinking the weakref.
10311061
callback: UnsafeCell<Option<PyObjectRef>>,
@@ -1072,15 +1102,15 @@ impl PyWeak {
10721102

10731103
/// Get the owned callback, which GC can retain after clearing the referent.
10741104
pub(crate) fn get_callback(&self) -> Option<PyObjectRef> {
1075-
let _lock = weakref_lock::lock(self.referent_addr);
1105+
let _lock = self.weakrefs_lock.lock();
10761106
// SAFETY: the original referent stripe protects the callback.
10771107
unsafe { &*self.callback.get() }.clone()
10781108
}
10791109

10801110
/// Unlink and steal the callback under the original referent stripe.
10811111
/// Return ownership so callers release it only after unlocking.
10821112
fn clear(&self) -> Option<PyObjectRef> {
1083-
let _lock = weakref_lock::lock(self.referent_addr);
1113+
let _lock = self.weakrefs_lock.lock();
10841114
// SAFETY: the original referent stripe protects both fields, including
10851115
// callbacks retained by GC after the weakref was already unlinked.
10861116
unsafe {
@@ -1091,7 +1121,7 @@ impl PyWeak {
10911121

10921122
/// _PyWeakref_ClearRef: unlink without releasing the callback.
10931123
fn clear_ref(&self) {
1094-
let _lock = weakref_lock::lock(self.referent_addr);
1124+
let _lock = self.weakrefs_lock.lock();
10951125
// SAFETY: the original referent stripe is held.
10961126
unsafe { self.clear_ref_lock_held() };
10971127
}
@@ -1132,7 +1162,7 @@ unsafe impl Traverse for PyWeak {
11321162
// Stripe critical sections never suspend or run Python, so a stopped
11331163
// thread cannot own this lock. Also synchronize gc.get_referents(),
11341164
// which traverses while other threads are running.
1135-
let _lock = weakref_lock::lock(self.referent_addr);
1165+
let _lock = self.weakrefs_lock.lock();
11361166
// SAFETY: the original referent stripe protects the callback.
11371167
unsafe { &*self.callback.get() }.traverse(tracer_fn);
11381168
}
@@ -3673,7 +3703,9 @@ mod tests {
36733703
PyWeak {
36743704
pointers: Pointers::new(),
36753705
wr_object: Radium::new(ptr::null_mut()),
3676-
referent_addr: target.as_object() as *const PyObject as usize,
3706+
weakrefs_lock: weakref_lock::select(
3707+
target.as_object() as *const PyObject as usize
3708+
),
36773709
callback: UnsafeCell::new(None),
36783710
hash: Radium::new(crate::common::hash::SENTINEL),
36793711
},

0 commit comments

Comments
 (0)