fix: tolerate a cycles balance increase on DTS resume (#11141)

Every kind of paused execution re-creates its helper from the current
clean canister state when it is resumed and thus compares the cycles
balance of that state with the balance at the start of the DTS
execution. The comparison used to reject any change of the balance; it
now only rejects a decrease, which the recorded steps replayed on the
clean canister state might no longer be able to cover.

An increase is safe: all cycles changes of a DTS execution are applied
relative to the balance of the clean canister state (the Wasm execution
reports a `CyclesBalanceChange` delta and the prepaid execution cycles
are refunded relative as well), so the additional cycles are preserved.

The new tests `dts_resume_succeeds_after_cycles_increase` and
`dts_install_code_resume_succeeds_after_cycles_increase` are the
counterparts of `dts_resume_fails_due_to_cycles_decrease` and
`dts_install_code_resume_fails_due_to_cycles_decrease`: they add cycles
to the canister while its execution is paused and assert that the
execution completes and that the added cycles are not lost. The former
runs every scenario twice, once without adding cycles, and asserts that
the two final balances differ by exactly the added cycles. The setup of
all the scenarios is now shared by the tests that decrease and increase
the cycles balance.

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: mraszyk

rs/execution_environment/src/execution/call_or_task.rs

@@ -443,8 +443,10 @@ impl CallOrTaskHelper {
     }
 
     /// Replays the previous update call steps on the given clean canister.
-    /// Returns an error if any step fails. Otherwise, it returns an instance of
-    /// the helper that can be used to continue the update call execution.
+    /// Returns an error if the cycles balance of the clean canister dropped
+    /// below the cycles balance at the start of the DTS execution or if any step
+    /// fails. Otherwise, it returns an instance of the helper that can be used
+    /// to continue the update call execution.
     fn resume(
         clean_canister: &CanisterState,
         original: &OriginalContext,
@@ -453,7 +455,13 @@ impl CallOrTaskHelper {
     ) -> Result<Self, UserError> {
         let mut helper = Self::new(clean_canister, original, deallocation_sender)?;
         helper.executed_wasm_instructions = paused.executed_wasm_instructions;
-        if helper.initial_cycles_balance != paused.initial_cycles_balance {
+        // The cycles balance of the clean canister must not decrease during the
+        // DTS execution: the recorded steps are replayed on the clean canister
+        // state and a lower balance might no longer be able to cover them.
+        // An increase is safe: all cycles changes of the DTS execution are
+        // applied relative to the balance of the clean canister state and hence
+        // the additional cycles are preserved.
+        if helper.initial_cycles_balance < paused.initial_cycles_balance {
             let msg = match original.call_or_task {
                 CanisterCallOrTask::Update(_) => {
                     "Mismatch in cycles balance when resuming an update call".to_string()

rs/execution_environment/src/execution/install_code.rs

@@ -249,7 +249,7 @@ impl InstallCodeHelper {
     }
 
     /// Replays the previous `install_code` steps on the given clean canister.
-    /// Returns an error if the cycles balance of the clean canister differs from
+    /// Returns an error if the cycles balance of the clean canister dropped below
     /// the cycles balance at the start of the DTS execution or if any step
     /// fails. Otherwise, it returns an instance of the helper that can be used
     /// to continue the `install_code` execution.
@@ -281,9 +281,13 @@ impl InstallCodeHelper {
                 .saturating_sub(executed_wasm_instructions.get()),
         );
 
-        // The cycles balance of the clean canister must not change during the
-        // DTS execution.
-        if helper.initial_cycles_balance != paused.initial_cycles_balance {
+        // The cycles balance of the clean canister must not decrease during the
+        // DTS execution: the recorded steps are replayed on the clean canister
+        // state and a lower balance might no longer be able to cover them.
+        // An increase is safe: all cycles changes of the DTS execution are
+        // applied relative to the balance of the clean canister state and hence
+        // the additional cycles are preserved.
+        if helper.initial_cycles_balance < paused.initial_cycles_balance {
             let msg = "Mismatch in cycles balance when resuming an install code".to_string();
             let err = HypervisorError::WasmEngineError(FailedToApplySystemChanges(msg));
             let err = (clean_canister.canister_id(), err).into();

rs/execution_environment/src/execution/install_code/tests.rs

@@ -23,7 +23,7 @@ use ic_test_utilities_metrics::fetch_int_counter;
 use ic_types::ingress::{IngressState, IngressStatus, WasmResult};
 use ic_types::messages::MessageId;
 use ic_types::{CanisterId, ComputeAllocation, MemoryAllocation, NumBytes, NumInstructions};
-use ic_types_cycles::{Cycles, CyclesUseCase, NominalCycles};
+use ic_types_cycles::{CompoundCycles, Cycles, CyclesUseCase, Instructions, NominalCycles};
 use ic_types_test_utils::ids::{canister_test_id, subnet_test_id, user_test_id};
 use ic_universal_canister::{UNIVERSAL_CANISTER_WASM, call_args, wasm};
 use maplit::btreemap;
@@ -139,19 +139,38 @@ fn consumed_cycles_for_instructions(
         .unwrap_or_default()
 }
 
-/// Analogously to `dts_resume_fails_due_to_cycles_decrease` for calls,
-/// replicated queries, callbacks, and tasks, resuming a paused `install_code`
-/// whose canister lost cycles while it was paused fails instead of replaying the
-/// recorded steps on a balance that can no longer cover them. The failed
-/// execution is charged for exactly the instructions it had already executed,
-/// including those of the paused Wasm execution.
-#[test]
-fn dts_install_code_resume_fails_due_to_cycles_decrease() {
-    const INSTRUCTION_LIMIT: u64 = 50_000_000;
-    const SLICE_INSTRUCTION_LIMIT: u64 = 132_000;
+/// The instruction limits of the DTS `install_code` tests that change the cycles
+/// balance of the canister while its execution is paused.
+const DTS_INSTALL_CODE_INSTRUCTION_LIMIT: u64 = 50_000_000;
+const DTS_INSTALL_CODE_SLICE_INSTRUCTION_LIMIT: u64 = 132_000;
+
+/// A canister with a paused `install_code` execution, along with a snapshot of
+/// the accounting counters of that canister taken before the execution started.
+///
+/// Shared by the tests that decrease and increase the cycles balance of the
+/// canister while its `install_code` execution is paused.
+struct PausedInstallCode {
+    test: ExecutionTest,
+    canister_id: CanisterId,
+    /// The ingress message of the `install_code` subnet message.
+    ingress_id: MessageId,
+    /// The cycles balance before the execution cycles were prepaid.
+    original_balance: Cycles,
+    /// The cycles consumed for instructions.
+    original_consumed_cycles: NominalCycles,
+    /// The instructions executed by all the slices of the canister.
+    original_executed_instructions: NumInstructions,
+    /// The accumulated cost of the instructions executed by the canister.
+    original_execution_cost: CompoundCycles<Instructions>,
+}
+
+/// Starts an `install_code` execution that pauses after its first slice: that
+/// slice compiles the Wasm module and executes its `(start)` function, which
+/// together exceed the slice instruction limit.
+fn install_code_paused_after_first_slice() -> PausedInstallCode {
     let mut test = ExecutionTestBuilder::new()
-        .with_install_code_instruction_limit(INSTRUCTION_LIMIT)
-        .with_install_code_slice_instruction_limit(SLICE_INSTRUCTION_LIMIT)
+        .with_install_code_instruction_limit(DTS_INSTALL_CODE_INSTRUCTION_LIMIT)
+        .with_install_code_slice_instruction_limit(DTS_INSTALL_CODE_SLICE_INSTRUCTION_LIMIT)
         .with_create_execution_state_base_cost(0)
         .with_manual_execution()
         .build();
@@ -169,6 +188,7 @@ fn dts_install_code_resume_fails_due_to_cycles_decrease() {
     let original_balance = test.canister_state(canister_id).system_state.balance();
     let original_consumed_cycles = consumed_cycles_for_instructions(&test, canister_id);
     let original_executed_instructions = test.canister_executed_instructions(canister_id);
+    let original_execution_cost = test.canister_execution_cost(canister_id);
 
     let ingress_id = test.dts_install_code(payload);
 
@@ -187,13 +207,42 @@ fn dts_install_code_resume_fails_due_to_cycles_decrease() {
             - test
                 .cycles_account_manager()
                 .execution_cost(
-                    NumInstructions::from(INSTRUCTION_LIMIT),
+                    NumInstructions::from(DTS_INSTALL_CODE_INSTRUCTION_LIMIT),
                     test.get_own_subnet_cycles_config(),
                     WASM_EXECUTION_MODE,
                 )
                 .real(),
     );
 
+    PausedInstallCode {
+        test,
+        canister_id,
+        ingress_id,
+        original_balance,
+        original_consumed_cycles,
+        original_executed_instructions,
+        original_execution_cost,
+    }
+}
+
+/// Analogously to `dts_resume_fails_due_to_cycles_decrease` for calls,
+/// replicated queries, callbacks, and tasks, resuming a paused `install_code`
+/// whose canister lost cycles while it was paused fails instead of replaying the
+/// recorded steps on a balance that can no longer cover them. The failed
+/// execution is charged for exactly the instructions it had already executed,
+/// including those of the paused Wasm execution.
+#[test]
+fn dts_install_code_resume_fails_due_to_cycles_decrease() {
+    let PausedInstallCode {
+        mut test,
+        canister_id,
+        ingress_id,
+        original_balance,
+        original_consumed_cycles,
+        original_executed_instructions,
+        ..
+    } = install_code_paused_after_first_slice();
+
     // Decrease the cycles balance of the clean canister.
     test.canister_state_mut(canister_id)
         .system_state
@@ -224,7 +273,10 @@ fn dts_install_code_resume_fails_due_to_cycles_decrease() {
     // more than the slice instruction limit.
     let executed_instructions =
         test.canister_executed_instructions(canister_id) - original_executed_instructions;
-    assert_gt!(executed_instructions.get(), SLICE_INSTRUCTION_LIMIT);
+    assert_gt!(
+        executed_instructions.get(),
+        DTS_INSTALL_CODE_SLICE_INSTRUCTION_LIMIT
+    );
 
     // The canister is charged exactly the cost of those instructions, including
     // the instructions of the paused Wasm execution: the rest of the prepaid
@@ -245,6 +297,62 @@ fn dts_install_code_resume_fails_due_to_cycles_decrease() {
     );
 }
 
+/// Counterpart of `dts_install_code_resume_fails_due_to_cycles_decrease`: while
+/// resuming a paused `install_code` whose canister lost cycles fails, an
+/// increase of the cycles balance while the execution is paused is tolerated and
+/// the additional cycles are not lost when the execution completes.
+#[test]
+fn dts_install_code_resume_succeeds_after_cycles_increase() {
+    const CYCLES_ADDED_WHILE_PAUSED: Cycles = Cycles::new(1_234_567_890);
+
+    let PausedInstallCode {
+        mut test,
+        canister_id,
+        ingress_id,
+        original_balance,
+        original_executed_instructions,
+        original_execution_cost,
+        ..
+    } = install_code_paused_after_first_slice();
+
+    // Increase the cycles balance of the clean canister.
+    test.canister_state_mut(canister_id)
+        .system_state
+        .add_cycles(CYCLES_ADDED_WHILE_PAUSED);
+
+    // The remaining slices resume the paused execution, which completes.
+    while test.canister_state(canister_id).next_execution() == NextExecution::ContinueInstallCode {
+        test.execute_slice(canister_id);
+    }
+    assert_eq!(
+        test.canister_state(canister_id).next_execution(),
+        NextExecution::None
+    );
+
+    let result = check_ingress_status(test.ingress_status(&ingress_id)).unwrap();
+    assert_eq!(result, WasmResult::Reply(EmptyBlob.encode()));
+
+    // The code has been installed.
+    assert!(test.canister_state(canister_id).execution_state.is_some());
+
+    // The execution spanned multiple slices.
+    let executed_instructions =
+        test.canister_executed_instructions(canister_id) - original_executed_instructions;
+    assert_gt!(
+        executed_instructions.get(),
+        DTS_INSTALL_CODE_SLICE_INSTRUCTION_LIMIT
+    );
+
+    // The canister is charged exactly the cost of the executed instructions: the
+    // cycles added while the execution was paused are not lost.
+    assert_eq!(
+        test.canister_state(canister_id).system_state.balance(),
+        original_balance
+            - (test.canister_execution_cost(canister_id) - original_execution_cost).real()
+            + CYCLES_ADDED_WHILE_PAUSED
+    );
+}
+
 #[test]
 fn dts_abort_works_in_install_code() {
     const INSTRUCTION_LIMIT: u64 = 50_000_000;

rs/execution_environment/src/execution/response.rs

@@ -339,8 +339,8 @@ impl ResponseHelper {
     /// call context, and execution state because it is not possible to invoke
     /// the cleanup callback in such cases.
     ///
-    /// It returns an error if the cycles balance of the clean canister differs
-    /// from the cycles balances at the start of the DTS execution.
+    /// It returns an error if the cycles balance of the clean canister dropped
+    /// below the cycles balance at the start of the DTS execution.
     #[allow(clippy::result_large_err)]
     fn resume(
         paused: PausedResponseHelper,
@@ -387,9 +387,13 @@ impl ResponseHelper {
             .validate(&call_context, original, round, round_limits)
             .expect("Failed to resume DTS response: validation");
 
-        // The cycles balance of the clean canister must not change during the
-        // DTS execution.
-        if helper.initial_cycles_balance != paused.initial_cycles_balance {
+        // The cycles balance of the clean canister must not decrease during the
+        // DTS execution: the initial steps are replayed on the clean canister
+        // state and a lower balance might no longer be able to cover them.
+        // An increase is safe: all cycles changes of the DTS execution are
+        // applied relative to the balance of the clean canister state and hence
+        // the additional cycles are preserved.
+        if helper.initial_cycles_balance < paused.initial_cycles_balance {
             let msg = "Mismatch in cycles balance when resuming a response call".to_string();
             let err = HypervisorError::WasmEngineError(FailedToApplySystemChanges(msg));
             return Err((helper, err));