chore: EXC-1691: Reserve cycles just for the canister snapshot increase

Closes EXC-1691

rs/execution_environment/src/canister_manager.rs

@@ -1789,7 +1789,7 @@ impl CanisterManager {
             return (Err(err), NumInstructions::new(0));
         };
 
-        match replace_snapshot {
+        let replace_snapshot_size = match replace_snapshot {
             // Check that replace snapshot ID exists if provided.
             Some(replace_snapshot) => {
                 match state.canister_snapshots.get(replace_snapshot) {
@@ -1814,6 +1814,7 @@ impl CanisterManager {
                                 NumInstructions::new(0),
                             );
                         }
+                        snapshot.size()
                     }
                 }
             }
@@ -1833,8 +1834,9 @@ impl CanisterManager {
                         NumInstructions::new(0),
                     );
                 }
+                0.into()
             }
-        }
+        };
 
         if self.config.rate_limiting_of_heap_delta == FlagStatus::Enabled
             && canister.scheduler_state.heap_delta_debit >= self.config.heap_delta_rate_limit
@@ -1850,6 +1852,11 @@ impl CanisterManager {
         }
 
         let new_snapshot_size = canister.snapshot_size_bytes();
+        let new_snapshot_diff = if new_snapshot_size > replace_snapshot_size {
+            new_snapshot_size - replace_snapshot_size
+        } else {
+            0.into()
+        };
 
         {
             // Run the following checks on memory usage and return an error
@@ -1861,7 +1868,7 @@ impl CanisterManager {
 
             // Calculate if any cycles will need to be reserved.
             let reservation_cycles = self.cycles_account_manager.storage_reservation_cycles(
-                new_snapshot_size,
+                new_snapshot_diff,
                 resource_saturation,
                 subnet_size,
             );
@@ -1915,13 +1922,13 @@ impl CanisterManager {
                         requested,
                         available,
                     } => CanisterManagerError::InsufficientCyclesInMemoryGrow {
-                        bytes: new_snapshot_size,
+                        bytes: new_snapshot_diff,
                         available,
                         threshold: requested,
                     },
                     ReservationError::ReservedLimitExceed { requested, limit } => {
                         CanisterManagerError::ReservedCyclesLimitExceededInMemoryGrow {
-                            bytes: new_snapshot_size,
+                            bytes: new_snapshot_diff,
                             requested,
                             limit,
                         }

rs/execution_environment/src/execution_environment/tests/canister_snapshots.rs

@@ -599,6 +599,101 @@ fn canister_request_take_canister_reserves_cycles() {
     );
 }
 
+#[test]
+fn canister_snapshot_reserves_cycles_difference() {
+    const CYCLES: Cycles = Cycles::new(200_000_000_000_000);
+    const CAPACITY: u64 = 2_000_000_000;
+    const THRESHOLD: u64 = CAPACITY / 4;
+    const WASM_PAGE_SIZE: u64 = 65_536;
+    // 7500 of stable memory pages is close to 500MB, but still leaves some room
+    // for Wasm memory of the universal canister.
+    const NUM_PAGES: u64 = 7_500;
+
+    let mut test = ExecutionTestBuilder::new()
+        .with_snapshots(FlagStatus::Enabled)
+        .with_heap_delta_rate_limit(NumBytes::new(1_000_000_000))
+        .with_subnet_execution_memory(CAPACITY as i64)
+        .with_subnet_memory_reservation(0)
+        .with_subnet_memory_threshold(THRESHOLD as i64)
+        .build();
+
+    let canister_id = test
+        .canister_from_cycles_and_binary(CYCLES, UNIVERSAL_CANISTER_WASM.into())
+        .unwrap();
+    test.canister_update_reserved_cycles_limit(canister_id, CYCLES)
+        .unwrap();
+    grow_stable_memory(&mut test, canister_id, WASM_PAGE_SIZE, NUM_PAGES);
+
+    // Get the reserve balance before taking a canister snapshot.
+    let initial_reserved_cycles = test
+        .canister_state(canister_id)
+        .system_state
+        .reserved_balance();
+    // Make sure there are no reserved cycles.
+    assert_eq!(initial_reserved_cycles, Cycles::zero());
+
+    // Take a snapshot 1 for the canister.
+    let args: TakeCanisterSnapshotArgs = TakeCanisterSnapshotArgs::new(canister_id, None);
+    let result = test
+        .subnet_message("take_canister_snapshot", args.encode())
+        .unwrap();
+    let snapshot_id_1 = CanisterSnapshotResponse::decode(&result.bytes())
+        .unwrap()
+        .snapshot_id();
+    let reserved_cycles_after_snapshot_1 = test
+        .canister_state(canister_id)
+        .system_state
+        .reserved_balance();
+
+    // Take a snapshot 2 for the canister by replacing previous snapshot.
+    let args: TakeCanisterSnapshotArgs =
+        TakeCanisterSnapshotArgs::new(canister_id, Some(snapshot_id_1));
+    let result = test
+        .subnet_message("take_canister_snapshot", args.encode())
+        .unwrap();
+    let snapshot_id_2 = CanisterSnapshotResponse::decode(&result.bytes())
+        .unwrap()
+        .snapshot_id();
+    let reserved_cycles_after_snapshot_2 = test
+        .canister_state(canister_id)
+        .system_state
+        .reserved_balance();
+    // Make sure the reserved cycles are the same.
+    assert_eq!(
+        reserved_cycles_after_snapshot_1,
+        reserved_cycles_after_snapshot_2
+    );
+
+    // Delete the Snapshot.
+    let args: DeleteCanisterSnapshotArgs =
+        DeleteCanisterSnapshotArgs::new(canister_id, snapshot_id_2);
+    test.subnet_message("delete_canister_snapshot", args.encode())
+        .unwrap();
+    let reserved_cycles_after_delete = test
+        .canister_state(canister_id)
+        .system_state
+        .reserved_balance();
+    // Make sure the reserved cycles are the same.
+    assert_eq!(
+        reserved_cycles_after_snapshot_2,
+        reserved_cycles_after_delete
+    );
+
+    // Take a new snapshot 3 for the canister.
+    let args: TakeCanisterSnapshotArgs = TakeCanisterSnapshotArgs::new(canister_id, None);
+    test.subnet_message("take_canister_snapshot", args.encode())
+        .unwrap();
+    let reserved_cycles_after_a_new_snapshot = test
+        .canister_state(canister_id)
+        .system_state
+        .reserved_balance();
+    // Make sure the reserved cycles are increased even more than before.
+    assert!(
+        reserved_cycles_after_a_new_snapshot
+            > reserved_cycles_after_snapshot_1 + reserved_cycles_after_snapshot_1
+    );
+}
+
 #[test]
 fn take_canister_snapshot_fails_subnet_memory_exceeded() {
     const CYCLES: Cycles = Cycles::new(20_000_000_000_000);