jit: re-mint the tracing sentinel when the GC is rebuilt

`set_tracing_rescall_dummy_gc_type_id` is called from `build_gc`
(`pyre-jit/src/eval.rs:1338,3639`), which `reset_gc_fresh_for_test` runs again
per GC-stress worker. The sentinel was a `OnceLock`, so a second heap kept the
address minted in the first — `is_managed_heap_object` no longer recognises it
once the heap it belongs to has been replaced, putting the traced
`virtual_token` / `vable_token` slots back on an address the collector does not
own. The assertion added with the previous commit turned that into a panic on
the second registration instead.

Hold the address in an `AtomicUsize` the setter clears, so the next request
mints in the heap that is now current, and publish it with a compare-exchange
so racing minters agree on one address.

The GC-stress harness produces this ordering the moment one of its programs
reaches a traced residual call; none does today, so the new test mints the
sentinel between two resets directly. It panics at the old assertion without
this change.

Reported by the Codex review bot on #1084.

Assisted-by: Claude

Author: youknowone

majit/majit-metainterp/src/virtualref.rs

@@ -13,7 +13,7 @@
 //!
 //! Mirrors `rpython/jit/metainterp/virtualref.py`.
 
-use std::sync::atomic::{AtomicU32, Ordering};
+use std::sync::atomic::{AtomicU32, AtomicUsize, Ordering};
 
 /// The value [`VREF_GC_TYPE_ID`] holds before `set_vref_gc_type_id` runs.  Zero
 /// is a legitimate id, so the sentinel has to be a value the registry never
@@ -195,10 +195,10 @@ pub const TOKEN_NONE: *mut u8 = std::ptr::null_mut();
 /// the translated identity word is pointer-sized too.
 pub const JITFRAME_DUMMY_VTABLE: usize = 0x4A46_444D; // "JFDM"
 
-/// Lazy initialisation of the `_dummy` address.  `OnceLock<usize>`
-/// (instead of `OnceLock<*mut u8>`) so the cell is `Sync` —
-/// raw-pointer types are not.
-static TRACING_RESCALL_DUMMY_PTR: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
+/// Lazy initialisation of the `_dummy` address, as a `usize` because raw
+/// pointers are not `Sync`.  Zero means "not minted yet"; a null sentinel would
+/// collide with `TOKEN_NONE`, so it is not a value this can ever hold.
+static TRACING_RESCALL_DUMMY_PTR: AtomicUsize = AtomicUsize::new(0);
 
 const TRACING_RESCALL_DUMMY_GC_TYPE_ID_UNSET: u32 = u32::MAX;
 static TRACING_RESCALL_DUMMY_GC_TYPE_ID: AtomicU32 =
@@ -207,15 +207,15 @@ static TRACING_RESCALL_DUMMY_GC_TYPE_ID: AtomicU32 =
 /// Publish the registered leaf type used by the prebuilt
 /// `virtualizable.py:326-330 JITFRAME_DUMMY` object.
 ///
-/// Must run before the sentinel is first requested — the address is minted
-/// once and never re-minted, so a late registration would leave the process
-/// with the unmanaged fallback for good.
+/// Registration comes from `build_gc`, so a second call means a second heap.
+/// A sentinel minted in the previous one is no longer part of the live heap —
+/// `is_managed_heap_object` would stop recognising it and the traced
+/// `virtual_token` / `vable_token` slots would be back to holding an address
+/// the collector does not own. Drop it so the next request mints in the heap
+/// that is now current.
 pub fn set_tracing_rescall_dummy_gc_type_id(type_id: u32) {
-    assert!(
-        TRACING_RESCALL_DUMMY_PTR.get().is_none(),
-        "JITFRAME_DUMMY type registered after the tracing sentinel was minted",
-    );
     TRACING_RESCALL_DUMMY_GC_TYPE_ID.store(type_id, Ordering::Relaxed);
+    TRACING_RESCALL_DUMMY_PTR.store(0, Ordering::Relaxed);
 }
 
 /// `virtualizable.py:327 _dummy = lltype.malloc(_DUMMY)` — allocate
@@ -254,11 +254,27 @@ fn allocate_tracing_rescall_dummy() -> *mut u8 {
 /// ```python
 /// TOKEN_TRACING_RESCALL = lltype.cast_opaque_ptr(llmemory.GCREF, _dummy)
 /// ```
-/// The returned address is a registered, rooted GC leaf and remains stable for
-/// the process lifetime.
+/// The returned address is a registered, rooted GC leaf, and stays the same for
+/// as long as the heap it was minted in is the current one.
 #[inline]
 pub fn token_tracing_rescall() -> *mut u8 {
-    *TRACING_RESCALL_DUMMY_PTR.get_or_init(|| allocate_tracing_rescall_dummy() as usize) as *mut u8
+    let minted = TRACING_RESCALL_DUMMY_PTR.load(Ordering::Relaxed);
+    if minted != 0 {
+        return minted as *mut u8;
+    }
+    let fresh = allocate_tracing_rescall_dummy() as usize;
+    // Racing minters both produce a valid sentinel, but the token protocol
+    // compares tokens by address, so exactly one may be published. The loser's
+    // object is immortal either way — upstream's `_dummy` is prebuilt.
+    match TRACING_RESCALL_DUMMY_PTR.compare_exchange(
+        0,
+        fresh,
+        Ordering::Relaxed,
+        Ordering::Relaxed,
+    ) {
+        Ok(_) => fresh as *mut u8,
+        Err(published) => published as *mut u8,
+    }
 }
 
 /// Virtual reference state for a single reference.

pyre/pyre-jit/tests/gc_stress.rs

@@ -1640,3 +1640,45 @@ while i < 40:
         "module dict move_to_end reentrant scan callback GC rooting program failed",
     );
 }
+
+/// `virtualizable.py:326-330` makes `TOKEN_TRACING_RESCALL` the address of a
+/// prebuilt GC object, so the sentinel has to belong to the heap that is
+/// current when a traced slot holds it. `reset_gc_fresh_for_test` builds a
+/// second heap and leaks the first, which leaves any sentinel minted in the
+/// first outside the live heap — `is_managed_heap_object` stops recognising it,
+/// and the traced `virtual_token` / `vable_token` slots would be back to
+/// holding an address the collector does not own.
+///
+/// Mint the sentinel between two resets and require a fresh address, which is
+/// what tells the two heaps apart. This is the ordering the GC-stress harness
+/// itself produces the moment any of its programs reaches a traced residual
+/// call; none does today, so nothing else in this binary covers it.
+#[test]
+fn tracing_sentinel_is_reminted_for_a_rebuilt_heap() {
+    let _serial = GC_STRESS_SERIAL.lock().unwrap_or_else(|e| e.into_inner());
+    let handle = std::thread::Builder::new()
+        .stack_size(256 * 1024 * 1024)
+        .spawn(|| {
+            init_jit_hooks();
+            reset_gc_fresh_for_test();
+            let first = majit_metainterp::virtualref::token_tracing_rescall();
+            assert!(!first.is_null(), "sentinel must never be TOKEN_NONE");
+
+            reset_gc_fresh_for_test();
+            let second = majit_metainterp::virtualref::token_tracing_rescall();
+            assert!(!second.is_null(), "sentinel must never be TOKEN_NONE");
+            assert_ne!(
+                first, second,
+                "the sentinel stayed in the heap that was replaced",
+            );
+
+            // Stable within one heap: the token protocol compares by address.
+            assert_eq!(
+                second,
+                majit_metainterp::virtualref::token_tracing_rescall(),
+                "the sentinel moved without the heap being rebuilt",
+            );
+        })
+        .expect("spawn worker thread");
+    handle.join().expect("worker thread panicked");
+}