feat: Detect a stale cold-pool aggregate during checkpoint validation

rs/replicated_state/src/canister_states.rs

@@ -160,6 +160,11 @@ impl ColdStats {
 /// 2. every canister in the `cold` pool satisfies `CanisterState::is_cold()`;
 /// 3. `cold_stats` matches a fresh recomputation over the `cold` pool.
 ///
+/// Invariant (3) is *additionally checked* in release builds during checkpoint
+/// validation, by [`Self::validate_cold_stats`]. That check is advisory: it logs
+/// a critical error and increments a counter, and does not abort or otherwise
+/// alter the checkpoint.
+///
 /// Additionally, the **strict** partition invariant — that every canister in
 /// the `hot` pool does *not* satisfy `is_cold()` — holds after
 /// [`Self::try_cool_all`] /
@@ -645,6 +650,45 @@ impl CanisterStates {
         Ok(())
     }
 
+    /// Validates that `cold_stats` matches a fresh recomputation over the `cold`
+    /// pool, i.e. that the sub-before / add-after bracketing around every
+    /// cold-pool mutation has been respected.
+    ///
+    /// Unlike the `debug_assert` in `debug_assert_invariants`, this is intended to
+    /// run in release builds during checkpoint validation, because the aggregates
+    /// are read into hashed replicated state: `cold_stats.memory_usage` feeds
+    /// `total_canister_memory_usage()` and hence
+    /// `SubnetMetrics::canister_state_bytes`, which is certified in the state tree
+    /// and today has no other check. A stale aggregate is therefore a divergence
+    /// risk, and this is the point at which in-memory and derived could differ.
+    ///
+    /// It runs only for a *locally produced* checkpoint, i.e. on the branch of
+    /// `validate_and_finalize_checkpoint_and_remove_unverified_marker` that has a
+    /// reference state, and not on the state-sync path. That is the only branch
+    /// where it could find anything: a `CanisterStates` freshly loaded from disk has
+    /// `cold_stats` recomputed by `CanisterStates::new`, so it is consistent by
+    /// construction; only the in-memory reference state can have drifted.
+    ///
+    /// Note that the caller's failure mode is **advisory**: `validate_eq_checkpoint`
+    /// logs a critical error and increments a counter, then finalizes the
+    /// checkpoint regardless. This detects and attributes a stale aggregate; it
+    /// does not prevent one from being used. The caller runs it *after* the
+    /// per-canister comparison and combines the two errors, so that an advisory
+    /// failure here does not mask the diagnostics that identify which canister
+    /// drifted.
+    ///
+    /// Complexity: `O(|cold canisters|)`.
+    pub fn validate_cold_stats(&self) -> Result<(), String> {
+        let recomputed = ColdStats::recompute(self.cold.values());
+        if recomputed != self.cold_stats {
+            return Err(format!(
+                "cold_stats out of sync with the cold pool: stored {:?}, recomputed {:?}",
+                self.cold_stats, recomputed
+            ));
+        }
+        Ok(())
+    }
+
     /// Debug-only consistency check, called at the end of every mutating operation.
     /// Verifies invariants (1)–(3) listed under [`CanisterStates`].
     ///

rs/replicated_state/src/canister_states/tests.rs

@@ -948,3 +948,67 @@ fn validate_strict_split_rejects_stale_cold_canister() {
         "unexpected error: {err}",
     );
 }
+
+/// Consumes `amount` cycles on `canister`, as storage / instruction charging
+/// does. Consuming cycles does not create work, so a cold canister stays cold.
+fn consume_cycles(canister: &mut Arc<CanisterState>, amount: u128) {
+    use ic_types_cycles::{CanisterCyclesCostSchedule, CompoundCycles, Instructions};
+
+    Arc::make_mut(canister)
+        .system_state
+        .consume_cycles(CompoundCycles::<Instructions>::new(
+            Cycles::new(amount),
+            CanisterCyclesCostSchedule::Normal,
+        ));
+}
+
+/// Folds `consumed_cycles` over every canister, hot and cold, without going
+/// through the `cold_stats` aggregate.
+fn direct_consumed_cycles_fold(states: &CanisterStates) -> ic_types_cycles::NominalCycles {
+    use ic_types_cycles::NominalCycles;
+
+    states
+        .all_values()
+        .fold(NominalCycles::zero(), |acc, canister| {
+            acc + canister.system_state.canister_metrics().consumed_cycles()
+        })
+}
+
+#[test]
+fn validate_cold_stats_accepts_consistent_stats() {
+    let mut states = CanisterStates::default();
+    states.insert(cold_canister(1));
+    states.insert(hot_canister(2));
+    states.insert(cold_canister(3));
+
+    assert_eq!(states.validate_cold_stats(), Ok(()));
+}
+
+#[test]
+fn validate_cold_stats_rejects_stale_stats() {
+    use ic_types_cycles::{NominalCycles, NominalCyclesTesting};
+
+    let mut states = CanisterStates::default();
+    let c = cold_canister(1);
+    states.insert(Arc::clone(&c));
+    assert_eq!(states.validate_cold_stats(), Ok(()));
+
+    // Bypass the public mutation entry points: mutate a cold canister's consumed
+    // cycles directly, behind the aggregate's back, simulating missing
+    // sub-before / add-after bracketing.
+    consume_cycles(states.cold.get_mut(&c.canister_id()).unwrap(), 42);
+    assert_eq!(states.hot.len(), 0);
+    assert_eq!(states.cold.len(), 1);
+
+    let err = states.validate_cold_stats().unwrap_err();
+    assert!(
+        err.contains("cold_stats out of sync with the cold pool"),
+        "unexpected error: {err}",
+    );
+    // The aggregate is stale, so the reported total is now wrong.
+    assert_eq!(states.cold_stats.consumed_cycles, NominalCycles::new(0));
+    assert_ne!(
+        states.total_consumed_cycles(),
+        direct_consumed_cycles_fold(&states)
+    );
+}

rs/state_manager/src/checkpoint.rs

@@ -555,31 +555,51 @@ impl CheckpointLoader {
                 .or_default()
                 .push(snapshot_id);
         }
-        maybe_parallel_map(thread_pool, ref_canister_ids.iter(), |canister_id| {
-            load_canister_state_from_checkpoint(
-                &self.checkpoint_layout,
-                canister_id,
-                snapshot_ids_per_canister
-                    .get(canister_id)
-                    .cloned()
-                    .unwrap_or_default(),
-                Arc::clone(&self.fd_factory),
-                &self.metrics,
-            )
-            .map_err(|err| {
-                format!(
-                    "Failed to load canister state for validation for key #{canister_id}: {err}"
+        let per_canister =
+            maybe_parallel_map(thread_pool, ref_canister_ids.iter(), |canister_id| {
+                load_canister_state_from_checkpoint(
+                    &self.checkpoint_layout,
+                    canister_id,
+                    snapshot_ids_per_canister
+                        .get(canister_id)
+                        .cloned()
+                        .unwrap_or_default(),
+                    Arc::clone(&self.fd_factory),
+                    &self.metrics,
                 )
-            })?
-            .0
-            .validate_eq(
-                ref_canister_states
-                    .get(canister_id)
-                    .expect("Failed to get canister from canister_states"),
-            )
-        })
-        .into_iter()
-        .try_for_each(identity)
+                .map_err(|err| {
+                    format!(
+                        "Failed to load canister state for validation for key #{canister_id}: {err}"
+                    )
+                })?
+                .0
+                .validate_eq(
+                    ref_canister_states
+                        .get(canister_id)
+                        .expect("Failed to get canister from canister_states"),
+                )
+            })
+            .into_iter()
+            .try_for_each(identity);
+
+        // Detect (and attribute) a stale cold-pool aggregate. Like every other
+        // check here, this is advisory: the caller logs a critical error and
+        // increments a counter, then finalizes the checkpoint regardless.
+        //
+        // Deliberately run *after* the per-canister comparison above, and combined
+        // with it rather than short-circuiting it: in the very scenario where this
+        // check fires, the per-canister diagnostics are what tell the operator
+        // *which* canister drifted, and an advisory check must not cost the
+        // operator that information.
+        let cold_stats = ref_canister_states
+            .validate_cold_stats()
+            .map_err(|err| format!("Canister Validation: {err}"));
+
+        match (per_canister, cold_stats) {
+            (Ok(()), Ok(())) => Ok(()),
+            (Err(err), Ok(())) | (Ok(()), Err(err)) => Err(err),
+            (Err(per_canister), Err(cold_stats)) => Err(format!("{per_canister}; {cold_stats}")),
+        }
     }
 
     fn validate_eq_canister_snapshots_ids(