Skip to content

Commit b5e11cf

Browse files
youknowoneyouknowone
committed
Use write_bytes for lock reinit after fork
Replace transmute+AtomicCell::swap with ptr::write_bytes(0) for resetting parking_lot mutexes in _at_fork_reinit methods. Simpler and avoids relying on AtomicCell's internal layout matching the mutex type.
1 parent 04cf5da commit b5e11cf

3 files changed

Lines changed: 29 additions & 47 deletions

File tree

‎crates/common/src/lock.rs‎

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -68,43 +68,46 @@ pub type PyMappedRwLockWriteGuard<'a, T> = MappedRwLockWriteGuard<'a, RawRwLock,
6868

6969
// can add fn const_{mutex,rw_lock}() if necessary, but we probably won't need to
7070

71-
/// Reset a `PyMutex` to its initial (unlocked) state after `fork()`.
71+
/// Reset a lock to its initial (unlocked) state by zeroing its bytes.
7272
///
73-
/// After `fork()`, locks held by dead parent threads would deadlock in the
74-
/// child. This writes `RawMutex::INIT` via the `Mutex::raw()` accessor,
75-
/// bypassing the normal unlock path which may interact with parking_lot's
76-
/// internal waiter queues.
73+
/// After `fork()`, any lock held by a now-dead thread would remain
74+
/// permanently locked. We zero the raw bytes (the unlocked state for all
75+
/// `parking_lot` raw lock types) instead of using the normal unlock path,
76+
/// which would interact with stale waiter queues.
7777
///
7878
/// # Safety
7979
///
8080
/// Must only be called from the single-threaded child process immediately
8181
/// after `fork()`, before any other thread is created.
82+
/// The type `T` must represent the unlocked state as all-zero bytes
83+
/// (true for `parking_lot::RawMutex`, `RawRwLock`, `RawReentrantMutex`, etc.).
8284
#[cfg(unix)]
83-
pub unsafe fn reinit_mutex_after_fork<T: ?Sized>(mutex: &PyMutex<T>) {
84-
// Use Mutex::raw() to access the underlying lock without layout assumptions.
85-
// parking_lot::RawMutex (AtomicU8) and RawCellMutex (Cell<bool>) both
86-
// represent the unlocked state as all-zero bytes.
85+
pub unsafe fn zero_reinit_after_fork<T>(lock: *const T) {
8786
unsafe {
88-
let raw = mutex.raw() as *const RawMutex as *mut u8;
89-
core::ptr::write_bytes(raw, 0, core::mem::size_of::<RawMutex>());
87+
core::ptr::write_bytes(lock as *mut u8, 0, core::mem::size_of::<T>());
9088
}
9189
}
9290

93-
/// Reset a `PyRwLock` to its initial (unlocked) state after `fork()`.
91+
/// Reset a `PyMutex` after `fork()`. See [`zero_reinit_after_fork`].
9492
///
95-
/// Same rationale as [`reinit_mutex_after_fork`] — dead threads' read or
96-
/// write locks would cause permanent deadlock in the child.
93+
/// # Safety
94+
///
95+
/// Must only be called from the single-threaded child process immediately
96+
/// after `fork()`, before any other thread is created.
97+
#[cfg(unix)]
98+
pub unsafe fn reinit_mutex_after_fork<T: ?Sized>(mutex: &PyMutex<T>) {
99+
unsafe { zero_reinit_after_fork(mutex.raw()) }
100+
}
101+
102+
/// Reset a `PyRwLock` after `fork()`. See [`zero_reinit_after_fork`].
97103
///
98104
/// # Safety
99105
///
100106
/// Must only be called from the single-threaded child process immediately
101107
/// after `fork()`, before any other thread is created.
102108
#[cfg(unix)]
103109
pub unsafe fn reinit_rwlock_after_fork<T: ?Sized>(rwlock: &PyRwLock<T>) {
104-
unsafe {
105-
let raw = rwlock.raw() as *const RawRwLock as *mut u8;
106-
core::ptr::write_bytes(raw, 0, core::mem::size_of::<RawRwLock>());
107-
}
110+
unsafe { zero_reinit_after_fork(rwlock.raw()) }
108111
}
109112

110113
/// Reset a `PyThreadMutex` to its initial (unlocked, unowned) state after `fork()`.

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

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,12 +47,7 @@ mod lock {
4747
pub(crate) unsafe fn reinit_after_fork() {
4848
if IMP_LOCK.is_locked() && !IMP_LOCK.is_owned_by_current_thread() {
4949
// Held by a dead thread — reset to unlocked.
50-
// Same pattern as RLock::_at_fork_reinit in thread.rs.
51-
unsafe {
52-
let old: &crossbeam_utils::atomic::AtomicCell<RawRMutex> =
53-
core::mem::transmute(&IMP_LOCK);
54-
old.swap(RawRMutex::INIT);
55-
}
50+
unsafe { rustpython_common::lock::zero_reinit_after_fork(&IMP_LOCK) };
5651
}
5752
}
5853
}

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

Lines changed: 7 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@ pub(crate) mod _thread {
2323
sync::{Arc, Weak},
2424
};
2525
use core::{cell::RefCell, time::Duration};
26-
use crossbeam_utils::atomic::AtomicCell;
2726
use parking_lot::{
2827
RawMutex, RawThreadId,
2928
lock_api::{RawMutex as RawMutexT, RawMutexTimed, RawReentrantMutex},
@@ -152,15 +151,9 @@ pub(crate) mod _thread {
152151

153152
#[pymethod]
154153
fn _at_fork_reinit(&self, _vm: &VirtualMachine) -> PyResult<()> {
155-
// Reset the mutex to unlocked by directly writing the INIT value.
156-
// Do NOT call unlock() here — after fork(), unlock_slow() would
157-
// try to unpark stale waiters from dead parent threads.
158-
let new_mut = RawMutex::INIT;
159-
unsafe {
160-
let old_mutex: &AtomicCell<RawMutex> = core::mem::transmute(&self.mu);
161-
old_mutex.swap(new_mut);
162-
}
163-
154+
// Overwrite lock state to unlocked. Do NOT call unlock() here —
155+
// after fork(), unlock_slow() would try to unpark stale waiters.
156+
unsafe { rustpython_common::lock::zero_reinit_after_fork(&self.mu) };
164157
Ok(())
165158
}
166159

@@ -252,16 +245,10 @@ pub(crate) mod _thread {
252245

253246
#[pymethod]
254247
fn _at_fork_reinit(&self, _vm: &VirtualMachine) -> PyResult<()> {
255-
// Reset the reentrant mutex to unlocked by directly writing INIT.
256-
// Do NOT call unlock() — after fork(), the slow path would try
257-
// to unpark stale waiters from dead parent threads.
248+
// Overwrite lock state to unlocked. Do NOT call unlock() here —
249+
// after fork(), unlock_slow() would try to unpark stale waiters.
258250
self.count.store(0, core::sync::atomic::Ordering::Relaxed);
259-
let new_mut = RawRMutex::INIT;
260-
unsafe {
261-
let old_mutex: &AtomicCell<RawRMutex> = core::mem::transmute(&self.mu);
262-
old_mutex.swap(new_mut);
263-
}
264-
251+
unsafe { rustpython_common::lock::zero_reinit_after_fork(&self.mu) };
265252
Ok(())
266253
}
267254

@@ -1021,10 +1008,7 @@ pub(crate) mod _thread {
10211008
/// Reset a parking_lot::Mutex to unlocked state after fork.
10221009
#[cfg(unix)]
10231010
fn reinit_parking_lot_mutex<T: ?Sized>(mutex: &parking_lot::Mutex<T>) {
1024-
unsafe {
1025-
let raw = mutex.raw() as *const parking_lot::RawMutex as *mut u8;
1026-
core::ptr::write_bytes(raw, 0, core::mem::size_of::<parking_lot::RawMutex>());
1027-
}
1011+
unsafe { rustpython_common::lock::zero_reinit_after_fork(mutex.raw()) };
10281012
}
10291013

10301014
// Thread handle state enum

0 commit comments

Comments
 (0)