mirrors/RustPython
0
0 Fork 0
mirror of https://github.com/RustPython/RustPython.git synced 2026-10-28 10:44:36 +08:00
Code Activity

Add interpreter-bound operation critical sections (#9215)

Browse Source 
* Factor shared interpreter stop-the-world guard

Move the existing collector barrier into the VM without changing its stop, restart or ownership behavior, preparing reuse by fork.

Assisted-by: dot:GPT-6.1

* Add interpreter-bound operation section primitive

Prepare attached-only ownership acquisition and detached observer parking with private per-interpreter scope stacks. Lifecycle wiring follows separately; no profiler consumer is enabled by this commit.

Assisted-by: dot:GPT-6.1

* Suspend operation sections at VM lifecycle boundaries

Integrate detach, stop-the-world, shutdown and interpreter-switch boundaries. Reuse the all-interpreter barrier for fork repair and add focused native ownership invariants.

Assisted-by: dot:GPT-6.1

* Replace operation sections with PyMutex critical sections

Add PyMutex with the lock.c protocol: LOCKED/HAS_PARKED bits, a bounded
spin, direct handoff to waiters parked longer than 1 ms, and detaching
while parked.

Rebuild critical sections on PyMutex with the stack kept as a tagged
pointer in ThreadSlot: single and two-mutex sections, recursive
top-most skipping, locking skipped while the world is stopped,
SuspendAll on detach/suspend/shutdown and Resume of the top-most
section on attach. Fork carries the stacks into the replacement slots.

Assisted-by: Claude Code:claude-opus-5-5

* Check attachment around critical-section suspension

Assert that a thread entering a critical section is attached, and
suspend sections in do_suspend only while the slot is still ATTACHED.

Assisted-by: Claude Code:claude-opus-5-5

---------

Co-authored-by: Jeong YunWon <jeong@youknowone.org>
This commit is contained in:
youknowdotandJeong YunWon authored and GitHub committed 2026-10-21 19:57:21 +09:00
1 parent b6b6446718
commit 97bd4af533
11 files changed
+1369 -144
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.cspell.dict/rust-more.txt
+2
View File
@@ -83,6 +83,8 @@ trai
ulonglong
unic
unistd
unparked
unparking
unraw
unsync
wasip1
Cargo.lock
Generated
+1
View File
@@ -4113,6 +4113,7 @@ dependencies = [
"num_enum",
"optional",
"parking_lot",
"parking_lot_core",
"paste",
"psm",
"rapidhash",
Cargo.toml
+1
View File
@@ -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"
crates/vm/Cargo.toml
+2 -1
View File
@@ -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 }
crates/vm/src/gc_state.rs
+4 -110
View File
@@ -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<GcPtrHasher>;
type GcSet<T> = std::collections::HashSet<T, GcBuildHasher>;
type GcMap<K, V> = std::collections::HashMap<K, V, GcBuildHasher>;
/// 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<crate::common::rc::PyRc<crate::vm::PyGlobalState>>,
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<parking_lot::MutexGuard<'static, ()>>,
/// Acquired before admission and retained until every state is restarted.
transaction: Option<crate::vm::StopTheWorldTransaction>,
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.
crates/vm/src/stdlib/posix.rs
+48 -23
View File
@@ -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<PyObjectRef> = 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<PyObjectRef> = 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))
}
}
crates/vm/src/vm/mod.rs
+5 -3
View File
@@ -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"));
}
crates/vm/src/vm/stop_the_world.rs
+175
View File
@@ -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<crate::common::rc::PyRc<crate::vm::PyGlobalState>>,
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<parking_lot::MutexGuard<'static, ()>>,
/// Acquired before admission and retained until every state is restarted.
transaction: Option<crate::vm::StopTheWorldTransaction>,
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<crate::vm::thread::CurrentFrameSlot>,
// 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();
}
}
crates/vm/src/vm/thread.rs
+65 -7
View File
@@ -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<crate::object::qsbr::QsbrSlot>,
/// `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<HashMap<i64, CurrentFrameSlot>> =
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<HashMap<i64, usize>> =
RefCell::new(HashMap::new());
/// Current thread's slot for the currently entered interpreter.
#[cfg(feature = "threading")]
static CURRENT_THREAD_SLOT: RefCell<Option<CurrentFrameSlot>> = 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<CurrentFrameSlot> {
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<CurrentFrameSlot>) {
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
});
});
}
crates/vm/src/vm/thread/critical_section.rs
+773
View File
@@ -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<R>(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<R>(
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<PyMutex>,
}
/// `PyCriticalSection2`
#[repr(C)]
struct PyCriticalSection2 {
base: PyCriticalSection,
mutex2: AtomicPtr<PyMutex>,
}
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<R>(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<R>(
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::<PyCriticalSection2>() };
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::<PyCriticalSection2>() };
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::<PyCriticalSection2>() };
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::<PyCriticalSection2>() };
(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::<i32>(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);
});
}
}
crates/vm/src/vm/thread/lock.rs
+293
View File
@@ -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());
}
}