feat: [DSM-103] Charge for storage after aborting DTS executions

rs/execution_environment/src/scheduler.rs

@@ -1457,14 +1457,6 @@ impl Scheduler for SchedulerImpl {
                 let _timer = self.metrics.round_finalization_ingress.start_timer();
                 final_state.prune_ingress_history();
             }
-            {
-                let _timer = self.metrics.round_finalization_charge.start_timer();
-                self.charge_canisters_for_resource_allocation_and_usage(
-                    &mut final_state,
-                    registry_settings.subnet_size,
-                    current_round,
-                );
-            }
 
             // Update canister priorities.
             {
@@ -1474,6 +1466,16 @@ impl Scheduler for SchedulerImpl {
 
             self.finish_round(&mut final_state, current_round_type);
 
+            // Charge canisters after (some) paused executions were aborted.
+            {
+                let _timer = self.metrics.round_finalization_charge.start_timer();
+                self.charge_canisters_for_resource_allocation_and_usage(
+                    &mut final_state,
+                    registry_settings.subnet_size,
+                    current_round,
+                );
+            }
+
             final_state
                 .metadata
                 .subnet_metrics

rs/execution_environment/tests/dts.rs

@@ -24,7 +24,6 @@ use ic_types_cycles::Cycles;
 use ic_universal_canister::{
     CallArgs, UNIVERSAL_CANISTER_NO_HEARTBEAT_WASM, UNIVERSAL_CANISTER_WASM, call_args, wasm,
 };
-use more_asserts::assert_ge;
 use std::sync::OnceLock;
 
 const INITIAL_CYCLES_BALANCE: Cycles = Cycles::new(100_000_000_000_000);
@@ -1794,8 +1793,11 @@ fn dts_ingress_status_of_update_with_call_is_correct() {
     env.await_ingress(update, 100).unwrap();
 }
 
-#[test]
-fn dts_canister_uninstalled_due_to_resource_charges_with_aborted_updrade() {
+/// Causes a canister to run out of cycles with a paused / aborted upgrade and
+/// returns the result of the upgrade.
+fn impl_dts_canister_with_upgrade_is_uninstalled_due_to_resource_charges(
+    aborted: bool,
+) -> Result<WasmResult, UserError> {
     let env = dts_env(
         NumInstructions::from(1_000_000_000),
         NumInstructions::from(10_000),
@@ -1815,108 +1817,98 @@ fn dts_canister_uninstalled_due_to_resource_charges_with_aborted_updrade() {
         .install_canister_with_cycles(binary.clone(), vec![], settings, INITIAL_CYCLES_BALANCE)
         .unwrap();
 
-    // Make sure that the upgrade message gets aborted after each round.
-    env.set_checkpoints_enabled(true);
+    // Checkpointing will abort the upgrade after each round.
+    env.set_checkpoints_enabled(aborted);
 
     let upgrade = {
         let args = InstallCodeArgs::new(CanisterInstallMode::Upgrade, canister, binary, vec![]);
         env.send_ingress(user_id, IC_00, Method::InstallCode, args.encode())
     };
-
     env.tick();
 
-    // Advance the time so that the canister gets uninstalled due to the
+    // Advance the time enough to get the canister uninstalled due to the
     // resource usage.
     env.advance_time(Duration::from_secs(10_000_000));
-
     env.tick();
 
     // Enable normal message execution.
     env.set_checkpoints_enabled(false);
 
-    let result = env.await_ingress(upgrade, 30).unwrap();
+    env.await_ingress(upgrade, 30)
+}
 
-    // The canister is uninstalled after the execution completes because an
-    // aborted install_code is always restarted and becomes a paused execution
-    // by the time we charge canister for resource allocation.
-    assert_eq!(result, WasmResult::Reply(EmptyBlob.encode()));
+#[test]
+fn dts_canister_with_aborted_upgrade_is_uninstalled_due_to_resource_charges() {
+    let result = impl_dts_canister_with_upgrade_is_uninstalled_due_to_resource_charges(true);
+
+    // Canister with aborted upgrade is uninstalled. Upgrade fails.
+    assert_matches!(
+        result,
+        Err(err) if err.code() == ErrorCode::CanisterWasmModuleNotFound
+    );
 }
 
 #[test]
-fn dts_canister_uninstalled_due_resource_charges_with_aborted_update() {
+fn dts_canister_with_paused_upgrade_is_not_uninstalled_due_to_resource_charges() {
+    let result = impl_dts_canister_with_upgrade_is_uninstalled_due_to_resource_charges(false);
+
+    // Canister is only uninstalled after paused upgrade completes successfully.
+    assert_eq!(result, Ok(WasmResult::Reply(EmptyBlob.encode())));
+}
+
+fn impl_dts_canister_with_update_is_uninstalled_due_to_resource_charges(
+    aborted: bool,
+) -> Result<WasmResult, UserError> {
     let env = dts_env(
         NumInstructions::from(1_000_000_000),
         NumInstructions::from(10_000),
     );
-
     let binary = wat2wasm(DTS_WAT);
-
     let user_id = PrincipalId::new_anonymous();
 
-    let n = 10;
-
-    let mut canisters = vec![];
-    for _ in 0..n {
-        let settings = Some(
-            CanisterSettingsArgsBuilder::new()
-                .with_compute_allocation(1)
-                .build(),
-        );
-
-        let id = env
-            .install_canister_with_cycles(binary.clone(), vec![], settings, INITIAL_CYCLES_BALANCE)
-            .unwrap();
-        canisters.push(id);
-    }
-
-    // Make sure that the update messages get aborted after each round.
-    env.set_checkpoints_enabled(true);
+    let settings = Some(
+        CanisterSettingsArgsBuilder::new()
+            .with_compute_allocation(1)
+            .build(),
+    );
+    let canister = env
+        .install_canister_with_cycles(binary.clone(), vec![], settings, INITIAL_CYCLES_BALANCE)
+        .unwrap();
 
-    let mut updates = vec![];
-    for canister in canisters.iter() {
-        let update = env.send_ingress(user_id, *canister, "update", vec![]);
-        updates.push(update);
-    }
+    // Checkpointing will abort the update after each round.
+    env.set_checkpoints_enabled(aborted);
 
-    // Ensure that each update message starts executing.
-    for _ in 0..n {
-        env.tick();
-    }
+    let update = env.send_ingress(user_id, canister, "update", vec![]);
+    env.tick();
 
     // Advance the time so that the canister gets uninstalled due to the
     // resource usage.
     env.advance_time(Duration::from_secs(10_000_000));
-
     env.tick();
 
     // Enable normal message execution.
     env.set_checkpoints_enabled(false);
 
-    let mut errors = 0;
+    env.await_ingress(update.clone(), 100)
+}
 
-    // Canisters that were chosen for execution before charging for resources
-    // become paused and don't get uninstalled until their execution completes.
-    // All other canister are uninstalled before resuming their aborted
-    // executions.
-    for i in 0..n {
-        match env.await_ingress(updates[i].clone(), 100) {
-            Ok(result) => {
-                assert_eq!(result, WasmResult::Reply(vec![]));
-            }
-            Err(err) => {
-                err.assert_contains(
-                    ErrorCode::CanisterWasmModuleNotFound,
-                    &format!(
-                        "Error from Canister {}: Attempted to execute a message, \
-                        but the canister contains no Wasm module.",
-                        canisters[i]
-                    ),
-                );
-                errors += 1;
-            }
-        }
-    }
-    assert_ge!(errors, 1);
+#[test]
+fn dts_canister_with_aborted_update_is_uninstalled_due_to_resource_charges() {
+    let result = impl_dts_canister_with_update_is_uninstalled_due_to_resource_charges(true);
+
+    // Canister with aborted update is uninstalled. Update fails.
+    assert_matches!(
+        result,
+        Err(err) if err.code() == ErrorCode::CanisterWasmModuleNotFound
+    );
+}
+
+#[test]
+fn dts_canister_with_paused_update_is_not_uninstalled_due_to_resource_charges() {
+    let result = impl_dts_canister_with_update_is_uninstalled_due_to_resource_charges(false);
+
+    // Canister is only uninstalled after paused update completes successfully.
+    assert_matches!(result, Ok(_));
 }
 
 #[test]