Skip to content

Commit 3e8d703

Browse files
youknowoneyouknowone
committed
Replace force_unlock with reinit_*_after_fork
Replace all force_unlock() + try_lock() patterns with zero-based reinit that bypasses parking_lot internals entirely. After fork(), the child is single-threaded so reinited locks won't contend. Add reinit_rwlock_after_fork to common::lock alongside the existing reinit_mutex_after_fork. Replace force_unlock_after_fork methods in codecs, intern, and gc_state with reinit_after_fork equivalents. This fixes after_fork_child silently dropping thread handles when try_lock() failed on per-handle Arc<Mutex> locks.
1 parent ce2be2a commit 3e8d703

File tree

7 files changed

+132
-180
lines changed

7 files changed

+132
-180
lines changed

‎crates/common/src/lock.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,3 +79,20 @@ pub unsafe fn reinit_mutex_after_fork<T: ?Sized>(mutex: &PyMutex<T>) {
7979
core::ptr::write_bytes(raw, 0, core::mem::size_of::<RawMutex>());
8080
}
8181
}
82+
83+
/// Reset a `PyRwLock` to its initial (unlocked) state after `fork()`.
84+
///
85+
/// Same rationale as [`reinit_mutex_after_fork`] — dead threads' read or
86+
/// write locks would cause permanent deadlock in the child.
87+
///
88+
/// # Safety
89+
///
90+
/// Must only be called from the single-threaded child process immediately
91+
/// after `fork()`, before any other thread is created.
92+
#[cfg(unix)]
93+
pub unsafe fn reinit_rwlock_after_fork<T: ?Sized>(rwlock: &PyRwLock<T>) {
94+
unsafe {
95+
let raw = rwlock.raw() as *const RawRwLock as *mut u8;
96+
core::ptr::write_bytes(raw, 0, core::mem::size_of::<RawRwLock>());
97+
}
98+
}

‎crates/vm/src/codecs.rs‎

Lines changed: 5 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -153,29 +153,14 @@ impl ToPyObject for PyCodec {
153153
}
154154

155155
impl CodecsRegistry {
156-
/// Force-unlock the inner RwLock after fork in the child process.
156+
/// Reset the inner RwLock to unlocked state after fork().
157157
///
158158
/// # Safety
159159
/// Must only be called after fork() in the child process when no other
160-
/// threads exist. The calling thread must NOT hold this lock.
161-
#[cfg(all(unix, feature = "host_env"))]
162-
pub(crate) unsafe fn force_unlock_after_fork(&self) {
163-
if self.inner.try_write().is_some() {
164-
return;
165-
}
166-
let is_shared = self.inner.try_read().is_some();
167-
if is_shared {
168-
loop {
169-
// SAFETY: Lock is shared-locked by dead thread(s).
170-
unsafe { self.inner.force_unlock_read() };
171-
if self.inner.try_write().is_some() {
172-
return;
173-
}
174-
}
175-
} else {
176-
// SAFETY: Lock is exclusively locked by a dead thread.
177-
unsafe { self.inner.force_unlock_write() };
178-
}
160+
/// threads exist.
161+
#[cfg(all(unix, feature = "threading"))]
162+
pub(crate) unsafe fn reinit_after_fork(&self) {
163+
unsafe { crate::common::lock::reinit_rwlock_after_fork(&self.inner) };
179164
}
180165

181166
pub(crate) fn new(ctx: &Context) -> Self {

‎crates/vm/src/gc_state.rs‎

Lines changed: 23 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -85,16 +85,14 @@ impl GcGeneration {
8585
guard.uncollectable += uncollectable;
8686
}
8787

88-
/// Force-unlock the stats mutex after fork() in the child process.
88+
/// Reset the stats mutex to unlocked state after fork().
8989
///
9090
/// # Safety
9191
/// Must only be called after fork() in the child process when no other
9292
/// threads exist.
93-
#[cfg(unix)]
94-
unsafe fn force_unlock_stats_after_fork(&self) {
95-
if self.stats.try_lock().is_none() {
96-
unsafe { self.stats.force_unlock() };
97-
}
93+
#[cfg(all(unix, feature = "threading"))]
94+
unsafe fn reinit_stats_after_fork(&self) {
95+
unsafe { crate::common::lock::reinit_mutex_after_fork(&self.stats) };
9896
}
9997
}
10098

@@ -793,65 +791,35 @@ impl GcState {
793791

794792
/// Force-unlock all locks after fork() in the child process.
795793
///
794+
/// Reset all locks to unlocked state after fork().
795+
///
796796
/// After fork(), only the forking thread survives. Any lock held by another
797-
/// thread is permanently stuck. This method releases all such stuck locks.
797+
/// thread is permanently stuck. This resets them by zeroing the raw bytes.
798798
///
799799
/// # Safety
800800
/// Must only be called after fork() in the child process when no other
801801
/// threads exist. The calling thread must NOT hold any of these locks.
802-
#[cfg(unix)]
803-
pub unsafe fn force_unlock_after_fork(&self) {
804-
// Force-unlock the collecting mutex
805-
if self.collecting.try_lock().is_none() {
806-
unsafe { self.collecting.force_unlock() };
807-
}
802+
#[cfg(all(unix, feature = "threading"))]
803+
pub unsafe fn reinit_after_fork(&self) {
804+
use crate::common::lock::{reinit_mutex_after_fork, reinit_rwlock_after_fork};
808805

809-
// Force-unlock garbage and callbacks mutexes
810-
if self.garbage.try_lock().is_none() {
811-
unsafe { self.garbage.force_unlock() };
812-
}
813-
if self.callbacks.try_lock().is_none() {
814-
unsafe { self.callbacks.force_unlock() };
815-
}
806+
unsafe {
807+
reinit_mutex_after_fork(&self.collecting);
808+
reinit_mutex_after_fork(&self.garbage);
809+
reinit_mutex_after_fork(&self.callbacks);
816810

817-
// Force-unlock generation stats mutexes
818-
for generation in &self.generations {
819-
unsafe { generation.force_unlock_stats_after_fork() };
820-
}
821-
unsafe { self.permanent.force_unlock_stats_after_fork() };
822-
823-
// Force-unlock RwLocks
824-
for rw in &self.generation_objects {
825-
unsafe { force_unlock_rwlock_after_fork(rw) };
826-
}
827-
unsafe { force_unlock_rwlock_after_fork(&self.permanent_objects) };
828-
unsafe { force_unlock_rwlock_after_fork(&self.tracked_objects) };
829-
unsafe { force_unlock_rwlock_after_fork(&self.finalized_objects) };
830-
}
831-
}
811+
for generation in &self.generations {
812+
generation.reinit_stats_after_fork();
813+
}
814+
self.permanent.reinit_stats_after_fork();
832815

833-
/// Force-unlock a PyRwLock after fork() in the child process.
834-
///
835-
/// # Safety
836-
/// Must only be called after fork() in the child process when no other
837-
/// threads exist. The calling thread must NOT hold this lock.
838-
#[cfg(unix)]
839-
unsafe fn force_unlock_rwlock_after_fork<T>(lock: &PyRwLock<T>) {
840-
if lock.try_write().is_some() {
841-
return;
842-
}
843-
let is_shared = lock.try_read().is_some();
844-
if is_shared {
845-
loop {
846-
// SAFETY: Lock is shared-locked by dead thread(s).
847-
unsafe { lock.force_unlock_read() };
848-
if lock.try_write().is_some() {
849-
return;
816+
for rw in &self.generation_objects {
817+
reinit_rwlock_after_fork(rw);
850818
}
819+
reinit_rwlock_after_fork(&self.permanent_objects);
820+
reinit_rwlock_after_fork(&self.tracked_objects);
821+
reinit_rwlock_after_fork(&self.finalized_objects);
851822
}
852-
} else {
853-
// SAFETY: Lock is exclusively locked by a dead thread.
854-
unsafe { lock.force_unlock_write() };
855823
}
856824
}
857825

‎crates/vm/src/intern.rs‎

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -31,30 +31,14 @@ impl Clone for StringPool {
3131
}
3232

3333
impl StringPool {
34-
/// Force-unlock the inner RwLock after fork in the child process.
34+
/// Reset the inner RwLock to unlocked state after fork().
3535
///
3636
/// # Safety
3737
/// Must only be called after fork() in the child process when no other
38-
/// threads exist. The calling thread must NOT hold this lock.
39-
#[cfg(all(unix, feature = "host_env"))]
40-
pub(crate) unsafe fn force_unlock_after_fork(&self) {
41-
if self.inner.try_write().is_some() {
42-
return;
43-
}
44-
// Lock is stuck from a thread that no longer exists.
45-
let is_shared = self.inner.try_read().is_some();
46-
if is_shared {
47-
loop {
48-
// SAFETY: Lock is shared-locked by dead thread(s).
49-
unsafe { self.inner.force_unlock_read() };
50-
if self.inner.try_write().is_some() {
51-
return;
52-
}
53-
}
54-
} else {
55-
// SAFETY: Lock is exclusively locked by a dead thread.
56-
unsafe { self.inner.force_unlock_write() };
57-
}
38+
/// threads exist.
39+
#[cfg(all(unix, feature = "threading"))]
40+
pub(crate) unsafe fn reinit_after_fork(&self) {
41+
unsafe { crate::common::lock::reinit_rwlock_after_fork(&self.inner) };
5842
}
5943

6044
#[inline]

‎crates/vm/src/stdlib/posix.rs‎

Lines changed: 43 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -713,37 +713,22 @@ pub mod module {
713713
}
714714

715715
fn py_os_after_fork_child(vm: &VirtualMachine) {
716-
// Reset low-level state before any Python code runs in the child.
717-
// Signal triggers from the parent must not fire in the child.
716+
// Phase 1: Reset all internal locks FIRST.
717+
// After fork(), locks held by dead parent threads would deadlock
718+
// if we try to acquire them. This must happen before anything else.
719+
#[cfg(feature = "threading")]
720+
reinit_locks_after_fork(vm);
721+
722+
// Phase 2: Reset low-level atomic state (no locks needed).
718723
crate::signal::clear_after_fork();
719724
crate::stdlib::signal::_signal::clear_wakeup_fd_after_fork();
720725

721726
// Reset weakref stripe locks that may have been held during fork.
722727
#[cfg(feature = "threading")]
723728
crate::object::reset_weakref_locks_after_fork();
724729

725-
// Force-unlock all global VM locks that may have been held by
726-
// threads that no longer exist in the child process after fork.
727-
// SAFETY: After fork, only the forking thread survives. Any lock
728-
// held by another thread is permanently stuck. The forking thread
729-
// does not hold these locks during fork() (a high-level Python op).
730-
unsafe {
731-
vm.ctx.string_pool.force_unlock_after_fork();
732-
vm.state.codec_registry.force_unlock_after_fork();
733-
force_unlock_mutex_after_fork(&vm.state.atexit_funcs);
734-
force_unlock_mutex_after_fork(&vm.state.before_forkers);
735-
force_unlock_mutex_after_fork(&vm.state.after_forkers_child);
736-
force_unlock_mutex_after_fork(&vm.state.after_forkers_parent);
737-
force_unlock_mutex_after_fork(&vm.state.global_trace_func);
738-
force_unlock_mutex_after_fork(&vm.state.global_profile_func);
739-
crate::gc_state::gc_state().force_unlock_after_fork();
740-
741-
// Import lock (ReentrantMutex) — was previously not reinit'd
742-
#[cfg(feature = "threading")]
743-
crate::stdlib::imp::reinit_imp_lock_after_fork();
744-
}
745-
746-
// Mark all other threads as done before running Python callbacks
730+
// Phase 3: Clean up thread state. Locks are now reinit'd so we can
731+
// acquire them normally instead of using try_lock().
747732
#[cfg(feature = "threading")]
748733
crate::stdlib::thread::after_fork_child(vm);
749734

@@ -752,18 +737,45 @@ pub mod module {
752737
vm.signal_handlers
753738
.get_or_init(crate::signal::new_signal_handlers);
754739

740+
// Phase 4: Run Python-level at-fork callbacks.
755741
let after_forkers_child: Vec<PyObjectRef> = vm.state.after_forkers_child.lock().clone();
756742
run_at_forkers(after_forkers_child, false, vm);
757743
}
758744

759-
/// Force-unlock a PyMutex if held by a dead thread after fork.
745+
/// Reset all parking_lot-based locks in the interpreter state after fork().
760746
///
761-
/// # Safety
762-
/// Must only be called after fork() in the child process.
763-
unsafe fn force_unlock_mutex_after_fork<T>(mutex: &crate::common::lock::PyMutex<T>) {
764-
if mutex.try_lock().is_none() {
765-
// SAFETY: Lock is held by a dead thread after fork.
766-
unsafe { mutex.force_unlock() };
747+
/// After fork(), only the calling thread survives. Any locks held by other
748+
/// (now-dead) threads would cause deadlocks. We unconditionally reset them
749+
/// to unlocked by zeroing the raw lock bytes.
750+
#[cfg(all(unix, feature = "threading"))]
751+
fn reinit_locks_after_fork(vm: &VirtualMachine) {
752+
use rustpython_common::lock::reinit_mutex_after_fork;
753+
754+
unsafe {
755+
// PyGlobalState PyMutex locks
756+
reinit_mutex_after_fork(&vm.state.before_forkers);
757+
reinit_mutex_after_fork(&vm.state.after_forkers_child);
758+
reinit_mutex_after_fork(&vm.state.after_forkers_parent);
759+
reinit_mutex_after_fork(&vm.state.atexit_funcs);
760+
reinit_mutex_after_fork(&vm.state.global_trace_func);
761+
reinit_mutex_after_fork(&vm.state.global_profile_func);
762+
763+
// PyGlobalState parking_lot::Mutex locks
764+
reinit_mutex_after_fork(&vm.state.thread_frames);
765+
reinit_mutex_after_fork(&vm.state.thread_handles);
766+
reinit_mutex_after_fork(&vm.state.shutdown_handles);
767+
768+
// Context-level RwLock
769+
vm.ctx.string_pool.reinit_after_fork();
770+
771+
// Codec registry RwLock
772+
vm.state.codec_registry.reinit_after_fork();
773+
774+
// GC state (multiple Mutex + RwLock)
775+
crate::gc_state::gc_state().reinit_after_fork();
776+
777+
// Import lock (RawReentrantMutex<RawMutex, RawThreadId>)
778+
crate::stdlib::imp::reinit_imp_lock_after_fork();
767779
}
768780
}
769781

0 commit comments

Comments
 (0)