@@ -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" ) ) ]
630647mod 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