Add import lock (IMP_LOCK) reinit after fork

youknowone · youknowone · commit 9be0543b3688 · 2026-03-02T16:53:52.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,13 @@ 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
+
+/// Return the current thread's parking_lot thread ID.
+///
+/// This is the same ID stored in the `owner` field of `RawReentrantMutex`
+/// when the current thread holds it.
+#[cfg(all(unix, feature = "threading"))]
+pub fn current_thread_id() -> core::num::NonZeroUsize {
+    use lock_api::GetThreadId;
+    RawThreadId.nonzero_thread_id()
+}
diff --git a/crates/vm/src/stdlib/imp.rs b/crates/vm/src/stdlib/imp.rs
@@ -33,6 +33,41 @@ 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() {
+        use core::sync::atomic::{AtomicUsize, Ordering};
+
+        unsafe {
+            // RawReentrantMutex layout: owner: AtomicUsize at offset 0
+            let owner_ptr = &IMP_LOCK as *const RawRMutex as *const AtomicUsize;
+            let owner = (*owner_ptr).load(Ordering::Relaxed);
+
+            if owner != 0 {
+                let current = rustpython_common::lock::current_thread_id().get();
+                if owner != current {
+                    // Held by a dead thread — reset to unlocked
+                    let ptr = &IMP_LOCK as *const RawRMutex as *mut u8;
+                    core::ptr::write_bytes(ptr, 0, core::mem::size_of::<RawRMutex>());
+                }
+            }
+        }
+    }
+}
+
+/// 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
