Fix _at_fork_reinit to write INIT directly instead of calling unlock() (#7312)

youknowone · web-flow · commit 9a12a8d532db · 2026-03-03T09:08:12.000+09:00
* Fix _at_fork_reinit to write INIT directly instead of calling unlock()

unlock() goes through unlock_slow() which accesses parking_lot's
global hash table to unpark waiters. After fork(), this hash table
contains stale entries from dead parent threads, making unlock_slow()
unsafe. Writing INIT directly bypasses parking_lot internals entirely.

* Add import lock (IMP_LOCK) reinit after fork

The import lock is a ReentrantMutex that was never reinit'd after
fork(). If a parent thread held it during fork, the child would
deadlock on any import. Only reset if the owner is a dead thread;
if the surviving thread held it, normal unlock still works.
diff --git a/crates/common/src/lock.rs b/crates/common/src/lock.rs
@@ -57,3 +57,24 @@ pub type PyRwLockWriteGuard<'a, T> = RwLockWriteGuard<'a, RawRwLock, T>;
 pub type PyMappedRwLockWriteGuard<'a, T> = MappedRwLockWriteGuard<'a, RawRwLock, T>;
 
 // can add fn const_{mutex,rw_lock}() if necessary, but we probably won't need to
+
+/// Reset a `PyMutex` to its initial (unlocked) state after `fork()`.
+///
+/// After `fork()`, locks held by dead parent threads would deadlock in the
+/// child. This zeroes the raw lock bytes directly, bypassing the normal unlock
+/// path which may interact with parking_lot's internal waiter queues.
+///
+/// # Safety
+///
+/// Must only be called from the single-threaded child process immediately
+/// after `fork()`, before any other thread is created.
+#[cfg(unix)]
+pub unsafe fn reinit_mutex_after_fork<T: ?Sized>(mutex: &PyMutex<T>) {
+    // lock_api::Mutex<R, T> layout: raw R at offset 0, then UnsafeCell<T>.
+    // Zeroing R resets to unlocked for both parking_lot::RawMutex (AtomicU8)
+    // and RawCellMutex (Cell<bool>).
+    unsafe {
+        let ptr = mutex as *const PyMutex<T> as *mut u8;
+        core::ptr::write_bytes(ptr, 0, core::mem::size_of::<RawMutex>());
+    }
+}
diff --git a/crates/vm/src/stdlib/imp.rs b/crates/vm/src/stdlib/imp.rs
@@ -33,6 +33,34 @@ mod lock {
     fn lock_held(_vm: &VirtualMachine) -> bool {
         IMP_LOCK.is_locked()
     }
+
+    /// Reset import lock after fork() — only if held by a dead thread.
+    ///
+    /// `IMP_LOCK` is a reentrant mutex. If the *current* (surviving) thread
+    /// held it at fork time, the child must be able to release it normally.
+    /// Only reset if a now-dead thread was the owner.
+    ///
+    /// # Safety
+    ///
+    /// Must only be called from single-threaded child after fork().
+    #[cfg(unix)]
+    pub(crate) unsafe fn reinit_after_fork() {
+        if IMP_LOCK.is_locked() && !IMP_LOCK.is_owned_by_current_thread() {
+            // Held by a dead thread — reset to unlocked.
+            // Same pattern as RLock::_at_fork_reinit in thread.rs.
+            unsafe {
+                let old: &crossbeam_utils::atomic::AtomicCell<RawRMutex> =
+                    core::mem::transmute(&IMP_LOCK);
+                old.swap(RawRMutex::INIT);
+            }
+        }
+    }
+}
+
+/// Re-export for fork safety code in posix.rs
+#[cfg(all(unix, feature = "threading"))]
+pub(crate) unsafe fn reinit_imp_lock_after_fork() {
+    unsafe { lock::reinit_after_fork() }
 }
 
 #[cfg(not(feature = "threading"))]
diff --git a/crates/vm/src/stdlib/posix.rs b/crates/vm/src/stdlib/posix.rs
@@ -737,6 +737,10 @@ pub mod module {
             force_unlock_mutex_after_fork(&vm.state.global_trace_func);
             force_unlock_mutex_after_fork(&vm.state.global_profile_func);
             crate::gc_state::gc_state().force_unlock_after_fork();
+
+            // Import lock (ReentrantMutex) — was previously not reinit'd
+            #[cfg(feature = "threading")]
+            crate::stdlib::imp::reinit_imp_lock_after_fork();
         }
 
         // Mark all other threads as done before running Python callbacks
diff --git a/crates/vm/src/stdlib/thread.rs b/crates/vm/src/stdlib/thread.rs
@@ -152,15 +152,9 @@ pub(crate) mod _thread {
 
         #[pymethod]
         fn _at_fork_reinit(&self, _vm: &VirtualMachine) -> PyResult<()> {
-            if self.mu.is_locked() {
-                unsafe {
-                    self.mu.unlock();
-                };
-            }
-            // Casting to AtomicCell is as unsafe as CPython code.
-            // Using AtomicCell will prevent compiler optimizer move it to somewhere later unsafe place.
-            // It will be not under the cell anymore after init call.
-
+            // Reset the mutex to unlocked by directly writing the INIT value.
+            // Do NOT call unlock() here — after fork(), unlock_slow() would
+            // try to unpark stale waiters from dead parent threads.
             let new_mut = RawMutex::INIT;
             unsafe {
                 let old_mutex: &AtomicCell<RawMutex> = core::mem::transmute(&self.mu);
@@ -252,11 +246,9 @@ pub(crate) mod _thread {
 
         #[pymethod]
         fn _at_fork_reinit(&self, _vm: &VirtualMachine) -> PyResult<()> {
-            if self.mu.is_locked() {
-                unsafe {
-                    self.mu.unlock();
-                };
-            }
+            // Reset the reentrant mutex to unlocked by directly writing INIT.
+            // Do NOT call unlock() — after fork(), the slow path would try
+            // to unpark stale waiters from dead parent threads.
             self.count.store(0, core::sync::atomic::Ordering::Relaxed);
             let new_mut = RawRMutex::INIT;
             unsafe {
