fix(session): stop a lazy gate initialisation from reverting a forced one

crates/onnx-runtime-session/src/executor/mod.rs

@@ -102,29 +102,60 @@ mod phase_profile {
     use std::sync::{Mutex, OnceLock};
     use std::time::Instant;
 
-    static STATE: AtomicU8 = AtomicU8::new(0); // 0 = unknown, 1 = off, 2 = on
+    pub(super) const UNKNOWN: u8 = 0;
+    pub(super) const OFF: u8 = 1;
+    pub(super) const ON: u8 = 2;
+
+    static STATE: AtomicU8 = AtomicU8::new(UNKNOWN);
     static PRINTED: AtomicBool = AtomicBool::new(false);
 
+    /// Publish an environment-derived gate value without overwriting a decision
+    /// that is already in force, and report the value that actually is.
+    ///
+    /// The lazy initialiser below is a *writer*, not just a reader, and it runs
+    /// on whatever thread first reaches the gate while holding no lock -
+    /// `run.rs` consults the planner gate on every executor run, so under the
+    /// parallel test runner that is any thread at all. Its load / read-env /
+    /// store is not atomic: an unconditional `store` can land *after* a
+    /// `force_*` call that did take [`globals_lock`], silently reverting it.
+    /// That is not hypothetical - it is how
+    /// `activation_memory_planner_reports_static_decode_graph_savings` failed
+    /// on `Rust coverage (Windows x86_64)`: the gate it forced on was reset to
+    /// `OFF` between the force and the run, so the run published no stats and
+    /// the `expect` on them panicked.
+    ///
+    /// [`globals_lock`] cannot fix that, because the losing writer is a plain
+    /// reader that must not be made to take a lock on the hot path. A
+    /// `compare_exchange` from [`UNKNOWN`] fixes it instead, by making the
+    /// initialiser lose the race rather than win it: it publishes only when
+    /// nothing else has, and otherwise reports what is in force.
+    pub(super) fn publish_env_derived(gate: &AtomicU8, on: bool) -> bool {
+        let desired = if on { ON } else { OFF };
+        match gate.compare_exchange(UNKNOWN, desired, Ordering::Relaxed, Ordering::Relaxed) {
+            Ok(_) => on,
+            Err(in_force) => in_force == ON,
+        }
+    }
+
     pub fn enabled() -> bool {
         match STATE.load(Ordering::Relaxed) {
-            1 => false,
-            2 => true,
+            OFF => false,
+            ON => true,
             _ => {
                 let on = std::env::var("NXRT_EXEC_PHASE_PROFILE")
                     .is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true"));
-                STATE.store(if on { 2 } else { 1 }, Ordering::Relaxed);
-                on
+                publish_env_derived(&STATE, on)
             }
         }
     }
 
     /// Test-only override of the env-derived enable state.
     #[cfg(test)]
     pub(super) fn force_enabled(on: bool) {
-        STATE.store(if on { 2 } else { 1 }, Ordering::Relaxed);
+        STATE.store(if on { ON } else { OFF }, Ordering::Relaxed);
     }
 
-    static PLAN_STATE: AtomicU8 = AtomicU8::new(0); // 0 = unknown, 1 = off, 2 = on
+    static PLAN_STATE: AtomicU8 = AtomicU8::new(UNKNOWN);
 
     /// Whether to run the activation-memory *planner* during a run.
     ///
@@ -144,30 +175,36 @@ mod phase_profile {
     /// path to these stats - keeps getting them.
     pub fn activation_plan_enabled() -> bool {
         match PLAN_STATE.load(Ordering::Relaxed) {
-            1 => false,
-            2 => true,
+            OFF => false,
+            ON => true,
             _ => {
                 let on = ["NXRT_ACTIVATION_MEMORY_PLAN", "NXRT_EXEC_PHASE_PROFILE"]
                     .iter()
                     .any(|key| {
                         std::env::var(key).is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
                     });
-                PLAN_STATE.store(if on { 2 } else { 1 }, Ordering::Relaxed);
-                on
+                publish_env_derived(&PLAN_STATE, on)
             }
         }
     }
 
     /// Opt a process into activation-memory planning without turning on phase
     /// profiling, for callers that want the stats and are willing to pay.
     pub fn enable_activation_plan_for_process() {
-        PLAN_STATE.store(2, Ordering::Relaxed);
+        PLAN_STATE.store(ON, Ordering::Relaxed);
     }
 
     /// Test-only override of the env-derived planner state.
     #[cfg(test)]
     pub(super) fn force_activation_plan_enabled(on: bool) {
-        PLAN_STATE.store(if on { 2 } else { 1 }, Ordering::Relaxed);
+        PLAN_STATE.store(if on { ON } else { OFF }, Ordering::Relaxed);
+    }
+
+    /// The live planner gate, so a test can drive the real one rather than a
+    /// copy of it.
+    #[cfg(test)]
+    pub(super) fn activation_plan_gate() -> &'static AtomicU8 {
+        &PLAN_STATE
     }
 
     /// Holds the planner on for one test and clears it on drop, so a leaked
@@ -199,14 +236,21 @@ mod phase_profile {
     /// Both are plain atomics with no per-test isolation, so without this the
     /// parallel runner lets one test's `enable`/`force` land inside another
     /// test's assertion window.
+    ///
+    /// It does **not** serialise every writer, and cannot: the lazy
+    /// initialiser in `enabled`/`activation_plan_enabled` also writes, from any
+    /// thread that reaches a gate first, and making a hot-path reader take a
+    /// mutex is not an option. Holding this lock is therefore *not* sufficient
+    /// to own a gate - see [`publish_env_derived`], which is what keeps that
+    /// third writer from reverting a forced value.
     #[cfg(test)]
     pub(super) fn globals_lock() -> std::sync::MutexGuard<'static, ()> {
         static LOCK: Mutex<()> = Mutex::new(());
         LOCK.lock().unwrap_or_else(|e| e.into_inner())
     }
 
     pub fn enable_for_process() {
-        STATE.store(2, Ordering::Relaxed);
+        STATE.store(ON, Ordering::Relaxed);
     }
 
     /// Test-only snapshot of a phase's accumulated `(total_ns, count)`.

crates/onnx-runtime-session/src/executor/tests.rs

@@ -1,4 +1,4 @@
-use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, AtomicUsize, Ordering};
 
 use onnx_runtime_ep_api::{
     CaptureSupport, Cost, EpConfig, EpError, ExecutionProviderCapabilities, Fence, Kernel,
@@ -266,6 +266,70 @@ fn device_validation_reset_failure_is_not_ignored_after_a_run() {
     assert_eq!(ep.resets.load(Ordering::Relaxed), 2);
 }
 
+#[test]
+fn a_late_lazy_gate_initialisation_cannot_revert_a_forced_gate() {
+    // Regression test for an intermittent failure of
+    // `activation_memory_planner_reports_static_decode_graph_savings` on
+    // `Rust coverage (Windows x86_64)`: it forced the planner gate on, and the
+    // run that followed published no stats, so the `expect` on them panicked.
+    //
+    // The third writer is the cause. `globals_lock` serialises the *tests* that
+    // force these gates, but `activation_plan_enabled` also writes, from any
+    // thread, without the lock - `run.rs` consults it on every executor run, so
+    // under the parallel runner a sibling test is a writer too. Its
+    // load / read-env / store is not atomic, so with an unconditional `store`
+    // its answer could land after a forced value and revert it.
+    //
+    // Driven through the real gate and the real publish path rather than a
+    // copy, so the test cannot pass against a seam that has drifted from the
+    // code it stands in for.
+    let _globals = phase_profile::globals_lock();
+
+    // The interleaving that failed: a reader entered the initialiser, was
+    // descheduled while reading the environment, and published `off` after the
+    // force. Both env answers, because the defect is in publishing at all.
+    for env_says in [false, true] {
+        phase_profile::force_activation_plan_enabled(true);
+        assert!(
+            phase_profile::publish_env_derived(phase_profile::activation_plan_gate(), env_says),
+            "a late initialiser must report the value in force, not its own"
+        );
+        assert!(
+            phase_profile::activation_plan_enabled(),
+            "a lazy initialisation that lands after a force (env_says={env_says}) \
+             reverted it; the forcing test then runs with the planner off and \
+             fails on stats it correctly asked for"
+        );
+    }
+
+    // Non-vacuity: the initialiser must still publish when nothing is in force,
+    // or the assertions above would hold for a `publish` that never stores.
+    // On a scratch gate, so an unowned process-global is never left `UNKNOWN`
+    // for a concurrent reader to re-derive.
+    for env_says in [false, true] {
+        let unowned = AtomicU8::new(phase_profile::UNKNOWN);
+        assert_eq!(
+            phase_profile::publish_env_derived(&unowned, env_says),
+            env_says,
+            "an uncontended initialiser must report its own env-derived answer"
+        );
+        let expected = if env_says {
+            phase_profile::ON
+        } else {
+            phase_profile::OFF
+        };
+        assert_eq!(
+            unowned.load(Ordering::Relaxed),
+            expected,
+            "an uncontended initialiser must publish, or every later reader \
+             pays the environment read again"
+        );
+    }
+
+    // Leave the planner off, as every other test in this file expects.
+    phase_profile::force_activation_plan_enabled(false);
+}
+
 #[test]
 fn phase_profile_gating_and_accumulation() {
     // The process-global gates have no per-test isolation, so every test that