Serialize stop-the-world transactions across interpreters (#9094)

* Serialize stop-the-world transactions across interpreters

Coordinate individual stop requests with process-wide collection, retain snapshot owners through rollback, and release transferred ownership after restart.

Assisted-by: OpenAI Codex: model identifier unavailable

* Reject reentrant stop-the-world transactions

Use non-reused logical requester identities to fail before waiting on the same thread, while preserving cross-thread restart and fork ownership.

Assisted-by: OpenAI Codex: model identifier unavailable

---------

Co-authored-by: dot <dot@youknowone.org>

crates/vm/src/gc_state.rs

@@ -207,22 +207,24 @@ type GcMap<K, V> = std::collections::HashMap<K, V, GcBuildHasher>;
 /// 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
-/// keeps exclusion acquisition ordered; the `collecting` mutex additionally
-/// serializes collections process-wide, so no second collector can take these
-/// exclusions in another 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 {
-    /// Stopped interpreter states, in stop order. Held as strong references so
-    /// an interpreter cannot be dropped between stop and restart, and kept past
-    /// the restart so that releasing the last one — which frees that
-    /// interpreter's objects, and so removes them from these lists — happens
-    /// after the collection has let go of the generation locks.
-    stopped: Vec<crate::common::rc::PyRc<crate::vm::PyGlobalState>>,
+    /// 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 `stopped`,
+    /// 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,
 }
 
@@ -237,8 +239,10 @@ impl CollectStopTheWorld {
         // keep the historical no-barrier fallback.
         if !crate::vm::thread::current_vm_is_set() {
             return Self {
-                stopped: Vec::new(),
+                states: Vec::new(),
+                stopped_count: 0,
                 admission: None,
+                transaction: None,
                 restarted: true,
             };
         }
@@ -247,14 +251,20 @@ impl CollectStopTheWorld {
         // `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 {
-            stopped: Vec::new(),
+            states: Vec::new(),
+            stopped_count: 0,
             admission: Some(crate::vm::runtime::lock_admission_for_stop()),
+            transaction: Some(transaction),
             restarted: false,
         };
-        for state in crate::vm::runtime::live_interpreter_states() {
-            state.stop_the_world.stop_the_world(&state);
-            guard.stopped.push(state);
+        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
     }
@@ -267,17 +277,18 @@ impl CollectStopTheWorld {
         self.restarted = true;
         // Reverse of the stop order. The references stay until this guard is
         // dropped; see the field comment.
-        for state in self.stopped.iter().rev() {
+        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.is_empty()
+        self.stopped_count != 0
     }
 }
 

crates/vm/src/vm/mod.rs

@@ -189,6 +189,68 @@ impl Default for ExceptionStack {
     }
 }
 
+/// Serializes complete stop/restart transactions across interpreters.
+/// A collector holds one permit while stopping every interpreter; other callers
+/// transfer their permit to the stopped state until the matching restart.
+#[cfg(feature = "threading")]
+pub(crate) struct StopTheWorldTransaction {
+    _private: (),
+}
+
+#[cfg(feature = "threading")]
+static STW_TRANSACTION: AtomicU64 = AtomicU64::new(0);
+
+#[cfg(feature = "threading")]
+static NEXT_STW_REQUESTER: AtomicU64 = AtomicU64::new(1);
+
+#[cfg(feature = "threading")]
+thread_local! {
+    // OS thread IDs can be reused after a requester exits, even while another
+    // thread still owes its matching restart. Logical IDs are never reused.
+    static STW_REQUESTER: u64 = NEXT_STW_REQUESTER
+        .try_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1))
+        .expect("stop-the-world requester IDs exhausted");
+}
+
+#[cfg(feature = "threading")]
+impl StopTheWorldTransaction {
+    pub(crate) fn acquire() -> Self {
+        let requester = STW_REQUESTER.with(|id| *id);
+        loop {
+            match STW_TRANSACTION.compare_exchange(
+                0,
+                requester,
+                Ordering::AcqRel,
+                Ordering::Relaxed,
+            ) {
+                Ok(_) => return Self { _private: () },
+                Err(owner) => assert_ne!(
+                    owner, requester,
+                    "stop-the-world transactions cannot be nested on the same thread"
+                ),
+            }
+            let wait = || {
+                while STW_TRANSACTION.load(Ordering::Acquire) != 0 {
+                    std::thread::yield_now();
+                }
+            };
+            // Wait detached from the CURRENT interpreter, not the interpreter
+            // we intend to stop. Reattach without owning the permit, then retry:
+            // a suspended or shutting-down waiter must not strand ownership.
+            if thread::try_with_current_vm(|vm| vm.allow_threads(wait)).is_none() {
+                wait();
+            }
+        }
+    }
+}
+
+#[cfg(feature = "threading")]
+impl Drop for StopTheWorldTransaction {
+    fn drop(&mut self) {
+        STW_TRANSACTION.store(0, Ordering::Release);
+    }
+}
+
 /// Stop-the-world state for fork safety. Before `fork()`, the requester
 /// stops all other Python threads so they are not holding internal locks.
 #[cfg(feature = "threading")]
@@ -199,10 +261,9 @@ pub struct StopTheWorldState {
     world_stopped: AtomicBool,
     /// Ident of the thread that requested the stop (like `stw->requester`)
     requester: AtomicU64,
-    /// Single exclusion held for the whole stop→start span. Fork and GC are
-    /// both stop-the-world requesters driving this shared state; only one may
-    /// hold it at a time. Acquired before any stop bookkeeping (see
-    /// `acquire_exclusion`) and released by `start_the_world`/`reset_after_fork`.
+    /// Ownership transferred by a successful public stop. GC retains its
+    /// transaction externally and leaves this false. Only a matching restart
+    /// or the forking owner's child reset may release this ownership.
     exclusion: AtomicBool,
     /// Signaled by suspending threads when their state transitions to SUSPENDED
     notify_mutex: std::sync::Mutex<()>,
@@ -401,50 +462,11 @@ impl StopTheWorldState {
         forced_parks != 0 && self.thread_countdown.load(Ordering::Acquire) == 0
     }
 
-    /// Acquire the single stop-the-world exclusion in a park-friendly way.
-    ///
-    /// Fork and GC both request stop-the-world through the same shared state;
-    /// without this exclusion their `requester`/`requested`/countdown words
-    /// could be clobbered by an interleaving requester, so the completion
-    /// check could never converge and a requester would wait on itself forever.
-    ///
-    /// The acquire must be park-friendly. While another requester's stop is in
-    /// progress it sets this thread's stop bit and waits for it to suspend;
-    /// blocking on a plain lock here would keep this thread from ever reaching
-    /// that safepoint, so the active requester would wait for this thread while
-    /// this thread waits for the lock — a deadlock swap. Instead we poll and
-    /// honor the suspend request between tries. Suspending here is safe as long
-    /// as any lock a spinning requester still holds is never acquired
-    /// attached-blocking by another thread. The fork requester holds IMP_LOCK,
-    /// but its acquisition detaches (`allow_threads`), so no attached thread
-    /// blocks on it; the GC requester holds only the `collecting` mutex, which
-    /// is only ever `try_lock`'d. The active requester therefore force-parks
-    /// this thread, finishes its whole stop→start span, releases the exclusion,
-    /// and only then does this thread resume and acquire it.
-    fn acquire_exclusion(&self, state: &PyGlobalState) {
-        if self
-            .exclusion
-            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
-            .is_ok()
-        {
-            return;
-        }
-        loop {
-            crate::vm::thread::suspend_if_needed(state);
-            std::thread::yield_now();
-            if self
-                .exclusion
-                .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
-                .is_ok()
-            {
-                return;
-            }
-        }
-    }
-
-    /// Release the stop-the-world exclusion taken by `acquire_exclusion`.
+    /// Release only a transaction owned by this state, after every wake/clear.
     fn release_exclusion(&self) {
-        self.exclusion.store(false, Ordering::Release);
+        if self.exclusion.swap(false, Ordering::AcqRel) {
+            drop(StopTheWorldTransaction { _private: () });
+        }
     }
 
     /// Stop all non-requester threads (`stop_the_world`).
@@ -454,11 +476,27 @@ impl StopTheWorldState {
     /// 3. Waits (polling with 1 ms condvar timeout) for attached threads
     ///    to self-suspend in `check_signals`.
     ///
-    /// Takes the shared exclusion first so at most one requester (fork or GC)
-    /// drives the stop→start span at a time; it is released by
-    /// `start_the_world`/`reset_after_fork`.
+    /// Public calls own one process-wide transaction until the matching start.
+    /// The caller must finish stop before invoking start, including when those
+    /// calls occur on different threads; concurrent duplicate starts are invalid.
+    /// Transactions are non-reentrant: a nested stop on the requesting thread
+    /// panics before changing any stop state or releasing the outer transaction.
     pub fn stop_the_world(&self, state: &PyGlobalState) {
-        self.acquire_exclusion(state);
+        let transaction = StopTheWorldTransaction::acquire();
+        self.stop_with_transaction(state, &transaction);
+        self.exclusion.store(true, Ordering::Release);
+        core::mem::forget(transaction);
+    }
+
+    /// Stop one state within a transaction retained by the caller (notably GC).
+    pub(crate) fn stop_with_transaction(
+        &self,
+        state: &PyGlobalState,
+        _transaction: &StopTheWorldTransaction,
+    ) {
+        // A partial stop must clear requests and wake parked slots before its
+        // caller releases the transaction or rolls back previously stopped VMs.
+        let rollback = scopeguard::guard((), |_| self.resume_threads(state, false));
         let start = std::time::Instant::now();
         let requester_ident = crate::stdlib::_thread::get_ident();
         self.requester.store(requester_ident, Ordering::Relaxed);
@@ -478,6 +516,7 @@ impl StopTheWorldState {
             stw_trace(format_args!(
                 "stop end requester={requester_ident} wait_ns=0 polls=0"
             ));
+            scopeguard::ScopeGuard::into_inner(rollback);
             return;
         }
 
@@ -526,10 +565,17 @@ impl StopTheWorldState {
         stw_trace(format_args!(
             "stop end requester={requester_ident} wait_ns={wait_ns} polls={polls}"
         ));
+        scopeguard::ScopeGuard::into_inner(rollback);
     }
 
     /// Resume all suspended threads (`start_the_world`).
     pub fn start_the_world(&self, state: &PyGlobalState) {
+        self.resume_threads(state, true);
+        self.release_exclusion();
+    }
+
+    /// Rollback also permits threads that never reached SUSPENDED.
+    fn resume_threads(&self, state: &PyGlobalState, fully_stopped: bool) {
         use thread::ThreadState;
         let requester = self.requester.load(Ordering::Relaxed);
         stw_trace(format_args!("start begin requester={requester}"));
@@ -560,7 +606,7 @@ impl StopTheWorldState {
                 continue;
             }
             debug_assert!(
-                state == ThreadState::Suspended as i32,
+                !fully_stopped || state == ThreadState::Suspended as i32,
                 "non-requester thread not suspended at start-the-world: id={id} state={state}"
             );
             if state == ThreadState::Suspended as i32 {
@@ -579,9 +625,6 @@ impl StopTheWorldState {
         crate::signal::clear_stop_bit();
         #[cfg(debug_assertions)]
         self.debug_assert_all_non_requester_detached(state);
-        // Release the exclusion last, ending the stop→start span so the next
-        // requester (fork or GC) can proceed.
-        self.release_exclusion();
         stw_trace(format_args!("start end requester={requester}"));
     }
 
@@ -595,8 +638,9 @@ impl StopTheWorldState {
         // Only one thread survives fork; any stop-the-world bit inherited
         // from the parent is stale.
         crate::signal::clear_stop_bit();
-        // The surviving child thread inherited the exclusion taken by the
-        // pre-fork `stop_the_world`; release it (no start_the_world runs here).
+        // 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.
         self.release_exclusion();
         stw_trace(format_args!("reset-after-fork"));
     }