From c6847128f3a872e0e084b2920bfcd21f881c69fa Mon Sep 17 00:00:00 2001
From: Andriy Berestovskyy <berestovskyy@gmail.com>
Date: Mon, 4 Nov 2024 09:53:38 +0100
Subject: [PATCH] fix: EXC-1775: Replace `delete_snapshots` functionality with
 `remove` (#2395)

Eliminate code duplication by replacing
`CanisterSnapshots::delete_snapshots` functionality with
`CanisterSnapshots::remove`.
---
 .../tests/execution_test.rs                   | 39 +++++++++++++++++++
 rs/replicated_state/src/canister_snapshots.rs |  7 +---
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/rs/execution_environment/tests/execution_test.rs b/rs/execution_environment/tests/execution_test.rs
index e55e03f5fab..aa5a45dec82 100644
--- a/rs/execution_environment/tests/execution_test.rs
+++ b/rs/execution_environment/tests/execution_test.rs
@@ -851,6 +851,45 @@ fn canister_snapshot_metrics_are_observed() {
     assert_eq!(gauge, 1.0);
 }
 
+#[test]
+fn canister_snapshot_metrics_are_consistent_after_canister_deletion() {
+    let env = StateMachineBuilder::new()
+        .with_subnet_type(SubnetType::Application)
+        .build();
+
+    let canister_id = create_universal_canister_with_cycles(
+        &env,
+        Some(CanisterSettingsArgsBuilder::new().build()),
+        INITIAL_CYCLES_BALANCE,
+    );
+
+    env.take_canister_snapshot(TakeCanisterSnapshotArgs::new(canister_id, None))
+        .unwrap();
+
+    let count = fetch_gauge(env.metrics_registry(), "scheduler_num_canister_snapshots").unwrap();
+    assert_eq!(count, 1.0);
+    let memory_usage = fetch_gauge(
+        env.metrics_registry(),
+        "scheduler_canister_snapshots_memory_usage_bytes",
+    )
+    .unwrap();
+    assert_gt!(memory_usage, 0.0);
+
+    env.stop_canister(canister_id)
+        .expect("Error stopping canister.");
+    env.delete_canister(canister_id)
+        .expect("Error deleting canister.");
+
+    let count = fetch_gauge(env.metrics_registry(), "scheduler_num_canister_snapshots").unwrap();
+    assert_eq!(count, 0.0);
+    let memory_usage = fetch_gauge(
+        env.metrics_registry(),
+        "scheduler_canister_snapshots_memory_usage_bytes",
+    )
+    .unwrap();
+    assert_eq!(memory_usage, 0.0);
+}
+
 fn assert_replied(result: Result<WasmResult, UserError>) {
     match result {
         Ok(wasm_result) => match wasm_result {
diff --git a/rs/replicated_state/src/canister_snapshots.rs b/rs/replicated_state/src/canister_snapshots.rs
index 3511c36e574..3370d109882 100644
--- a/rs/replicated_state/src/canister_snapshots.rs
+++ b/rs/replicated_state/src/canister_snapshots.rs
@@ -126,12 +126,9 @@ impl CanisterSnapshots {
     /// Additionally, new items are added to the `unflushed_changes`,
     /// representing the deleted backups since the last flush to the disk.
     pub fn delete_snapshots(&mut self, canister_id: CanisterId) {
-        if let Some(snapshot_ids) = self.snapshot_ids.remove(&canister_id) {
+        if let Some(snapshot_ids) = self.snapshot_ids.get(&canister_id).cloned() {
             for snapshot_id in snapshot_ids {
-                debug_assert!(self.snapshots.contains_key(&snapshot_id));
-                self.snapshots.remove(&snapshot_id).unwrap();
-                self.unflushed_changes
-                    .push(SnapshotOperation::Delete(snapshot_id));
+                self.remove(snapshot_id);
             }
         }
     }