jit: apply every elapsed counter-decay interval, and drain it before every table access

majit/majit-trace/src/counter.rs

@@ -154,12 +154,25 @@ impl JitCounter {
     /// TODO: no RPython counterpart. Read-only peek
     /// used by warmstate's cold fast path to avoid GreenKey allocation.
     pub fn would_tick_fire(&self, hash: u64, increment: f64) -> bool {
+        let elapsed = DECAY_GENERATION
+            .load(Ordering::Relaxed)
+            .wrapping_sub(self.last_decay_generation);
         let index = self._get_index(hash);
         let subhash = Self::_get_subhash(hash);
         let entry = &self.timetable[index];
         for i in 0..ASSOCIATIVITY {
             if entry.subhashes[i] == subhash {
-                return entry.times[i] as f64 + increment >= 1.0;
+                // This predicate is &self, so decay-adjust the read instead of
+                // mutating the table to apply the pending generations. Step
+                // through them one at a time rather than raising the multiplier
+                // to `elapsed`: `decay_all_counters` rounds back to f32 after
+                // every step, and this answer has to be the one a `tick` would
+                // give once it drains the same generations.
+                let mut time = entry.times[i];
+                for _ in 0..elapsed {
+                    time = (time as f64 * self.decay_by_mult) as f32;
+                }
+                return time as f64 + increment >= 1.0;
             }
         }
         increment >= 1.0
@@ -169,24 +182,19 @@ impl JitCounter {
     #[inline(always)]
     pub fn tick(&mut self, hash: u64, increment: f64) -> bool {
         // counter.py:104-121 applies the decay synchronously inside the minor
-        // collection. pyre defers it to the next tick because the metainterp
-        // owns the counter table, and reaching it from inside the collector
-        // would re-enter a borrow the GC does not hold. Counters are only read
-        // at tick time.
+        // collection, where the hook closes over the process's one JitCounter.
+        // pyre defers it to the next table access instead: the counter is
+        // reached through the `JIT_DRIVER` cell in eval.rs, whose accessor
+        // mints a `&'static mut JitDriverPair`, and a minor collection can be
+        // triggered by an allocation the metainterp makes while already holding
+        // one. Decaying from inside the collector would alias it.
         //
         // Each JitCounter keeps its own last-seen generation because JIT_DRIVER
-        // in eval.rs is thread-local, giving each mutator thread its own counter.
-        // This still saturates: if two or more 32-collection intervals elapse
-        // between ticks, counter.py has applied that many decays while this
-        // applies one. Carrying the elapsed count would match it, but would move
-        // inline_chain_depth_typeflip's recorded jit-stats (bridges_compiled
-        // 19 -> 18, guard_failures 3820 -> 3681); left for a change that can
-        // re-record them on every platform.
-        let generation = DECAY_GENERATION.load(Ordering::Relaxed);
-        if generation != self.last_decay_generation {
-            self.last_decay_generation = generation;
-            self.decay_all_counters();
-        }
+        // is thread-local, giving each mutator thread its own counter.
+        // Every elapsed interval is applied before every mutating table access;
+        // would_tick_fire decay-adjusts its read. A value written after a
+        // collection is therefore not retro-decayed.
+        self.apply_pending_decay();
 
         let index = self._get_index(hash);
         let subhash = Self::_get_subhash(hash);
@@ -213,6 +221,8 @@ impl JitCounter {
 
     /// counter.py:204-230 change_current_fraction(hash, new_fraction)
     pub fn change_current_fraction(&mut self, hash: u64, new_fraction: f64) {
+        self.apply_pending_decay();
+
         let index = self._get_index(hash);
         let subhash = Self::_get_subhash(hash);
         let entry = &mut self.timetable[index];
@@ -232,6 +242,8 @@ impl JitCounter {
 
     /// counter.py:232-237 reset(hash)
     pub fn reset(&mut self, hash: u64) {
+        self.apply_pending_decay();
+
         let index = self._get_index(hash);
         let subhash = Self::_get_subhash(hash);
         let entry = &mut self.timetable[index];
@@ -245,13 +257,17 @@ impl JitCounter {
     /// TODO: no RPython equivalent.
     /// Zero all timetable entries.
     pub fn reset_all(&mut self) {
+        self.apply_pending_decay();
+
         for entry in &mut self.timetable {
             *entry = Entry::default();
         }
     }
 
     /// counter.py:258-264 set_decay(decay)
     pub fn set_decay(&mut self, decay: i32) {
+        self.apply_pending_decay();
+
         let clamped = decay.clamp(0, 1000);
         self.decay_by_mult = 1.0_f64 - (clamped as f64 * 0.001);
     }
@@ -265,6 +281,18 @@ impl JitCounter {
             }
         }
     }
+
+    /// Apply every 32-collection interval that elapsed since this counter
+    /// last looked. counter.py:104-121 decays inside the collection, so a
+    /// pending decay must land before anything reads or writes the table —
+    /// otherwise it would decay values written after the collection.
+    fn apply_pending_decay(&mut self) {
+        let generation = DECAY_GENERATION.load(Ordering::Relaxed);
+        while self.last_decay_generation != generation {
+            self.last_decay_generation = self.last_decay_generation.wrapping_add(1);
+            self.decay_all_counters();
+        }
+    }
 }
 
 /// counter.py:309 DeterministicJitCounter — test-only, NOT_RPYTHON.
@@ -371,6 +399,25 @@ impl DeterministicJitCounter {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::sync::Mutex;
+
+    static DECAY_GENERATION_TEST_LOCK: Mutex<()> = Mutex::new(());
+
+    fn advance_decay_generation(intervals: usize) {
+        DECAY_GENERATION.fetch_add(intervals, Ordering::Relaxed);
+    }
+
+    fn counter_time(counter: &JitCounter, hash: u64) -> f32 {
+        let index = counter._get_index(hash);
+        let subhash = JitCounter::_get_subhash(hash);
+        let entry = &counter.timetable[index];
+        for i in 0..ASSOCIATIVITY {
+            if entry.subhashes[i] == subhash {
+                return entry.times[i];
+            }
+        }
+        0.0
+    }
 
     #[test]
     fn test_basic_counting() {
@@ -434,6 +481,46 @@ mod tests {
         assert!(time > 0.7 && time < 0.8, "time={}", time);
     }
 
+    #[test]
+    fn test_tick_applies_every_pending_decay_generation() {
+        let _generation_guard = DECAY_GENERATION_TEST_LOCK
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        let mut counter = JitCounter::new(DEFAULT_SIZE);
+        counter.set_decay(40);
+        let h = 3u64 << counter.shift;
+        counter.change_current_fraction(h, 0.5);
+
+        advance_decay_generation(2);
+        let increment = 0.001;
+        assert!(!counter.tick(h, increment));
+
+        let expected =
+            (((0.5f32 as f64 * 0.96) as f32 as f64 * 0.96) as f32 as f64 + increment) as f32;
+        let actual = counter_time(&counter, h);
+        assert!((actual - expected).abs() < 1.0e-6, "actual={actual}");
+    }
+
+    #[test]
+    fn test_change_current_fraction_is_not_retro_decayed() {
+        let _generation_guard = DECAY_GENERATION_TEST_LOCK
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        let mut counter = JitCounter::new(DEFAULT_SIZE);
+        counter.set_decay(40);
+        let h = 4u64 << counter.shift;
+        counter.change_current_fraction(h, 0.5);
+
+        advance_decay_generation(1);
+        counter.change_current_fraction(h, 0.98);
+        let increment = 0.001;
+        assert!(!counter.tick(h, increment));
+
+        let expected = (0.98f32 as f64 + increment) as f32;
+        let actual = counter_time(&counter, h);
+        assert!((actual - expected).abs() < 1.0e-6, "actual={actual}");
+    }
+
     #[test]
     fn test_auto_reset_on_fire() {
         let mut counter = JitCounter::new(DEFAULT_SIZE);

pyre/bench/synth/exception_traceback_loop_forms.wasm.jitstats

@@ -8,7 +8,8 @@ fbw_rolled_back_with_effects=0
 fbw_store_journal_rollback_failed=0
 field_pos_attached_misplaced=0
 field_pos_spec_misplaced=0
-guard_failures=810
+guard_failures=811
 internal_compile_panics=0
 loops_aborted=0
 loops_compiled=8
+retraces_compiled=0

pyre/bench/synth/inline_chain_depth_typeflip.cranelift.jitstats

@@ -1,4 +1,4 @@
-bridges_compiled=19
+bridges_compiled=18
 descr_set_absent=0
 descr_set_ambiguous=0
 descr_set_stale_absent=0
@@ -8,7 +8,8 @@ fbw_rolled_back_with_effects=0
 fbw_store_journal_rollback_failed=0
 field_pos_attached_misplaced=0
 field_pos_spec_misplaced=0
-guard_failures=3820
+guard_failures=3681
 internal_compile_panics=0
 loops_aborted=0
 loops_compiled=6
+retraces_compiled=0

pyre/bench/synth/inline_chain_depth_typeflip.dynasm.jitstats

@@ -1,4 +1,4 @@
-bridges_compiled=19
+bridges_compiled=18
 descr_set_absent=0
 descr_set_ambiguous=0
 descr_set_stale_absent=0
@@ -8,7 +8,8 @@ fbw_rolled_back_with_effects=0
 fbw_store_journal_rollback_failed=0
 field_pos_attached_misplaced=0
 field_pos_spec_misplaced=0
-guard_failures=3820
+guard_failures=3681
 internal_compile_panics=0
 loops_aborted=0
 loops_compiled=6
+retraces_compiled=0

pyre/check.py

@@ -1537,11 +1537,17 @@ def _jitstats_baseline_path(self, backend, script):
         # what let eight baselines go stale unnoticed.
         #
         # Measured, not assumed, and read back from the runner jobs rather than
-        # predicted. A fresh local build and ubuntu-24.04 report str_fstring
-        # guard_failures 658/659 on dynasm/cranelift. macos-latest reports the
-        # inverse 659/658, and windows-latest reports 658/658. Every other gated
-        # counter is identical (six loops and three bridges), and each mismatch
-        # reproduced in both stability reruns.
+        # predicted. ubuntu-24.04 reports str_fstring guard_failures 658/659 on
+        # dynasm/cranelift. macos-latest reports the inverse 659/658, and
+        # windows-latest reports 658/658. Every other gated counter is identical
+        # (six loops and three bridges), and each mismatch reproduced in both
+        # stability reruns.
+        #
+        # A local arm64 macOS build observes the macos-latest pair rather than
+        # ubuntu's: two full local gates pass against plain `.darwin` overlays
+        # holding dynasm 659 and cranelift 658. Darwin is therefore carried
+        # platform-wide, and the darwin GitHub-runner overlay that once held the
+        # same cranelift value was dropped as unreachable.
         #
         # This is a collection-schedule boundary, not a compile-shape change.
         # On a local arm64 macOS cranelift build, adding check.py's two wasm
@@ -1552,10 +1558,7 @@ def _jitstats_baseline_path(self, backend, script):
         # is where a collection-triggered breaker lands, not how many loops or
         # bridges compile. The runner did not enable either census, so exactly
         # which host input moves that placement remains unidentified; the
-        # observed counter and stable runner split are direct measurements. The
-        # fresh local macOS build matches ubuntu rather than the macOS runner,
-        # which is why these are GitHub-runner overlays rather than
-        # platform-wide ones.
+        # observed counter and stable runner split are direct measurements.
         #
         # An overlay shadows the shared baseline permanently, so one written
         # against a value that later converges becomes a failure main does not