Add import lock (IMP_LOCK) reinit after fork

youknowone · youknowone · commit 2893430ee137 · 2026-03-02T20:55:59.000+09:00
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
