From a7cbf1843f58507b558de54bd5fc1982d0d6388b Mon Sep 17 00:00:00 2001 From: Alin Sinpalean Date: Sun, 24 May 2026 07:48:04 +0000 Subject: [PATCH 1/2] feat: [DSM-103] Charge for storage after aborting DTS executions When charging canisters for storage, we don't charge canisters with paused executions because they should have a constant cycle balance throughuot the DTS execution. But we do charge canisters with aborted executions. So swap things around, so that we charge for storage after aborting paused executions (whether executions above the limit on regular rounds; or all paused executions on checkpoint rounds). This ensures that we can charge all canisters on checkpoint rounds (and makes it more likely to charge the average busy canister on regular rounds), making it possible to evenrually mote away from trying to charge every canister every round just to make sure we succeed. --- rs/execution_environment/src/scheduler.rs | 18 +-- rs/execution_environment/tests/dts.rs | 128 ++++++++++------------ 2 files changed, 70 insertions(+), 76 deletions(-) diff --git a/rs/execution_environment/src/scheduler.rs b/rs/execution_environment/src/scheduler.rs index ba3cacc71ab3..5901de86c5e1 100644 --- a/rs/execution_environment/src/scheduler.rs +++ b/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 diff --git a/rs/execution_environment/tests/dts.rs b/rs/execution_environment/tests/dts.rs index 7dbd7e4a3250..15208f2e7e29 100644 --- a/rs/execution_environment/tests/dts.rs +++ b/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 { 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 { 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] From 22a93ae23387f6f40c88d1ced337061f16fbeaf6 Mon Sep 17 00:00:00 2001 From: Alin Sinpalean Date: Mon, 1 Jun 2026 11:25:35 +0000 Subject: [PATCH 2/2] Apply code review suggestions. --- rs/execution_environment/tests/dts.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rs/execution_environment/tests/dts.rs b/rs/execution_environment/tests/dts.rs index 15208f2e7e29..b68788ab3857 100644 --- a/rs/execution_environment/tests/dts.rs +++ b/rs/execution_environment/tests/dts.rs @@ -1889,7 +1889,7 @@ fn impl_dts_canister_with_update_is_uninstalled_due_to_resource_charges( // Enable normal message execution. env.set_checkpoints_enabled(false); - env.await_ingress(update.clone(), 100) + env.await_ingress(update, 100) } #[test] @@ -1908,7 +1908,7 @@ 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(_)); + assert_matches!(result, Ok(WasmResult::Reply(_))); } #[test]