diff --git a/.cspell.dict/rust-more.txt b/.cspell.dict/rust-more.txt index c8614f12e..60cd912d9 100644 --- a/.cspell.dict/rust-more.txt +++ b/.cspell.dict/rust-more.txt @@ -83,6 +83,8 @@ trai ulonglong unic unistd +unparked +unparking unraw unsync wasip1 diff --git a/Cargo.lock b/Cargo.lock index 1bdf0ca28..6dd22c89c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4113,6 +4113,7 @@ dependencies = [ "num_enum", "optional", "parking_lot", + "parking_lot_core", "paste", "psm", "rapidhash", diff --git a/Cargo.toml b/Cargo.toml index c2a0a7e74..47edef74d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -275,6 +275,7 @@ openssl = "0.10.80" openssl-sys = "0.9.110" optional = "0.5" parking_lot = "0.12.3" +parking_lot_core = "0.9.12" paste = "1.0.15" pem-rfc7468 = "1.0" pkcs8 = "0.11" diff --git a/crates/vm/Cargo.toml b/crates/vm/Cargo.toml index 862929c1a..7753e30cb 100644 --- a/crates/vm/Cargo.toml +++ b/crates/vm/Cargo.toml @@ -21,7 +21,7 @@ opcode-histogram = [] flame-it = ["flame", "flamer"] freeze-stdlib = ["encodings"] jit = ["rustpython-jit"] -threading = ["rustpython-common/threading"] +threading = ["rustpython-common/threading", "dep:parking_lot_core"] gc = [] compiler = ["parser", "codegen", "rustpython-compiler"] ast = ["ruff_python_ast", "ruff_text_size"] @@ -68,6 +68,7 @@ num-integer = { workspace = true } num-traits = { workspace = true } num_enum = { workspace = true } parking_lot = { workspace = true } +parking_lot_core = { workspace = true, optional = true } paste = { workspace = true } rapidhash = { workspace = true } scopeguard = { workspace = true } diff --git a/crates/vm/src/gc_state.rs b/crates/vm/src/gc_state.rs index 38d31aec7..70ee04d46 100644 --- a/crates/vm/src/gc_state.rs +++ b/crates/vm/src/gc_state.rs @@ -5,6 +5,8 @@ use crate::common::linked_list::LinkedList; use crate::common::lock::{PyMutex, PyRwLock}; use crate::object::{GC_NO_OWNER, GC_PERMANENT, GC_REACHABLE, GC_UNTRACKED, GcLink, GcOwner}; +#[cfg(feature = "threading")] +use crate::vm::stop_the_world::StopTheWorldGroup; use crate::{AsObject, PyObject, PyObjectRef}; use core::ptr::NonNull; use core::sync::atomic::{AtomicBool, AtomicU16, AtomicU32, AtomicUsize, Ordering}; @@ -191,114 +193,6 @@ type GcBuildHasher = core::hash::BuildHasherDefault; type GcSet = std::collections::HashSet; type GcMap = std::collections::HashMap; -/// RAII barrier that parks every other thread for the pointer-reading phases -/// of a collection and lets them run again before finalizers execute. -/// -/// Reference subtraction, the reachability walk and the strong-reference -/// snapshot dereference the interpreter state of every tracked object, -/// including the `localsplus` of frames that other threads are actively -/// executing. Those writes carry no synchronization, so the reads are only -/// well-defined while all other threads are parked at a safepoint. Restarting -/// happens explicitly once the snapshot has pinned every object; `Drop` is a -/// backstop that also restarts on the early-return paths. -/// -/// A collection acts on one interpreter's objects, but its candidates include -/// the ones no interpreter owns, which every interpreter can reference and so -/// incref. Reading a refcount that another interpreter is changing is what -/// makes an object look unreachable when it is not, so every live interpreter -/// is stopped, not just the collecting one. Stopping in `runtime` id order -/// is deterministic. One process-wide STW transaction covers the entire group, -/// preventing a single-interpreter requester from stopping a collector while -/// that collector waits for a different interpreter. -#[cfg(feature = "threading")] -struct CollectStopTheWorld { - /// Entire interpreter snapshot, retained through restart even on unwind. - /// Dropping the last state can release Python objects and run finalizers, - /// so both stopped and not-yet-stopped states must outlive the transaction - /// and the collection's generation locks. - states: Vec>, - stopped_count: usize, - /// Keeps interpreters from registering between the snapshot below and the - /// restart. One registered in that window would be missing from `states`, - /// so its bootstrap would keep running — and mutating the shared generation - /// lists — while this collection reads them. - admission: Option>, - /// Acquired before admission and retained until every state is restarted. - transaction: Option, - restarted: bool, -} - -#[cfg(feature = "threading")] -impl CollectStopTheWorld { - /// Request stop-the-world on every live interpreter when the current thread - /// has an attached VM. Falls back to no barrier when no VM is attached (the - /// tracked-object reads then run without other threads only if the caller - /// guarantees it). - fn new() -> Self { - // No attached VM means no interpreter is running Python on this thread; - // keep the historical no-barrier fallback. - if !crate::vm::thread::current_vm_is_set() { - return Self { - states: Vec::new(), - stopped_count: 0, - admission: None, - transaction: None, - restarted: true, - }; - } - - // Accumulate into a live `Self` rather than a bare Vec: if a later - // `stop_the_world` unwinds, dropping this guard restarts the - // interpreters already stopped, instead of leaving their threads parked - // and their exclusion held forever. - let transaction = crate::vm::StopTheWorldTransaction::acquire(); - let mut guard = Self { - states: Vec::new(), - stopped_count: 0, - admission: Some(crate::vm::runtime::lock_admission_for_stop()), - transaction: Some(transaction), - restarted: false, - }; - guard.states = crate::vm::runtime::live_interpreter_states(); - for state in &guard.states { - state - .stop_the_world - .stop_with_transaction(state, guard.transaction.as_ref().unwrap()); - guard.stopped_count += 1; - } - guard - } - - /// Restart the world. Idempotent. - fn restart(&mut self) { - if self.restarted { - return; - } - self.restarted = true; - // Reverse of the stop order. The references stay until this guard is - // dropped; see the field comment. - for state in self.states[..self.stopped_count].iter().rev() { - state.stop_the_world.start_the_world(state); - } - // Nothing is parked any more, so registration may resume. - self.admission = None; - self.transaction = None; - } - - /// Whether this collection actually stopped the world. - #[cfg(all(unix, debug_assertions))] - fn is_stopped(&self) -> bool { - self.stopped_count != 0 - } -} - -#[cfg(feature = "threading")] -impl Drop for CollectStopTheWorld { - fn drop(&mut self) { - self.restart(); - } -} - /// The process-wide object lists every interpreter's collections walk. /// /// Interpreter-owned policy and results live in [`GcInterpreterState`]; what is @@ -639,7 +533,7 @@ impl GcState { // Closing it fully would require making those locks stop-the-world // aware; the exclusion above only serializes the fork/GC requesters. #[cfg(feature = "threading")] - let mut stw = CollectStopTheWorld::new(); + let mut stw = StopTheWorldGroup::new(); // Step 1: Gather objects from generations 0..=generation // Hold read locks for the entire scan to prevent concurrent modifications. @@ -1048,7 +942,7 @@ impl GcState { // liveness check and weakref closure in the same stopped phase. // No callbacks, finalizers or tp_clear may run under this guard. #[cfg(feature = "threading")] - let mut clear_stw = CollectStopTheWorld::new(); + let mut clear_stw = StopTheWorldGroup::new(); // Untrack dead objects BEFORE clearing them, mirroring the // untrack-then-clear ordering of the refcount dealloc path. diff --git a/crates/vm/src/stdlib/posix.rs b/crates/vm/src/stdlib/posix.rs index 8e142c417..69cf11e3f 100644 --- a/crates/vm/src/stdlib/posix.rs +++ b/crates/vm/src/stdlib/posix.rs @@ -681,7 +681,12 @@ pub mod module { } } - fn py_os_before_fork(vm: &VirtualMachine) { + #[cfg(feature = "threading")] + use crate::vm::stop_the_world::ForkStopTheWorld; + #[cfg(not(feature = "threading"))] + struct ForkStopTheWorld; + + fn py_os_before_fork(vm: &VirtualMachine) -> ForkStopTheWorld { let before_forkers: Vec = vm.state.before_forkers.lock().clone(); // functions must be executed in reversed order as they are registered // only for before_forkers, refer: test_register_at_fork in test_posix @@ -692,10 +697,21 @@ pub mod module { crate::stdlib::_imp::acquire_imp_lock_for_fork(vm); #[cfg(feature = "threading")] - vm.state.stop_the_world.stop_the_world(&vm.state); + { + ForkStopTheWorld::new() + } + #[cfg(not(feature = "threading"))] + { + ForkStopTheWorld + } } - fn py_os_after_fork_child(vm: &VirtualMachine) { + fn py_os_after_fork_child(vm: &VirtualMachine, _stopped: &mut ForkStopTheWorld) { + #[cfg(feature = "threading")] + unsafe { + _stopped.begin_child_repair(); + } + // The interpreter registry is reachable from every thread, so repair it // before anything enumerates interpreters. #[cfg(all(unix, feature = "threading"))] @@ -711,20 +727,14 @@ pub mod module { // if we try to acquire them. This must happen before anything else. #[cfg(feature = "threading")] reinit_locks_after_fork(vm); + #[cfg(feature = "threading")] + _stopped.retain_thread_slots(&vm.state); // The collector stops every interpreter, so interpreters other than the // forking one must be repaired too; otherwise the child's first // collection waits for threads that did not survive the fork. #[cfg(all(unix, feature = "threading"))] - reinit_other_interpreters_after_fork(vm); - - // Reinit per-object IO buffer locks on std streams. - // BufferedReader/Writer/TextIOWrapper use PyThreadMutex which can be - // held by dead parent threads, causing deadlocks on any IO in the child. - #[cfg(feature = "threading")] - unsafe { - crate::stdlib::_io::reinit_std_streams_after_fork(vm) - }; + reinit_other_interpreters_after_fork(vm, _stopped); // Phase 2: Reset low-level atomic state (no locks needed). crate::signal::clear_after_fork(); @@ -774,6 +784,20 @@ pub mod module { vm.signal_handlers .get_or_init(crate::signal::SignalHandlers::default); + #[cfg(feature = "threading")] + _stopped.finish_child_repair(); + + // Reinit per-object IO buffer locks on std streams. + // BufferedReader/Writer/TextIOWrapper use PyThreadMutex which can be + // held by dead parent threads, causing deadlocks on any IO in the child. + // Lookup can call Python (for example sys.__getattr__), so core repair + // and transaction release must precede it. Such hooks must still avoid + // using inherited IO locks before this existing repair completes. + #[cfg(feature = "threading")] + unsafe { + crate::stdlib::_io::reinit_std_streams_after_fork(vm) + }; + // Phase 4: Run Python-level at-fork callbacks. let after_forkers_child: Vec = vm.state.after_forkers_child.lock().clone(); run_at_forkers(after_forkers_child, false, vm); @@ -831,7 +855,7 @@ pub mod module { /// # Safety /// Must only be called after `fork()` in the child, when no other threads exist. #[cfg(all(unix, feature = "threading"))] - fn reinit_other_interpreters_after_fork(vm: &VirtualMachine) { + fn reinit_other_interpreters_after_fork(vm: &VirtualMachine, stopped: &mut ForkStopTheWorld) { use rustpython_common::lock::reinit_mutex_after_fork; for state in crate::vm::runtime::live_interpreter_states() { @@ -857,6 +881,7 @@ pub mod module { } state.stop_the_world.reset_after_fork(); + stopped.retain_thread_slots(&state); // Every thread registered here belongs to the parent, including any // slot the forking thread itself registered before the fork. @@ -868,9 +893,9 @@ pub mod module { crate::vm::thread::purge_other_interpreter_slots_after_fork(vm.state.interpreter_id); } - fn py_os_after_fork_parent(vm: &VirtualMachine) { + fn py_os_after_fork_parent(vm: &VirtualMachine, _stopped: &mut ForkStopTheWorld) { #[cfg(feature = "threading")] - vm.state.stop_the_world.start_the_world(&vm.state); + _stopped.restart(); #[cfg(feature = "threading")] crate::stdlib::_imp::release_imp_lock_after_fork_parent(); @@ -955,25 +980,25 @@ pub mod module { vm.audit("os.fork", || ())?; - py_os_before_fork(vm); + let mut stopped = py_os_before_fork(vm); let pid = rustpython_host_env::posix::fork(); match pid { Ok(0) => { - py_os_after_fork_child(vm); + py_os_after_fork_child(vm, &mut stopped); Ok(0) } Ok(pid) => { // Capture this before parent after-fork hooks in case those // hooks start threads. let num_os_threads = get_number_of_os_threads(); - py_os_after_fork_parent(vm); + py_os_after_fork_parent(vm, &mut stopped); // Warn only after parent callback path resumes the world. warn_if_multi_threaded("fork", num_os_threads, vm)?; Ok(pid) } Err(err) => { - py_os_after_fork_parent(vm); + py_os_after_fork_parent(vm, &mut stopped); Err(err.into_pyexception(vm)) } } @@ -1000,22 +1025,22 @@ pub mod module { vm.audit("os.forkpty", || ())?; - py_os_before_fork(vm); + let mut stopped = py_os_before_fork(vm); let result = rustpython_host_env::posix::forkpty(); match result { Ok((0, master)) => { - py_os_after_fork_child(vm); + py_os_after_fork_child(vm, &mut stopped); Ok((0, master)) } Ok((pid, master)) => { let num_os_threads = get_number_of_os_threads(); - py_os_after_fork_parent(vm); + py_os_after_fork_parent(vm, &mut stopped); warn_if_multi_threaded("forkpty", num_os_threads, vm)?; Ok((pid, master)) } Err(err) => { - py_os_after_fork_parent(vm); + py_os_after_fork_parent(vm, &mut stopped); Err(err.into_pyexception(vm)) } } diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index ad748f63a..2d9f0d56e 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -18,6 +18,8 @@ mod python_run; pub use python_run::InteractiveLine; pub mod runtime; mod setting; +#[cfg(feature = "threading")] +pub(crate) mod stop_the_world; pub mod thread; mod vm_new; mod vm_object; @@ -644,9 +646,9 @@ impl StopTheWorldState { // Only one thread survives fork; any stop-the-world bit inherited // from the parent is stale. crate::signal::clear_stop_bit(); - // before_fork owns the sole transaction. Release it exactly once via - // that state's ownership flag; resetting other interpreters must not - // release a later transaction. No unconditional global reset is needed. + // A single-state stop owns the transaction through this flag. A fork + // group instead retains its shared transaction through child repair; + // its member states do not own exclusion and must not release it here. self.release_exclusion(); stw_trace(format_args!("reset-after-fork")); } diff --git a/crates/vm/src/vm/stop_the_world.rs b/crates/vm/src/vm/stop_the_world.rs new file mode 100644 index 000000000..a8b43c595 --- /dev/null +++ b/crates/vm/src/vm/stop_the_world.rs @@ -0,0 +1,175 @@ +/// RAII barrier that parks every other thread for the pointer-reading phases +/// of a collection and lets them run again before finalizers execute. +/// +/// Reference subtraction, the reachability walk and the strong-reference +/// snapshot dereference the interpreter state of every tracked object, +/// including the `localsplus` of frames that other threads are actively +/// executing. Those writes carry no synchronization, so the reads are only +/// well-defined while all other threads are parked at a safepoint. Restarting +/// happens explicitly once the snapshot has pinned every object; `Drop` is a +/// backstop that also restarts on the early-return paths. +/// +/// A collection acts on one interpreter's objects, but its candidates include +/// the ones no interpreter owns, which every interpreter can reference and so +/// incref. Reading a refcount that another interpreter is changing is what +/// makes an object look unreachable when it is not, so every live interpreter +/// is stopped, not just the collecting one. Stopping in `runtime` id order +/// is deterministic. One process-wide STW transaction covers the entire group, +/// preventing a single-interpreter requester from stopping a collector while +/// that collector waits for a different interpreter. +#[cfg(feature = "threading")] +pub(crate) struct StopTheWorldGroup { + /// Entire interpreter snapshot, retained through restart even on unwind. + /// Dropping the last state can release Python objects and run finalizers, + /// so both stopped and not-yet-stopped states must outlive the transaction + /// and the collection's generation locks. + states: Vec>, + stopped_count: usize, + /// Keeps interpreters from registering between the snapshot below and the + /// restart. One registered in that window would be missing from `states`, + /// so its bootstrap would keep running — and mutating the shared generation + /// lists — while this collection reads them. + admission: Option>, + /// Acquired before admission and retained until every state is restarted. + transaction: Option, + restarted: bool, +} + +#[cfg(feature = "threading")] +impl StopTheWorldGroup { + /// Request stop-the-world on every live interpreter when the current thread + /// has an attached VM. Falls back to no barrier when no VM is attached (the + /// tracked-object reads then run without other threads only if the caller + /// guarantees it). + pub(crate) fn new() -> Self { + // No attached VM means no interpreter is running Python on this thread; + // keep the historical no-barrier fallback. + if !crate::vm::thread::current_vm_is_set() { + return Self { + states: Vec::new(), + stopped_count: 0, + admission: None, + transaction: None, + restarted: true, + }; + } + + // Accumulate into a live `Self` rather than a bare Vec: if a later + // `stop_the_world` unwinds, dropping this guard restarts the + // interpreters already stopped, instead of leaving their threads parked + // and their exclusion held forever. + let transaction = crate::vm::StopTheWorldTransaction::acquire(); + let mut guard = Self { + states: Vec::new(), + stopped_count: 0, + admission: Some(crate::vm::runtime::lock_admission_for_stop()), + transaction: Some(transaction), + restarted: false, + }; + guard.states = crate::vm::runtime::live_interpreter_states(); + for state in &guard.states { + state + .stop_the_world + .stop_with_transaction(state, guard.transaction.as_ref().unwrap()); + guard.stopped_count += 1; + } + guard + } + + /// Restart the world. Idempotent. + pub(crate) fn restart(&mut self) { + if self.restarted { + return; + } + self.restarted = true; + // Reverse of the stop order. The references stay until this guard is + // dropped; see the field comment. + for state in self.states[..self.stopped_count].iter().rev() { + state.stop_the_world.start_the_world(state); + } + // Nothing is parked any more, so registration may resume. + self.admission = None; + self.transaction = None; + } + + /// Whether this collection actually stopped the world. + #[cfg(all(unix, debug_assertions))] + pub(crate) fn is_stopped(&self) -> bool { + self.stopped_count != 0 + } +} + +#[cfg(feature = "threading")] +impl Drop for StopTheWorldGroup { + fn drop(&mut self) { + self.restart(); + } +} + +/// Fork owns the same barrier as GC, with child-only repair and deferred drops. +#[cfg(all(unix, feature = "host_env"))] +pub(crate) struct ForkStopTheWorld { + // ThreadSlot owns Python exception/trace/profile references. Keep removed + // slots alive until repair has finished and the transaction is released. + retained_slots: Vec, + // Field order keeps interpreter states alive while deferred Python roots + // above are destroyed, after Drop has released the transaction. + group: StopTheWorldGroup, + child_repair_incomplete: bool, +} + +#[cfg(all(unix, feature = "host_env"))] +impl ForkStopTheWorld { + pub(crate) fn new() -> Self { + let group = StopTheWorldGroup::new(); + Self { + retained_slots: Vec::new(), + group, + child_repair_incomplete: false, + } + } + + pub(crate) fn restart(&mut self) { + self.group.restart(); + } + + /// # Safety + /// Call only in the single surviving child immediately after fork. + pub(crate) unsafe fn begin_child_repair(&mut self) { + self.child_repair_incomplete = true; + // Never restart inherited thread registries on child unwind. + self.group.restarted = true; + // The dependency's automatic atfork registration can race its first + // parker. Explicitly reset queues before even unlocking admission. + parking_lot_core::reinit_after_fork(); + // runtime::reinit_after_fork must not zero a mutex with a live guard. + self.group.admission = None; + crate::vm::thread::retain_cached_slots_after_fork(&mut self.retained_slots); + } + + /// Preserve registry-owned Python references before child cleanup clears it. + /// The registry's lock must already have been reset in the child. + pub(crate) fn retain_thread_slots(&mut self, state: &crate::vm::PyGlobalState) { + self.retained_slots + .extend(state.thread_frames.lock().values().cloned()); + } + + /// Repair must be complete before Python callbacks or deferred slot drops. + pub(crate) fn finish_child_repair(&mut self) { + self.group.transaction = None; + self.child_repair_incomplete = false; + } +} + +#[cfg(all(unix, feature = "host_env"))] +impl Drop for ForkStopTheWorld { + fn drop(&mut self) { + if self.child_repair_incomplete { + // Partial child repair cannot support Rust unwind recovery or + // Python destruction. Like a fatal child-setup failure, abort + // before destroying the retained interpreter/thread roots. + rustpython_host_env::os::abort(); + } + self.group.restart(); + } +} diff --git a/crates/vm/src/vm/thread.rs b/crates/vm/src/vm/thread.rs index 339218d79..3cf07c981 100644 --- a/crates/vm/src/vm/thread.rs +++ b/crates/vm/src/vm/thread.rs @@ -27,6 +27,9 @@ use itertools::Itertools; use std::collections::HashMap; use std::thread_local; +pub(crate) mod critical_section; +pub(crate) mod lock; + /// Thread states for stop-the-world support (`_Py_THREAD_*`). /// /// DETACHED: not executing Python bytecode (in native code, or idle) @@ -96,6 +99,9 @@ pub struct ThreadSlot { pub thread: std::thread::Thread, /// QSBR state for deferred memory reclamation. pub(crate) qsbr: Arc, + /// `tstate->critical_section`: tagged pointer to the top-most critical + /// section. Only the owning thread reads or writes it. + pub(crate) critical_section: AtomicUsize, } #[cfg(feature = "threading")] @@ -156,6 +162,12 @@ thread_local! { static INTERP_THREAD_SLOTS: RefCell> = RefCell::new(HashMap::new()); + /// Critical sections of the slots `purge_other_interpreter_slots_after_fork` + /// dropped, carried into the slot the next enter of that interpreter builds. + #[cfg(feature = "threading")] + static FORKED_CRITICAL_SECTIONS: RefCell> = + RefCell::new(HashMap::new()); + /// Current thread's slot for the currently entered interpreter. #[cfg(feature = "threading")] static CURRENT_THREAD_SLOT: RefCell> = const { RefCell::new(None) }; @@ -468,6 +480,11 @@ fn ensure_thread_slot(vm: &VirtualMachine) -> CurrentFrameSlot { stop_requested: core::sync::atomic::AtomicBool::new(false), thread: std::thread::current(), qsbr: crate::object::qsbr::QSBR.register(), + critical_section: AtomicUsize::new( + FORKED_CRITICAL_SECTIONS + .with(|forked| forked.borrow_mut().remove(&interp_id)) + .unwrap_or(0), + ), }); registry.insert(thread_id, new_slot.clone()); drop(registry); @@ -483,6 +500,13 @@ pub fn current_thread_slot() -> Option { CURRENT_THREAD_SLOT.with(|slot| slot.borrow().clone()) } +/// Keep cached/orphaned slots alive until child repair permits Python drops. +#[cfg(all(unix, feature = "threading", feature = "host_env"))] +pub(crate) fn retain_cached_slots_after_fork(retained: &mut Vec) { + INTERP_THREAD_SLOTS.with(|slots| retained.extend(slots.borrow().values().cloned())); + CURRENT_THREAD_SLOT.with(|slot| retained.extend(slot.borrow().iter().cloned())); +} + /// Make `slot` the current thread slot (and the cached top-frame pointer). #[cfg(feature = "threading")] fn set_current_thread_slot(slot: CurrentFrameSlot) { @@ -586,6 +610,7 @@ fn hang_thread() -> ! { pub fn hang_current_thread(state: &PyGlobalState) -> ! { CURRENT_THREAD_SLOT.with(|slot| { if let Some(s) = slot.borrow().as_ref() { + critical_section::suspend_all(s); let prev = s .state .swap(ThreadState::ShuttingDown as i32, Ordering::AcqRel); @@ -663,6 +688,11 @@ fn attach_thread(vm: &VirtualMachine) { } } }); + if let Some(s) = current_thread_slot() { + // Resume previous critical section. This acquires the lock(s) from the + // top-most critical section. + critical_section::resume(&s); + } // A stop-the-world may have been requested while this thread was detached. // Honoring it here (rather than only at the next bytecode safepoint) keeps // a thread doing rapid allow_threads calls from re-attaching and running @@ -679,6 +709,7 @@ fn attach_thread(vm: &VirtualMachine) { fn detach_thread() { CURRENT_THREAD_SLOT.with(|slot| { if let Some(s) = slot.borrow().as_ref() { + critical_section::suspend_all(s); match s.state.compare_exchange( ThreadState::Attached as i32, ThreadState::Detached as i32, @@ -828,10 +859,10 @@ pub fn suspend_if_needed(state: &PyGlobalState) { #[cold] fn do_suspend(state: &PyGlobalState) { let stw = &state.stop_the_world; - CURRENT_THREAD_SLOT.with(|slot| { + let reattached = CURRENT_THREAD_SLOT.with(|slot| { let borrowed = slot.borrow(); let Some(s) = borrowed.as_ref() else { - return; + return false; }; // Decide whether to park while holding the thread registry. Both edges @@ -845,6 +876,12 @@ fn do_suspend(state: &PyGlobalState) { let park = { let _registry = state.thread_frames.lock(); if stw.requested.load(Ordering::Acquire) { + // Release critical sections before publishing SUSPENDED. Only + // this thread moves its slot out of ATTACHED, except to + // SHUTTING_DOWN, which hangs it below. + if s.state.load(Ordering::Acquire) == ThreadState::Attached as i32 { + critical_section::suspend_all(s); + } Some(s.state.compare_exchange( ThreadState::Attached as i32, ThreadState::Suspended as i32, @@ -861,7 +898,7 @@ fn do_suspend(state: &PyGlobalState) { match park { None => { super::stw_trace(format_args!("suspend skip not-requested")); - return; + return false; } Some(Ok(_)) => { // Consumed this thread's stop request bit. @@ -871,13 +908,13 @@ fn do_suspend(state: &PyGlobalState) { Some(ThreadState::Detached) => { // Leaving VM; caller will re-check on next entry. super::stw_trace(format_args!("suspend skip DETACHED")); - return; + return false; } Some(ThreadState::Suspended) => { // Already parked by another path. s.stop_requested.store(false, Ordering::Release); super::stw_trace(format_args!("suspend skip already-suspended")); - return; + return false; } Some(ThreadState::ShuttingDown) => { s.stop_requested.store(false, Ordering::Release); @@ -886,7 +923,7 @@ fn do_suspend(state: &PyGlobalState) { } _ => { debug_assert!(false, "unexpected thread state in suspend: {state}"); - return; + return false; } }, } @@ -931,7 +968,11 @@ fn do_suspend(state: &PyGlobalState) { } s.stop_requested.store(false, Ordering::Release); super::stw_trace(format_args!("suspend resume -> ATTACHED")); + true }); + if reattached && let Some(s) = current_thread_slot() { + critical_section::resume(&s); + } } #[cfg(feature = "threading")] @@ -1126,6 +1167,12 @@ pub fn cleanup_current_thread_frames(vm: &VirtualMachine) { let slot_for_interp = INTERP_THREAD_SLOTS.with(|slots| slots.borrow_mut().remove(&interp_id)); let current_slot = CURRENT_THREAD_SLOT.with(|slot| slot.borrow().as_ref().cloned()); let slot_to_clean = slot_for_interp.or(current_slot); + debug_assert!( + slot_to_clean + .as_ref() + .is_none_or(|slot| slot.critical_section.load(Ordering::Relaxed) == 0), + "thread cleanup cannot abandon a critical section" + ); // A dying thread should not remain logically ATTACHED while its // thread-state slot is being removed. @@ -1222,6 +1269,9 @@ pub fn reinit_frame_slot_after_fork(vm: &VirtualMachine) { } }; let top_iframe_ptr = get_current_frame() as usize; + // The surviving thread keeps its critical sections. + let critical_section = + current_thread_slot().map_or(0, |slot| slot.critical_section.load(Ordering::Relaxed)); let new_slot = Arc::new(ThreadSlot { // The surviving child thread keeps executing its current frame chain. // Only publish heavy frames for signal safety. @@ -1237,6 +1287,7 @@ pub fn reinit_frame_slot_after_fork(vm: &VirtualMachine) { stop_requested: core::sync::atomic::AtomicBool::new(false), thread: std::thread::current(), qsbr: crate::object::qsbr::QSBR.register(), + critical_section: AtomicUsize::new(critical_section), }); FRAME_SLOT_CACHE.with(|cache| { #[cfg(unix)] @@ -1269,7 +1320,14 @@ pub fn reinit_frame_slot_after_fork(vm: &VirtualMachine) { #[cfg(feature = "threading")] pub fn purge_other_interpreter_slots_after_fork(keep_id: i64) { INTERP_THREAD_SLOTS.with(|slots| { - slots.borrow_mut().retain(|&id, _| id == keep_id); + slots.borrow_mut().retain(|&id, slot| { + let critical_section = slot.critical_section.load(Ordering::Relaxed); + if id != keep_id && critical_section != 0 { + FORKED_CRITICAL_SECTIONS + .with(|forked| forked.borrow_mut().insert(id, critical_section)); + } + id == keep_id + }); }); } diff --git a/crates/vm/src/vm/thread/critical_section.rs b/crates/vm/src/vm/thread/critical_section.rs new file mode 100644 index 000000000..19cc49e60 --- /dev/null +++ b/crates/vm/src/vm/thread/critical_section.rs @@ -0,0 +1,773 @@ +//! Critical sections over [`PyMutex`]es. +//! +//! Each thread state keeps a stack of the critical sections it has entered. +//! Detaching releases the mutexes of every section on the stack, and attaching +//! again reacquires only the top-most one; the others are reacquired as the +//! sections above them end. Code inside a section must therefore revalidate +//! its state after anything that can detach, such as calling Python. +//! +//! Without the `threading` feature critical sections are no-ops. + +use super::lock::PyMutex; +use crate::VirtualMachine; + +/// `Py_BEGIN_CRITICAL_SECTION_MUTEX(m)` ... `Py_END_CRITICAL_SECTION()` +#[cfg_attr( + not(all(test, feature = "threading")), + expect(dead_code, reason = "Removed when the profiler consumer is wired") +)] +pub(crate) fn with(vm: &VirtualMachine, m: &PyMutex, f: impl FnOnce() -> R) -> R { + #[cfg(feature = "threading")] + { + threaded::with(vm, m, f) + } + #[cfg(not(feature = "threading"))] + { + let _ = (vm, m); + f() + } +} + +/// `Py_BEGIN_CRITICAL_SECTION2_MUTEX(m1, m2)` ... `Py_END_CRITICAL_SECTION2()` +#[cfg_attr( + not(all(test, feature = "threading")), + expect(dead_code, reason = "Removed when the profiler consumer is wired") +)] +pub(crate) fn with2( + vm: &VirtualMachine, + m1: &PyMutex, + m2: &PyMutex, + f: impl FnOnce() -> R, +) -> R { + #[cfg(feature = "threading")] + { + threaded::with2(vm, m1, m2, f) + } + #[cfg(not(feature = "threading"))] + { + let _ = (vm, m1, m2); + f() + } +} + +#[cfg(feature = "threading")] +pub(super) use threaded::{resume, suspend_all}; + +#[cfg(feature = "threading")] +mod threaded { + use super::PyMutex; + use crate::VirtualMachine; + use crate::vm::thread::{self, ThreadSlot}; + use core::ptr; + use core::sync::atomic::{AtomicPtr, AtomicUsize, Ordering}; + + /// The section's mutexes are released; it is resumed when it becomes the + /// top-most section of an attached thread. + const INACTIVE: usize = 0x1; + /// The tagged pointer refers to a [`PyCriticalSection2`]. + const TWO_MUTEXES: usize = 0x2; + const MASK: usize = 0x3; + + /// `PyCriticalSection` + /// + /// Only the owning thread touches a section; the fields are atomics + /// because the thread state that links them is shared. + #[repr(C, align(4))] + struct PyCriticalSection { + /// Tagged pointer to the previous critical section. + prev: AtomicUsize, + /// Mutex held by this section. Null while it is being acquired, and + /// for a section that skipped locking. + mutex: AtomicPtr, + } + + /// `PyCriticalSection2` + #[repr(C)] + struct PyCriticalSection2 { + base: PyCriticalSection, + mutex2: AtomicPtr, + } + + impl PyCriticalSection { + const fn new() -> Self { + Self { + prev: AtomicUsize::new(0), + mutex: AtomicPtr::new(ptr::null_mut()), + } + } + } + + fn mutex_ptr(m: &PyMutex) -> *mut PyMutex { + ptr::from_ref(m).cast_mut() + } + + /// # Safety + /// A non-null `m` must be a mutex of a live critical section. + unsafe fn lock(m: *mut PyMutex) { + unsafe { &*m }.lock(); + } + + /// # Safety + /// A non-null `m` must be a mutex of a live critical section. + unsafe fn unlock(m: *mut PyMutex) { + unsafe { &*m }.unlock(); + } + + fn untag(tag: usize) -> *const PyCriticalSection { + (tag & !MASK) as *const PyCriticalSection + } + + /// `_PyCriticalSection_IsActive` + fn is_active(tag: usize) -> bool { + tag != 0 && tag & INACTIVE == 0 + } + + fn tstate() -> thread::CurrentFrameSlot { + thread::current_thread_slot().expect("critical section requires a thread state") + } + + /// The thread state of a thread entering a critical section, which must + /// be attached: a detached thread would push a section that attaching + /// again tries to resume while it is still held. + fn attached_tstate() -> thread::CurrentFrameSlot { + let tstate = tstate(); + debug_assert_eq!( + tstate.state.load(Ordering::Acquire), + thread::ThreadState::Attached as i32, + "critical section requires an attached thread state" + ); + tstate + } + + fn world_stopped(vm: &VirtualMachine) -> bool { + vm.state + .stop_the_world + .world_stopped + .load(Ordering::Acquire) + } + + pub(super) fn with(vm: &VirtualMachine, m: &PyMutex, f: impl FnOnce() -> R) -> R { + let c = PyCriticalSection::new(); + begin_mutex(vm, &attached_tstate(), &c, m); + // The thread state is looked up again at the end, because fork + // replaces the slot of the surviving thread. + let _end = scopeguard::guard(&c, |c| end(&tstate(), c)); + f() + } + + pub(super) fn with2( + vm: &VirtualMachine, + m1: &PyMutex, + m2: &PyMutex, + f: impl FnOnce() -> R, + ) -> R { + let c = PyCriticalSection2 { + base: PyCriticalSection::new(), + mutex2: AtomicPtr::new(ptr::null_mut()), + }; + begin2_mutex(vm, &attached_tstate(), &c, m1, m2); + let _end = scopeguard::guard(&c, |c| end2(&tstate(), c)); + f() + } + + /// `_PyCriticalSection_BeginMutex` + fn begin_mutex(vm: &VirtualMachine, tstate: &ThreadSlot, c: &PyCriticalSection, m: &PyMutex) { + if m.lock_fast() { + c.mutex.store(mutex_ptr(m), Ordering::Relaxed); + c.prev.store( + tstate.critical_section.load(Ordering::Relaxed), + Ordering::Relaxed, + ); + tstate + .critical_section + .store(ptr::from_ref(c) as usize, Ordering::Relaxed); + } else { + begin_slow(vm, tstate, c, m); + } + } + + /// `_PyCriticalSection_BeginSlow` + #[cold] + fn begin_slow(vm: &VirtualMachine, tstate: &ThreadSlot, c: &PyCriticalSection, m: &PyMutex) { + // As an optimisation for locking the same mutex recursively, skip + // locking if the mutex is currently locked by the top-most critical + // section. If the top-most critical section is a two-mutex critical + // section, then locking is skipped if either mutex is m. + let top = tstate.critical_section.load(Ordering::Relaxed); + if top != 0 { + // SAFETY: the stack links live sections of this thread. + let prev = unsafe { &*untag(top) }; + if prev.mutex.load(Ordering::Relaxed) == mutex_ptr(m) { + c.mutex.store(ptr::null_mut(), Ordering::Relaxed); + c.prev.store(0, Ordering::Relaxed); + return; + } + if top & TWO_MUTEXES != 0 { + // SAFETY: the tag marks a two-mutex section. + let prev2 = unsafe { &*untag(top).cast::() }; + if prev2.mutex2.load(Ordering::Relaxed) == mutex_ptr(m) { + c.mutex.store(ptr::null_mut(), Ordering::Relaxed); + c.prev.store(0, Ordering::Relaxed); + return; + } + } + } + // If the world is stopped, we don't need to acquire the lock because + // there are no other threads that could be accessing the object. + // Without this check, acquiring a critical section while the world is + // stopped could lead to a deadlock. + if world_stopped(vm) { + c.mutex.store(ptr::null_mut(), Ordering::Relaxed); + c.prev.store(0, Ordering::Relaxed); + return; + } + c.mutex.store(ptr::null_mut(), Ordering::Relaxed); + c.prev.store(top, Ordering::Relaxed); + tstate + .critical_section + .store(ptr::from_ref(c) as usize, Ordering::Relaxed); + + m.lock(); + c.mutex.store(mutex_ptr(m), Ordering::Relaxed); + } + + /// `_PyCriticalSection2_BeginMutex` + fn begin2_mutex( + vm: &VirtualMachine, + tstate: &ThreadSlot, + c: &PyCriticalSection2, + m1: &PyMutex, + m2: &PyMutex, + ) { + if ptr::eq(m1, m2) { + // If the two mutex arguments are the same, treat this as a + // critical section with a single mutex. + c.mutex2.store(ptr::null_mut(), Ordering::Relaxed); + begin_mutex(vm, tstate, &c.base, m1); + return; + } + + // Sort the mutexes so that the lower address is locked first. The + // exact order does not matter, but we need to acquire the mutexes in + // a consistent order to avoid lock ordering deadlocks. + let (m1, m2) = if mutex_ptr(m2) < mutex_ptr(m1) { + (m2, m1) + } else { + (m1, m2) + }; + + if m1.lock_fast() { + if m2.lock_fast() { + c.base.mutex.store(mutex_ptr(m1), Ordering::Relaxed); + c.mutex2.store(mutex_ptr(m2), Ordering::Relaxed); + c.base.prev.store( + tstate.critical_section.load(Ordering::Relaxed), + Ordering::Relaxed, + ); + tstate + .critical_section + .store(ptr::from_ref(c) as usize | TWO_MUTEXES, Ordering::Relaxed); + } else { + begin2_slow(vm, tstate, c, m1, m2, true); + } + } else { + begin2_slow(vm, tstate, c, m1, m2, false); + } + } + + /// `_PyCriticalSection2_BeginSlow` + #[cold] + fn begin2_slow( + vm: &VirtualMachine, + tstate: &ThreadSlot, + c: &PyCriticalSection2, + m1: &PyMutex, + m2: &PyMutex, + is_m1_locked: bool, + ) { + if world_stopped(vm) { + c.base.mutex.store(ptr::null_mut(), Ordering::Relaxed); + c.mutex2.store(ptr::null_mut(), Ordering::Relaxed); + c.base.prev.store(0, Ordering::Relaxed); + return; + } + // Same optimization as in begin_slow: skip locking when recursively + // acquiring the same locks. + let top = tstate.critical_section.load(Ordering::Relaxed); + if top != 0 && top & TWO_MUTEXES != 0 { + // SAFETY: the tag marks a live two-mutex section of this thread. + let prev2 = unsafe { &*untag(top).cast::() }; + debug_assert!(mutex_ptr(m1) < mutex_ptr(m2)); + debug_assert!( + prev2.base.mutex.load(Ordering::Relaxed) < prev2.mutex2.load(Ordering::Relaxed) + ); + if prev2.base.mutex.load(Ordering::Relaxed) == mutex_ptr(m1) + && prev2.mutex2.load(Ordering::Relaxed) == mutex_ptr(m2) + { + c.base.mutex.store(ptr::null_mut(), Ordering::Relaxed); + c.mutex2.store(ptr::null_mut(), Ordering::Relaxed); + c.base.prev.store(0, Ordering::Relaxed); + return; + } + } + c.base.mutex.store(ptr::null_mut(), Ordering::Relaxed); + c.mutex2.store(ptr::null_mut(), Ordering::Relaxed); + c.base.prev.store(top, Ordering::Relaxed); + tstate + .critical_section + .store(ptr::from_ref(c) as usize | TWO_MUTEXES, Ordering::Relaxed); + + if !is_m1_locked { + m1.lock(); + } + m2.lock(); + c.base.mutex.store(mutex_ptr(m1), Ordering::Relaxed); + c.mutex2.store(mutex_ptr(m2), Ordering::Relaxed); + } + + /// `_PyCriticalSection_Pop`: removes the top-most critical section from + /// the thread's stack. If the new top-most section is inactive, then it + /// is resumed. + fn pop(tstate: &ThreadSlot, c: &PyCriticalSection) { + let prev = c.prev.load(Ordering::Relaxed); + tstate.critical_section.store(prev, Ordering::Relaxed); + if prev & INACTIVE != 0 { + resume_top(tstate); + } + } + + /// `_PyCriticalSection_End` + fn end(tstate: &ThreadSlot, c: &PyCriticalSection) { + // If the mutex is null, we used the fast path in begin_slow for locks + // already held in the top-most critical section, and we shouldn't + // unlock or pop this critical section. + let m = c.mutex.load(Ordering::Relaxed); + if m.is_null() { + return; + } + // SAFETY: an active section holds its mutex. + unsafe { unlock(m) }; + pop(tstate, c); + } + + /// `_PyCriticalSection2_End` + fn end2(tstate: &ThreadSlot, c: &PyCriticalSection2) { + // If mutex1 is null, we used the fast path in either begin_slow or + // begin2_slow for mutexes that are already held, and mutex2 should + // also be null. + let m1 = c.base.mutex.load(Ordering::Relaxed); + if m1.is_null() { + debug_assert!(c.mutex2.load(Ordering::Relaxed).is_null()); + return; + } + let m2 = c.mutex2.load(Ordering::Relaxed); + // SAFETY: an active section holds its mutexes. + unsafe { + if !m2.is_null() { + unlock(m2); + } + unlock(m1); + } + pop(tstate, &c.base); + } + + /// `_PyCriticalSection_SuspendAll`: release all locks held by critical + /// sections. Called before the thread stops being ATTACHED. + pub(crate) fn suspend_all(tstate: &ThreadSlot) { + let mut tagptr = &tstate.critical_section; + loop { + let tag = tagptr.load(Ordering::Relaxed); + if !is_active(tag) { + break; + } + // SAFETY: the stack links live sections of this thread. + let c = unsafe { &*untag(tag) }; + let m = c.mutex.load(Ordering::Relaxed); + if !m.is_null() { + // SAFETY: an active section holds its mutexes. + unsafe { unlock(m) }; + if tag & TWO_MUTEXES != 0 { + // SAFETY: the tag marks a two-mutex section. + let c2 = unsafe { &*untag(tag).cast::() }; + let m2 = c2.mutex2.load(Ordering::Relaxed); + if !m2.is_null() { + // SAFETY: as above. + unsafe { unlock(m2) }; + } + } + } + + tagptr.store(tag | INACTIVE, Ordering::Relaxed); + tagptr = &c.prev; + } + } + + /// Resume the top-most critical section after the thread attached. + pub(crate) fn resume(tstate: &ThreadSlot) { + if tstate.critical_section.load(Ordering::Relaxed) != 0 { + resume_top(tstate); + } + } + + /// `_PyCriticalSection_Resume` + fn resume_top(tstate: &ThreadSlot) { + let p = tstate.critical_section.load(Ordering::Relaxed); + // SAFETY: the stack links live sections of this thread. + let c = unsafe { &*untag(p) }; + debug_assert!(!is_active(p)); + + let m1 = c.mutex.swap(ptr::null_mut(), Ordering::Relaxed); + let (c2, m2) = if p & TWO_MUTEXES != 0 { + // SAFETY: the tag marks a two-mutex section. + let c2 = unsafe { &*untag(p).cast::() }; + (Some(c2), c2.mutex2.swap(ptr::null_mut(), Ordering::Relaxed)) + } else { + (None, ptr::null_mut()) + }; + + // SAFETY: the mutexes belong to a live section that released them. + unsafe { + if !m1.is_null() { + lock(m1); + } + if !m2.is_null() { + lock(m2); + } + } + + c.mutex.store(m1, Ordering::Relaxed); + if let Some(c2) = c2 + && !m2.is_null() + { + c2.mutex2.store(m2, Ordering::Relaxed); + } + + let p = tstate.critical_section.load(Ordering::Relaxed); + tstate + .critical_section + .store(p & !INACTIVE, Ordering::Relaxed); + } +} + +#[cfg(test)] +#[cfg(feature = "threading")] +mod tests { + use super::*; + #[cfg(all(unix, feature = "host_env"))] + use crate::AsObject; + use crate::Interpreter; + use crate::vm::stop_the_world::StopTheWorldGroup; + use alloc::sync::Arc; + use core::{ + panic::AssertUnwindSafe, + sync::atomic::{AtomicBool, AtomicUsize, Ordering}, + }; + + fn assert_empty() { + let slot = crate::vm::thread::current_thread_slot().unwrap(); + assert_eq!(slot.critical_section.load(Ordering::Relaxed), 0); + } + + #[test] + fn nested_sections_restore_ownership() { + Interpreter::builder(Default::default()) + .build() + .enter(|vm| { + let first = PyMutex::new(); + let second = PyMutex::new(); + with(vm, &first, || { + with(vm, &first, || assert!(first.is_locked())); + with(vm, &second, || { + assert!(first.is_locked()); + // Not the top-most section: locking it again detaches, + // which suspends every section of this thread. + with(vm, &first, || { + assert!(first.is_locked()); + assert!(!second.is_locked()); + }); + assert!(second.is_locked()); + assert!(!first.is_locked()); + }); + assert!(first.is_locked()); + }); + assert!(!first.is_locked()); + assert!(!second.is_locked()); + assert_empty(); + }); + } + + #[test] + fn two_mutex_sections() { + Interpreter::builder(Default::default()) + .build() + .enter(|vm| { + let first = PyMutex::new(); + let second = PyMutex::new(); + with2(vm, &second, &first, || { + assert!(first.is_locked() && second.is_locked()); + with2(vm, &first, &second, || {}); + with(vm, &second, || {}); + vm.allow_threads(|| assert!(!first.is_locked() && !second.is_locked())); + assert!(first.is_locked() && second.is_locked()); + }); + with2(vm, &first, &first, || assert!(first.is_locked())); + assert!(!first.is_locked() && !second.is_locked()); + assert_empty(); + }); + } + + #[test] + fn detachment_and_unwind_restore_sections() { + Interpreter::builder(Default::default()) + .build() + .enter(|vm| { + let m = PyMutex::new(); + let result = std::panic::catch_unwind(AssertUnwindSafe(|| { + with(vm, &m, || { + vm.allow_threads(|| assert!(!m.is_locked())); + assert!(m.is_locked()); + panic!("unwind the critical section"); + }); + })); + assert!(result.is_err()); + assert!(!m.is_locked()); + assert_empty(); + with(vm, &m, || assert!(m.is_locked())); + }); + } + + #[test] + fn interpreter_switch_restores_only_its_own_sections() { + let first = Interpreter::builder(Default::default()).build(); + let second = Interpreter::builder(Default::default()).build(); + first.enter(|first_vm| { + let first_lock = PyMutex::new(); + with(first_vm, &first_lock, || { + second.enter(|second_vm| { + assert!(!first_lock.is_locked()); + let second_lock = PyMutex::new(); + with(second_vm, &second_lock, || { + first.enter(|restored_vm| { + assert!(first_lock.is_locked()); + assert!(!second_lock.is_locked()); + with(restored_vm, &first_lock, || {}); + }); + assert!(second_lock.is_locked()); + }); + }); + assert!(first_lock.is_locked()); + }); + }); + } + + #[test] + fn opposite_order_contention_completes() { + Interpreter::builder(Default::default()) + .build() + .enter(|vm| { + let first = Arc::new(PyMutex::new()); + let second = Arc::new(PyMutex::new()); + let completed = Arc::new(AtomicUsize::new(0)); + let workers: Vec<_> = [false, true] + .into_iter() + .map(|reverse| { + let (outer, inner) = if reverse { + (second.clone(), first.clone()) + } else { + (first.clone(), second.clone()) + }; + let completed = completed.clone(); + vm.start_thread(move |vm| { + for _ in 0..1000 { + with(vm, &outer, || { + with(vm, &inner, || { + completed.fetch_add(1, Ordering::Relaxed); + }) + }); + } + }) + }) + .collect(); + vm.allow_threads(|| { + for worker in workers { + worker.join().unwrap(); + } + }); + assert_eq!(completed.load(Ordering::Relaxed), 2000); + assert!(!first.is_locked()); + assert!(!second.is_locked()); + }); + } + + #[test] + fn stop_the_world_releases_critical_sections() { + Interpreter::builder(Default::default()) + .build() + .enter(|vm| { + let m = Arc::new(PyMutex::new()); + let finished = Arc::new(AtomicBool::new(false)); + let (ready, started) = std::sync::mpsc::channel(); + let worker = { + let m = m.clone(); + let finished = finished.clone(); + vm.start_thread(move |vm| { + with(vm, &m, || { + ready.send(()).unwrap(); + while !finished.load(Ordering::Acquire) { + crate::vm::thread::suspend_if_needed(&vm.state); + std::thread::yield_now(); + } + assert!(m.is_locked()); + }); + }) + }; + vm.allow_threads(|| started.recv().unwrap()); + let mut stopped = StopTheWorldGroup::new(); + assert!(!m.is_locked()); + finished.store(true, Ordering::Release); + stopped.restart(); + vm.allow_threads(|| worker.join().unwrap()); + assert!(!m.is_locked()); + }); + } + + #[test] + fn stopped_world_skips_locking() { + Interpreter::builder(Default::default()) + .build() + .enter(|vm| { + let m = PyMutex::new(); + m.lock(); + let mut stopped = StopTheWorldGroup::new(); + with(vm, &m, assert_empty); + stopped.restart(); + m.unlock(); + }); + } + + // Run this exact test alone in an owned process under an external timeout. + // It forks while a critical section is active in another interpreter, and + // keeps lexical sections in both interpreters alive across the child's + // slot repair. + #[cfg(all(unix, feature = "host_env"))] + #[test] + #[ignore = "requires an isolated process and external process-tree timeout"] + fn fork_preserves_nested_interpreter_sections() { + let first = Interpreter::builder(Default::default()).build(); + let second = Interpreter::builder(Default::default()).build(); + first.enter(|first_vm| { + let outer = PyMutex::new(); + let worker_lock = Arc::new(PyMutex::new()); + let finished = Arc::new(AtomicBool::new(false)); + let (ready, started) = std::sync::mpsc::channel(); + let worker = { + let m = worker_lock.clone(); + let finished = finished.clone(); + first_vm.start_thread(move |vm| { + with(vm, &m, || { + ready.send(()).unwrap(); + while !finished.load(Ordering::Acquire) { + crate::vm::thread::suspend_if_needed(&vm.state); + std::thread::yield_now(); + } + }); + }) + }; + first_vm.allow_threads(|| started.recv().unwrap()); + with(first_vm, &outer, || { + second.enter(|vm| { + let inner = PyMutex::new(); + let posix = vm.import("posix", 0).unwrap(); + let lookups = Arc::new(AtomicUsize::new(0)); + let hook = { + let lookups = lookups.clone(); + vm.new_function( + "__getattr__", + move |_name: crate::PyObjectRef, + vm: &VirtualMachine| + -> crate::PyResult { + // Stream repair can reenter Python. Its native + // callback must be able to acquire a fresh STW + // transaction after core child repair. + let mut stopped = StopTheWorldGroup::new(); + stopped.restart(); + lookups.fetch_add(1, Ordering::Relaxed); + Err(vm.new_attribute_error("missing test stream")) + }, + ) + }; + vm.sys_module.set_attr("__getattr__", hook, vm).unwrap(); + let _ = vm.sys_module.dict().del_item("stdin", vm); + with(vm, &inner, || { + lookups.store(0, Ordering::Relaxed); + let pid = vm + .call_method(posix.as_object(), "fork", ()) + .unwrap() + .try_into_value::(vm) + .unwrap(); + if pid == 0 { + assert!(lookups.load(Ordering::Relaxed) > 0); + assert!(inner.is_locked()); + assert!(!worker_lock.is_locked()); + first.enter(|vm| { + assert!(outer.is_locked()); + assert!(!inner.is_locked()); + with(vm, &worker_lock, || {}); + }); + assert!(inner.is_locked()); + // The inherited shared transaction was released. + let mut stopped = StopTheWorldGroup::new(); + stopped.restart(); + // SAFETY: this owned child must not return into the + // parent test runner or join the vanished worker. + unsafe { libc::_exit(0) }; + } + finished.store(true, Ordering::Release); + let mut status = 0; + vm.allow_threads(|| { + rustpython_host_env::posix::waitpid(pid, &mut status, 0).unwrap(); + }); + assert!(libc::WIFEXITED(status)); + assert_eq!(libc::WEXITSTATUS(status), 0); + }); + }); + assert!(outer.is_locked()); + }); + first_vm.allow_threads(|| worker.join().unwrap()); + }); + } + + // The worker deliberately remains parked until the owned test process exits. + #[test] + #[ignore = "requires an isolated process and external process-tree timeout"] + fn shutting_down_releases_critical_sections() { + Interpreter::builder(Default::default()) + .build() + .enter(|vm| { + let m = Arc::new(PyMutex::new()); + let (ready, started) = std::sync::mpsc::channel(); + let worker = { + let m = m.clone(); + vm.start_thread(move |vm| { + with(vm, &m, || { + ready + .send(crate::vm::thread::current_thread_slot().unwrap()) + .unwrap(); + crate::vm::thread::hang_current_thread(&vm.state); + }); + }) + }; + let slot = vm.allow_threads(|| started.recv().unwrap()); + vm.allow_threads(|| { + while slot.state.load(Ordering::Acquire) + != crate::vm::thread::ThreadState::ShuttingDown as i32 + { + std::thread::yield_now(); + } + }); + assert!(!m.is_locked()); + with(vm, &m, || assert!(m.is_locked())); + drop(worker); + }); + } +} diff --git a/crates/vm/src/vm/thread/lock.rs b/crates/vm/src/vm/thread/lock.rs new file mode 100644 index 000000000..b517be9a7 --- /dev/null +++ b/crates/vm/src/vm/thread/lock.rs @@ -0,0 +1,293 @@ +//! One-byte mutex whose blocking acquire detaches the waiting thread. +//! +//! Only the two least significant bits of the state are used: +//! 0b00 unlocked, 0b01 locked, 0b10 unlocked with parked threads, +//! 0b11 locked with parked threads. + +use core::sync::atomic::AtomicU8; +#[cfg(feature = "threading")] +use core::sync::atomic::Ordering; + +#[cfg(feature = "threading")] +const UNLOCKED: u8 = 0; +#[cfg(feature = "threading")] +const LOCKED: u8 = 1; +#[cfg(feature = "threading")] +const HAS_PARKED: u8 = 2; + +/// `PyMutex`. Zero is the unlocked state. +#[derive(Debug, Default)] +#[repr(transparent)] +pub(crate) struct PyMutex { + bits: AtomicU8, +} + +impl PyMutex { + #[cfg_attr( + not(all(test, feature = "threading")), + expect(dead_code, reason = "Removed when the profiler consumer is wired") + )] + pub(crate) const fn new() -> Self { + Self { + bits: AtomicU8::new(0), + } + } +} + +#[cfg(feature = "threading")] +mod threaded { + use super::{HAS_PARKED, LOCKED, Ordering, PyMutex, UNLOCKED}; + use core::cell::Cell; + use core::time::Duration; + use parking_lot_core::{FilterOp, ParkResult, ParkToken, UnparkToken}; + use std::time::Instant; + + /// A waiter blocked longer than this is handed the lock directly by the + /// unlocking thread, which avoids starvation. + const TIME_TO_BE_FAIR: Duration = Duration::from_millis(1); + + /// Spin for a bit before parking the thread. + const MAX_SPIN_COUNT: usize = 40; + + const TOKEN_NORMAL: UnparkToken = UnparkToken(0); + const TOKEN_HANDED_OFF: UnparkToken = UnparkToken(1); + + /// `struct mutex_entry`; its address is the waiter's park token. + struct MutexEntry { + /// After this point the unlocking thread hands off ownership directly. + time_to_be_fair: Instant, + } + + impl PyMutex { + /// `PyMutex_LockFast` + #[inline] + pub(crate) fn lock_fast(&self) -> bool { + self.bits + .compare_exchange(UNLOCKED, LOCKED, Ordering::Acquire, Ordering::Relaxed) + .is_ok() + } + + /// `PyMutex_Lock`. A thread attached to an interpreter detaches while + /// it is parked. + #[inline] + pub(crate) fn lock(&self) { + if !self.lock_fast() { + self.lock_slow(); + } + } + + /// `_PyMutex_LockTimed` with no timeout and `_PY_LOCK_DETACH`. + #[cold] + fn lock_slow(&self) { + let mut v = self.bits.load(Ordering::Relaxed); + if v & LOCKED == 0 + && self + .bits + .compare_exchange(v, v | LOCKED, Ordering::Acquire, Ordering::Relaxed) + .is_ok() + { + return; + } + + let entry = MutexEntry { + time_to_be_fair: Instant::now() + TIME_TO_BE_FAIR, + }; + + let mut spin_count = 0; + loop { + if v & LOCKED == 0 { + // The lock is unlocked. Try to grab it. + match self.bits.compare_exchange( + v, + v | LOCKED, + Ordering::Acquire, + Ordering::Relaxed, + ) { + Ok(_) => return, + Err(actual) => { + v = actual; + continue; + } + } + } + + if v & HAS_PARKED == 0 && spin_count < MAX_SPIN_COUNT { + // Spin for a bit. + std::thread::yield_now(); + spin_count += 1; + continue; + } + + let mut newv = v; + if v & HAS_PARKED == 0 { + // We are the first waiter. Set the HAS_PARKED flag. + newv = v | HAS_PARKED; + if let Err(actual) = + self.bits + .compare_exchange(v, newv, Ordering::Relaxed, Ordering::Relaxed) + { + v = actual; + continue; + } + } + + if self.park(newv, &entry) == ParkResult::Unparked(TOKEN_HANDED_OFF) { + // We own the lock now. + debug_assert!(self.bits.load(Ordering::Relaxed) & LOCKED != 0); + return; + } + + v = self.bits.load(Ordering::Relaxed); + } + } + + /// `_PyParkingLot_Park` on this mutex, detached when attached. + fn park(&self, expected: u8, entry: &MutexEntry) -> ParkResult { + let result = Cell::new(ParkResult::Invalid); + crate::vm::thread::wait_detached_from_interpreter(&|| { + // SAFETY: the key is this mutex's address, used only by this + // mutex. The validate and timed-out callbacks only load an + // atomic or do nothing, so they neither panic nor reenter the + // parking lot. The entry outlives the park, and the unparking + // thread reads it under the bucket lock while this thread is + // still queued. + result.set(unsafe { + parking_lot_core::park( + self.key(), + || self.bits.load(Ordering::Relaxed) == expected, + || {}, + |_, _| {}, + ParkToken(core::ptr::from_ref(entry) as usize), + None, + ) + }); + }); + result.get() + } + + /// `PyMutex_Unlock` + #[inline] + pub(crate) fn unlock(&self) { + if self + .bits + .compare_exchange(LOCKED, UNLOCKED, Ordering::Release, Ordering::Relaxed) + .is_err() + { + self.unlock_slow(); + } + } + + /// `_PyMutex_UnlockSlow` + #[cold] + fn unlock_slow(&self) { + assert!( + self.try_unlock().is_ok(), + "unlocking mutex that is not locked" + ); + } + + /// `_PyMutex_TryUnlock` + fn try_unlock(&self) -> Result<(), ()> { + let mut v = self.bits.load(Ordering::Acquire); + loop { + if v & LOCKED == 0 { + // error: the mutex is not locked + return Err(()); + } else if v & HAS_PARKED != 0 { + // wake up a single thread + self.unpark(); + return Ok(()); + } + match self + .bits + .compare_exchange(v, UNLOCKED, Ordering::Release, Ordering::Relaxed) + { + // fast-path: no waiters + Ok(_) => return Ok(()), + Err(actual) => v = actual, + } + } + } + + /// `_PyParkingLot_Unpark` with `mutex_unpark`. + fn unpark(&self) { + let waiter = Cell::new(None); + // SAFETY: the key is this mutex's address. The filter and callback + // run under the bucket lock, do not panic and do not reenter the + // parking lot. A token is the address of a `MutexEntry` its waiter + // keeps alive while it is queued. + unsafe { + parking_lot_core::unpark_filter( + self.key(), + |token| { + if waiter.get().is_some() { + return FilterOp::Stop; + } + waiter.set(Some(token.0 as *const MutexEntry)); + FilterOp::Unpark + }, + |result| { + let mut v = UNLOCKED; + let mut token = TOKEN_NORMAL; + if let Some(entry) = waiter.get() { + let should_be_fair = Instant::now() > (*entry).time_to_be_fair; + if should_be_fair { + v |= LOCKED; + token = TOKEN_HANDED_OFF; + } + if result.have_more_threads { + v |= HAS_PARKED; + } + } + self.bits.store(v, Ordering::Release); + token + }, + ); + } + } + + /// `PyMutex_IsLocked` + #[cfg_attr( + not(test), + expect(dead_code, reason = "Removed when the profiler consumer is wired") + )] + pub(crate) fn is_locked(&self) -> bool { + self.bits.load(Ordering::Acquire) & LOCKED != 0 + } + + fn key(&self) -> usize { + core::ptr::from_ref(&self.bits) as usize + } + } +} + +#[cfg(test)] +#[cfg(feature = "threading")] +mod tests { + use super::{HAS_PARKED, PyMutex}; + use alloc::sync::Arc; + use core::{sync::atomic::Ordering, time::Duration}; + + #[test] + fn unlock_hands_off_to_starved_waiter() { + let m = Arc::new(PyMutex::new()); + m.lock(); + let waiter = { + let m = m.clone(); + std::thread::spawn(move || { + m.lock(); + m.unlock(); + }) + }; + while m.bits.load(Ordering::Relaxed) & HAS_PARKED == 0 { + std::thread::yield_now(); + } + // Long past the waiter's fairness deadline. + std::thread::sleep(Duration::from_millis(20)); + m.unlock(); + // The waiter owns the lock before it has even run again. + assert!(!m.lock_fast()); + waiter.join().unwrap(); + assert!(!m.is_locked()); + } +}