jit: root the forced-virtual caches, store them inside handle_async_forcing, and zero PyFrame.vable_token on JIT-allocated frames

majit/majit-gc/src/collector.rs

@@ -2857,7 +2857,7 @@ impl MiniMarkGC {
         // An owner outside old-gen is either immortal (`malloc_typed`, no
         // header to read) or, under a non-moving major, still in the live
         // nursery; neither can be proven dead here, so both are kept.
-        crate::shadow_stack::prune_ephemeron_tables(&mut |owner| {
+        let mut classify_owner = |owner: usize| -> Option<usize> {
             if owner == 0 || !self.oldgen.contains(owner) {
                 return Some(owner);
             }
@@ -2867,7 +2867,19 @@ impl MiniMarkGC {
             } else {
                 None
             }
-        });
+        };
+        crate::shadow_stack::prune_ephemeron_tables(&mut classify_owner);
+        // The same question for tables a single mutator owns in its own TLS.
+        // Those cannot go through the global registration: it names no thread,
+        // so a major driven here would leave every other mutator's dead-owner
+        // entries pinned. Reach exactly as far as this collection's own root
+        // walk did (`enumerate_root_walker_values`) — foreign TLS only while
+        // this thread owns STW.
+        if crate::gc_sync::mutators_quiesced() {
+            crate::shadow_stack::prune_all_mutator_areas(&mut classify_owner);
+        } else {
+            crate::shadow_stack::prune_my_mutator_areas(&mut classify_owner);
+        }
         // incminimark.py:2510-2511 — run destructors of dying old objects
         // before the sweep frees them (VISITED still distinguishes
         // survivors from the dying at this point).

majit/majit-gc/src/shadow_stack.rs

@@ -263,6 +263,7 @@ struct MutatorEntry {
     bh_regs_stack: *const RefCell<Vec<BhRegsEntry>>,
     resume_ref_roots_stack: *const RefCell<Vec<(*mut i64, usize)>>,
     extra_areas: Vec<MutatorExtraArea>,
+    pruners: Vec<MutatorPruner>,
 }
 
 /// Walker for one opaque root area owned by a registered mutator.
@@ -277,6 +278,18 @@ struct MutatorExtraArea {
     data: *const (),
 }
 
+/// Pruner for one owner-keyed side table owned by a registered mutator.
+///
+/// Same contract as [`MutatorExtraWalkFn`]: it runs on the collecting thread
+/// and must derive everything from `data`, never from caller TLS.
+pub type MutatorPrunerFn = unsafe fn(*const (), &mut dyn FnMut(usize) -> Option<usize>);
+
+#[derive(Clone, Copy)]
+struct MutatorPruner {
+    prune: MutatorPrunerFn,
+    data: *const (),
+}
+
 // The raw pointers refer to TLS owned by `thread_id`. The registry only moves
 // pointer values between threads; dereferencing them requires the STW
 // quiescence established by gc_sync.
@@ -307,6 +320,7 @@ pub fn register_mutator() {
         bh_regs_stack,
         resume_ref_roots_stack,
         extra_areas: Vec::new(),
+        pruners: Vec::new(),
     });
 }
 
@@ -328,6 +342,72 @@ pub unsafe fn register_mutator_extra_area(walk: MutatorExtraWalkFn, data: *const
     entry.extra_areas.push(MutatorExtraArea { walk, data });
 }
 
+/// Append an owner-keyed-table pruner to the current registered mutator.
+///
+/// The ephemeron half of [`register_mutator_extra_area`]: a table whose keys are
+/// owner addresses and whose values a walker roots needs its dead-owner entries
+/// dropped, and needs it with the same per-mutator reach the root walk has.
+/// [`register_ephemeron_pruner`] cannot serve TLS-owned state — it hands the
+/// classifier no way to name a thread, so a major driven by one thread would
+/// leave every other thread's dead-owner entries pinned.
+///
+/// # Safety
+///
+/// Same as [`register_mutator_extra_area`]: `data` must stay valid until
+/// [`unregister_mutator`] runs on this thread, and `prune` must derive every
+/// address it dereferences from `data`, never from caller TLS.
+pub unsafe fn register_mutator_pruner(prune: MutatorPrunerFn, data: *const ()) {
+    let thread_id = std::thread::current().id();
+    let mut registry = MUTATOR_REGISTRY.lock().unwrap();
+    let entry = registry
+        .iter_mut()
+        .find(|entry| entry.thread_id == thread_id)
+        .expect("register_mutator_pruner called before register_mutator");
+    entry.pruners.push(MutatorPruner { prune, data });
+}
+
+/// Prune every registered mutator's owner-keyed tables during STW.
+///
+/// Called from the same pre-sweep point as [`prune_ephemeron_tables`] and with
+/// the same classifier; see that function for why only a major prunes. Callers
+/// pick between this and [`prune_my_mutator_areas`] on
+/// `gc_sync::mutators_quiesced()`, exactly as the root walk picks between
+/// [`walk_all_extra_areas`] and [`walk_my_extra_areas`] — so a collection's
+/// prune reach always equals its own root-walk reach.
+pub fn prune_all_mutator_areas(classify: &mut dyn FnMut(usize) -> Option<usize>) {
+    debug_assert!(
+        crate::gc_sync::mutators_quiesced(),
+        "prune_all_mutator_areas reaches foreign mutator TLS; caller must own collector-side STW",
+    );
+    let registry = MUTATOR_REGISTRY.lock().unwrap();
+    for mutator in registry.iter() {
+        for pruner in mutator.pruners.iter() {
+            // SAFETY: gc_sync has quiesced every registered owner, and each
+            // pruner's data remains valid until its MutatorEntry is removed.
+            unsafe { (pruner.prune)(pruner.data, classify) };
+        }
+    }
+}
+
+/// Prune the current mutator's owner-keyed tables.
+///
+/// The single-thread collection path, mirroring [`walk_my_extra_areas`]:
+/// callers without a registered mutator have no per-thread tables and are a
+/// no-op. A collection that only walked its own roots must only prune its own
+/// tables — another mutator's owner was never marked here, so its entries
+/// cannot be classified.
+pub fn prune_my_mutator_areas(classify: &mut dyn FnMut(usize) -> Option<usize>) {
+    let thread_id = std::thread::current().id();
+    let registry = MUTATOR_REGISTRY.lock().unwrap();
+    let Some(mutator) = registry.iter().find(|entry| entry.thread_id == thread_id) else {
+        return;
+    };
+    for pruner in mutator.pruners.iter() {
+        // SAFETY: this is the owning thread's synchronous collection path.
+        unsafe { (pruner.prune)(pruner.data, classify) };
+    }
+}
+
 /// Walk every registered mutator's opaque extra root areas during STW.
 pub fn walk_all_extra_areas(mut visitor: impl FnMut(&mut GcRef)) {
     debug_assert!(

majit/majit-metainterp/src/jitdriver.rs

@@ -4590,14 +4590,14 @@ impl<S: JitState> JitDriver<S> {
     /// virtuals through the same resume allocator used by ordinary guard
     /// failure.  In particular, a `jit.virtual_ref` frame must not be decoded
     /// through `NullAllocator`, or its `forced` writeback remains null.
-    pub fn force_virtualizable_token(&mut self, token: u64) -> Option<(Vec<i64>, Vec<i64>)> {
+    pub fn force_virtualizable_token(&mut self, token: u64) {
         let fallback_alloc = crate::resume::NullAllocator;
         let allocator: &dyn crate::resume::BlackholeAllocator = self
             .blackhole_allocator
             .as_deref()
             .unwrap_or(&fallback_alloc);
         self.meta
-            .force_virtualizable_token_with_allocator(token, allocator)
+            .force_virtualizable_token_with_allocator(token, allocator);
     }
 
     fn prepare_exit_resume_heap_with_blackhole_allocator(
@@ -4643,6 +4643,18 @@ impl<S: JitState> JitDriver<S> {
         self.meta.walk_compile_snapshot_refs(visitor);
     }
 
+    /// GC walker for the forced-virtual caches awaiting a `GUARD_NOT_FORCED`.
+    /// See `MetaInterp::walk_forced_virtuals_refs`.
+    pub fn walk_forced_virtuals_refs(&mut self, visitor: impl FnMut(&mut majit_ir::GcRef)) {
+        self.meta.walk_forced_virtuals_refs(visitor);
+    }
+
+    /// Drop forced-virtual caches whose owner frame died.
+    /// See `MetaInterp::prune_forced_virtuals`.
+    pub fn prune_forced_virtuals(&mut self, classify: &mut dyn FnMut(usize) -> Option<usize>) {
+        self.meta.prune_forced_virtuals(classify);
+    }
+
     pub fn run_compiled_detailed_keyed(
         &mut self,
         green_key: u64,