From eeafb6694fc845e2a21121388b1a2272759eb619 Mon Sep 17 00:00:00 2001 From: Dimitris Sarlis Date: Mon, 30 Mar 2026 09:59:35 +0000 Subject: [PATCH 1/3] refactor: Separate adding/removing cycles from refunding/consuming cycles --- .../src/cycles_account_manager.rs | 4 +- .../tests/cycles_account_manager.rs | 29 +---- rs/embedders/fuzz/src/wasm_executor.rs | 2 +- .../system_api/sandbox_safe_system_state.rs | 35 ++--- .../tests/sandbox_safe_system_state.rs | 9 +- .../src/canister_manager.rs | 12 +- .../src/canister_manager/tests.rs | 2 - .../src/execution/call_or_task/tests.rs | 7 +- .../src/execution/response.rs | 27 ++-- .../src/execution_environment.rs | 16 +-- .../src/query_handler/query_cache/tests.rs | 12 +- rs/execution_environment/src/scheduler.rs | 2 - .../src/scheduler/tests/charging.rs | 11 +- .../src/scheduler/tests/metrics.rs | 14 +- .../src/routing/stream_builder/tests.rs | 3 +- rs/messaging/src/routing/stream_handler.rs | 2 +- .../src/routing/stream_handler/tests.rs | 20 +-- rs/replicated_state/src/canister_state.rs | 3 - .../src/canister_state/system_state.rs | 123 +++++++++++++----- .../src/canister_state/tests.rs | 72 +++------- rs/replicated_state/src/replicated_state.rs | 16 +-- rs/replicated_state/tests/canister_state.rs | 2 - rs/replicated_state/tests/replicated_state.rs | 14 +- rs/state_machine_tests/src/lib.rs | 11 +- .../execution_environment/src/lib.rs | 3 - 25 files changed, 191 insertions(+), 260 deletions(-) diff --git a/rs/cycles_account_manager/src/cycles_account_manager.rs b/rs/cycles_account_manager/src/cycles_account_manager.rs index c213d59d95cf..e66461e6cd77 100644 --- a/rs/cycles_account_manager/src/cycles_account_manager.rs +++ b/rs/cycles_account_manager/src/cycles_account_manager.rs @@ -563,7 +563,7 @@ impl CyclesAccountManager { cost_schedule, ) .min(prepaid_execution_cycles); - system_state.add_cycles(cycles_to_refund); + system_state.refund_cycles(prepaid_execution_cycles, cycles_to_refund); } /// Returns the cost of compute allocation for the given duration. @@ -1018,7 +1018,7 @@ impl CyclesAccountManager { )?; debug_assert_ne!(use_case, CyclesUseCase::NonConsumed); - system_state.remove_cycles(cycles); + system_state.consume_cycles(cycles); Ok(()) } diff --git a/rs/cycles_account_manager/tests/cycles_account_manager.rs b/rs/cycles_account_manager/tests/cycles_account_manager.rs index 574ce01b234d..4d711b481e35 100644 --- a/rs/cycles_account_manager/tests/cycles_account_manager.rs +++ b/rs/cycles_account_manager/tests/cycles_account_manager.rs @@ -25,7 +25,7 @@ use ic_types::{ }; use ic_types_cycles::{ CanisterCyclesCostSchedule, CompoundCycles, Cycles, Instructions, Memory, NominalCycles, - NominalCyclesTesting, NonConsumed, Uninstall, + NominalCyclesTesting, Uninstall, }; use prometheus::IntCounter; use std::{convert::TryFrom, time::Duration}; @@ -97,12 +97,7 @@ fn test_can_charge_application_subnets() { .memory_cost(memory, duration, subnet_size, cost_schedule) .real(); let initial_cycles = expected_fee; - canister - .system_state - .add_cycles(CompoundCycles::::new( - initial_cycles, - cost_schedule, - )); + canister.system_state.add_cycles(initial_cycles); assert_eq!(canister.system_state.balance(), initial_cycles); cycles_account_manager .charge_canister_for_resource_allocation_and_usage( @@ -1145,10 +1140,7 @@ fn consume_cycles_for_memory_drains_reserved_balance() { let mut system_state = SystemStateBuilder::new() .initial_cycles(Cycles::zero()) .build(); - system_state.add_cycles(CompoundCycles::::new( - Cycles::new(4_000_000), - cost_schedule, - )); + system_state.add_cycles(Cycles::new(4_000_000)); system_state.reserve_cycles(Cycles::new(1_000_000)).unwrap(); cam.consume_with_threshold( &mut system_state, @@ -1170,10 +1162,7 @@ fn consume_cycles_for_compute_drains_reserved_balance() { let mut system_state = SystemStateBuilder::new() .initial_cycles(Cycles::zero()) .build(); - system_state.add_cycles(CompoundCycles::::new( - Cycles::new(4_000_000), - cost_schedule, - )); + system_state.add_cycles(Cycles::new(4_000_000)); system_state.reserve_cycles(Cycles::new(1_000_000)).unwrap(); cam.consume_with_threshold( &mut system_state, @@ -1198,10 +1187,7 @@ fn consume_cycles_for_uninstall_drains_reserved_balance() { let mut system_state = SystemStateBuilder::new() .initial_cycles(Cycles::zero()) .build(); - system_state.add_cycles(CompoundCycles::::new( - Cycles::new(4_000_000), - cost_schedule, - )); + system_state.add_cycles(Cycles::new(4_000_000)); system_state.reserve_cycles(Cycles::new(1_000_000)).unwrap(); cam.consume_with_threshold( &mut system_state, @@ -1223,10 +1209,7 @@ fn consume_cycles_for_execution_does_not_drain_reserved_balance() { let mut system_state = SystemStateBuilder::new() .initial_cycles(Cycles::zero()) .build(); - system_state.add_cycles(CompoundCycles::::new( - Cycles::new(4_000_000), - cost_schedule, - )); + system_state.add_cycles(Cycles::new(4_000_000)); system_state.reserve_cycles(Cycles::new(1_000_000)).unwrap(); cam.consume_with_threshold( &mut system_state, diff --git a/rs/embedders/fuzz/src/wasm_executor.rs b/rs/embedders/fuzz/src/wasm_executor.rs index aeda48c429a5..3129fa169a0d 100644 --- a/rs/embedders/fuzz/src/wasm_executor.rs +++ b/rs/embedders/fuzz/src/wasm_executor.rs @@ -110,7 +110,7 @@ pub fn run_fuzzer(module: ICWasmModule) { } canister_state_changes .system_state_modifications - .apply_balance_changes(&mut system_state, CanisterCyclesCostSchedule::Normal); + .apply_balance_changes(&mut system_state); } WasmExecutionResult::Paused(_, _) => (), // Only possible via execute_dts } diff --git a/rs/embedders/src/wasmtime_embedder/system_api/sandbox_safe_system_state.rs b/rs/embedders/src/wasmtime_embedder/system_api/sandbox_safe_system_state.rs index f4bde051c1cd..b03524a89421 100644 --- a/rs/embedders/src/wasmtime_embedder/system_api/sandbox_safe_system_state.rs +++ b/rs/embedders/src/wasmtime_embedder/system_api/sandbox_safe_system_state.rs @@ -34,7 +34,7 @@ use ic_types::{ }; use ic_types_cycles::{ BurnedCycles, CanisterCyclesCostSchedule, CompoundCycles, Cycles, CyclesUseCaseKind, - Instructions, NonConsumed, RequestAndResponseTransmission, + Instructions, RequestAndResponseTransmission, }; use ic_wasm_types::WasmEngineError; use serde::{Deserialize, Serialize}; @@ -412,11 +412,8 @@ impl SystemStateModifications { .append_delta_log(&mut self.canister_log); // Verify total cycle change is not positive and update cycles balance. - let cost_schedule = network_topology - .get_cost_schedule(&own_subnet_id) - .unwrap_or_default(); self.validate_cycle_change(system_state.canister_id() == CYCLES_MINTING_CANISTER_ID)?; - self.apply_balance_changes(system_state, cost_schedule); + self.apply_balance_changes(system_state); if let Some(hook_condition_check_result) = self.on_low_wasm_memory_hook_condition_check_result @@ -605,11 +602,7 @@ impl SystemStateModifications { } /// Applies the balance change to the given state. - pub fn apply_balance_changes( - &self, - state: &mut SystemState, - cost_schedule: CanisterCyclesCostSchedule, - ) { + pub fn apply_balance_changes(&self, state: &mut SystemState) { let initial_balance = state.balance(); // `self.cycles_balance_change` consists of: @@ -631,12 +624,8 @@ impl SystemStateModifications { // Apply the main cycles balance change without the consumed and reserved cycles. match adjusted_balance_change { - CyclesBalanceChange::Added(added) => { - state.add_cycles(CompoundCycles::::new(added, cost_schedule)) - } - CyclesBalanceChange::Removed(removed) => { - state.remove_cycles(CompoundCycles::::new(removed, cost_schedule)) - } + CyclesBalanceChange::Added(added) => state.add_cycles(added), + CyclesBalanceChange::Removed(removed) => state.remove_cycles(removed), } // Apply the consumed cycles with the use case metrics recording. @@ -646,13 +635,13 @@ impl SystemStateModifications { request_and_response_transmission, } = self.consumed_cycles_by_use_case; if let Some(x) = burned { - state.remove_cycles(x); + state.consume_cycles(x); } if let Some(x) = instructions { - state.remove_cycles(x); + state.consume_cycles(x); } if let Some(x) = request_and_response_transmission { - state.remove_cycles(x); + state.consume_cycles(x); } // Apply the reserved cycles. This must succeed because the cycle @@ -1618,7 +1607,7 @@ mod tests { }, ); - system_state_modifications.apply_balance_changes(&mut system_state, cost_schedule); + system_state_modifications.apply_balance_changes(&mut system_state); assert_eq!(initial_cycles_balance - removed, system_state.balance()); @@ -1637,7 +1626,7 @@ mod tests { }, ); - system_state_modifications.apply_balance_changes(&mut system_state, cost_schedule); + system_state_modifications.apply_balance_changes(&mut system_state); assert_eq!(initial_cycles_balance - removed, system_state.balance()); @@ -1656,7 +1645,7 @@ mod tests { }, ); - system_state_modifications.apply_balance_changes(&mut system_state, cost_schedule); + system_state_modifications.apply_balance_changes(&mut system_state); assert_eq!(initial_cycles_balance + added, system_state.balance()); @@ -1675,7 +1664,7 @@ mod tests { }, ); - system_state_modifications.apply_balance_changes(&mut system_state, cost_schedule); + system_state_modifications.apply_balance_changes(&mut system_state); assert_eq!(initial_cycles_balance + added, system_state.balance()); } diff --git a/rs/embedders/tests/sandbox_safe_system_state.rs b/rs/embedders/tests/sandbox_safe_system_state.rs index 3a1a2ba1cb27..2790db68fce4 100644 --- a/rs/embedders/tests/sandbox_safe_system_state.rs +++ b/rs/embedders/tests/sandbox_safe_system_state.rs @@ -26,7 +26,7 @@ use ic_types::methods::{Callback, WasmClosure}; use ic_types::time::UNIX_EPOCH; use ic_types::{ComputeAllocation, NumInstructions}; use ic_types_cycles::{ - CanisterCyclesCostSchedule, CompoundCycles, Cycles, CyclesUseCase, NominalCycles, NonConsumed, + CanisterCyclesCostSchedule, CompoundCycles, Cycles, CyclesUseCase, NominalCycles, }; use prometheus::IntCounter; use std::collections::BTreeSet; @@ -317,7 +317,7 @@ fn correct_charging_source_canister_for_a_request() { cost_schedule, ); - system_state.add_cycles(refund_cycles); + system_state.refund_cycles(prepayment_for_response_transmission, refund_cycles); // MAX_NUM_INSTRUCTIONS also gets partially refunded in the real // ExecutionEnvironmentImpl::execute_canister_response() @@ -395,10 +395,7 @@ fn mint_cycles_large_value() { .canister_id(CYCLES_MINTING_CANISTER_ID) .build(); - system_state.add_cycles(CompoundCycles::::new( - Cycles::from(1_000_000_000_000_000_u128), - CanisterCyclesCostSchedule::Normal, - )); + system_state.add_cycles(Cycles::from(1_000_000_000_000_000_u128)); let api_type = ApiTypeBuilder::build_update_api(); let mut api = get_system_api(api_type, &system_state, cycles_account_manager); diff --git a/rs/execution_environment/src/canister_manager.rs b/rs/execution_environment/src/canister_manager.rs index 5772cadb6d1f..30083eaf8ab3 100644 --- a/rs/execution_environment/src/canister_manager.rs +++ b/rs/execution_environment/src/canister_manager.rs @@ -69,7 +69,7 @@ use ic_types::{ }; use ic_types_cycles::{ CanisterCreation, CanisterCyclesCostSchedule, CompoundCycles, Cycles, CyclesUseCase, - Instructions, NonConsumed, + Instructions, }; use ic_wasm_types::WasmHash; use more_asserts::{debug_assert_ge, debug_assert_le}; @@ -1489,7 +1489,7 @@ impl CanisterManager { Arc::clone(&self.fd_factory), ); - system_state.remove_cycles(creation_fee); + system_state.consume_cycles(creation_fee); let mut new_canister = CanisterState::new( system_state, None, @@ -1567,7 +1567,6 @@ impl CanisterManager { cycles_amount: Option, canister: &mut CanisterState, provisional_whitelist: &ProvisionalWhitelist, - cost_schedule: CanisterCyclesCostSchedule, ) -> Result<(), CanisterManagerError> { if !provisional_whitelist.contains(&sender) { return Err(CanisterManagerError::SenderNotInWhitelist(sender)); @@ -1578,12 +1577,7 @@ impl CanisterManager { None => self.config.default_provisional_cycles_balance, }; - canister - .system_state - .add_cycles(CompoundCycles::::new( - cycles_amount, - cost_schedule, - )); + canister.system_state.add_cycles(cycles_amount); Ok(()) } diff --git a/rs/execution_environment/src/canister_manager/tests.rs b/rs/execution_environment/src/canister_manager/tests.rs index 0a6b599dded5..25931cb15805 100644 --- a/rs/execution_environment/src/canister_manager/tests.rs +++ b/rs/execution_environment/src/canister_manager/tests.rs @@ -2226,7 +2226,6 @@ fn add_cycles_sender_in_whitelist() { Some(123), canister, &ProvisionalWhitelist::Set(btreeset! { canister_test_id(1).get() }), - CanisterCyclesCostSchedule::Normal, ) .unwrap(); @@ -2256,7 +2255,6 @@ fn add_cycles_sender_not_in_whitelist() { Some(123), canister, &ProvisionalWhitelist::Set(BTreeSet::new()), - CanisterCyclesCostSchedule::Normal, ), Err(CanisterManagerError::SenderNotInWhitelist(sender)) ); diff --git a/rs/execution_environment/src/execution/call_or_task/tests.rs b/rs/execution_environment/src/execution/call_or_task/tests.rs index 12e77885f83b..90fa9a517077 100644 --- a/rs/execution_environment/src/execution/call_or_task/tests.rs +++ b/rs/execution_environment/src/execution/call_or_task/tests.rs @@ -13,7 +13,7 @@ use ic_sys::PAGE_SIZE; use ic_types::ingress::IngressState; use ic_types::messages::{CallbackId, RequestMetadata}; use ic_types::{NumBytes, NumInstructions, NumOsPages}; -use ic_types_cycles::{CanisterCyclesCostSchedule, CompoundCycles, Cycles, NonConsumed}; +use ic_types_cycles::{CanisterCyclesCostSchedule, Cycles}; use ic_universal_canister::{call_args, wasm}; use more_asserts::assert_gt; use std::time::Duration; @@ -760,10 +760,7 @@ fn dts_replicated_execution_resume_fails_due_to_cycles_change() { let balance = test.canister_state(a_id).system_state.balance(); test.canister_state_mut(a_id) .system_state - .add_cycles(CompoundCycles::::new( - balance + Cycles::new(1), - CanisterCyclesCostSchedule::Normal, - )); + .add_cycles(balance + Cycles::new(1)); test.execute_slice(a_id); diff --git a/rs/execution_environment/src/execution/response.rs b/rs/execution_environment/src/execution/response.rs index c64768616dfc..a69a9843f967 100644 --- a/rs/execution_environment/src/execution/response.rs +++ b/rs/execution_environment/src/execution/response.rs @@ -27,7 +27,7 @@ use ic_types::messages::{ use ic_types::methods::{Callback, FuncRef, WasmClosure}; use ic_types::{NumBytes, NumInstructions, Time}; use ic_types_cycles::{ - CanisterCyclesCostSchedule, CompoundCycles, Cycles, Instructions, NonConsumed, + CanisterCyclesCostSchedule, CompoundCycles, Cycles, Instructions, RequestAndResponseTransmission, }; use ic_utils_thread::deallocator_thread::DeallocationSender; @@ -111,6 +111,7 @@ const RESERVED_CLEANUP_INSTRUCTIONS_IN_PERCENT: u64 = 5; #[derive(Debug)] struct PausedResponseHelper { refund_for_sent_cycles: Cycles, + prepayment_for_response_transmission: CompoundCycles, refund_for_response_transmission: CompoundCycles, initial_cycles_balance: Cycles, response_sender: CanisterId, @@ -121,6 +122,7 @@ struct PausedResponseHelper { struct ResponseHelper { canister: CanisterState, refund_for_sent_cycles: Cycles, + prepayment_for_response_transmission: CompoundCycles, refund_for_response_transmission: CompoundCycles, initial_cycles_balance: Cycles, response_sender: CanisterId, @@ -185,6 +187,9 @@ impl ResponseHelper { let mut helper = Self { canister, refund_for_sent_cycles, + prepayment_for_response_transmission: original + .callback + .prepayment_for_response_transmission, refund_for_response_transmission, initial_cycles_balance, response_sender, @@ -200,17 +205,15 @@ impl ResponseHelper { /// /// These are the only state changes to the initial canister state before /// executing Wasm code. - fn apply_initial_refunds(&mut self, cost_schedule: CanisterCyclesCostSchedule) { + fn apply_initial_refunds(&mut self) { self.canister .system_state - .add_cycles(CompoundCycles::::new( - self.refund_for_sent_cycles, - cost_schedule, - )); + .add_cycles(self.refund_for_sent_cycles); - self.canister - .system_state - .add_cycles(self.refund_for_response_transmission); + self.canister.system_state.refund_cycles( + self.prepayment_for_response_transmission, + self.refund_for_response_transmission, + ); } /// Checks that the canister has not been uninstalled: @@ -284,6 +287,7 @@ impl ResponseHelper { self.deallocation_sender.send(Box::new(self.canister)); PausedResponseHelper { refund_for_sent_cycles: self.refund_for_sent_cycles, + prepayment_for_response_transmission: self.prepayment_for_response_transmission, refund_for_response_transmission: self.refund_for_response_transmission, initial_cycles_balance: self.initial_cycles_balance, response_sender: self.response_sender, @@ -323,6 +327,7 @@ impl ResponseHelper { let mut helper = Self { canister: clean_canister.clone(), refund_for_sent_cycles: paused.refund_for_sent_cycles, + prepayment_for_response_transmission: paused.prepayment_for_response_transmission, refund_for_response_transmission: paused.refund_for_response_transmission, initial_cycles_balance: clean_canister.system_state.balance(), response_sender: paused.response_sender, @@ -332,7 +337,7 @@ impl ResponseHelper { helper.apply_subnet_memory_reservation(round_limits); - helper.apply_initial_refunds(round.cost_schedule); + helper.apply_initial_refunds(); // This validation succeeded in `execute_response()` and we expect it to // succeed here too. @@ -980,7 +985,7 @@ pub fn execute_response( round_limits, deallocation_sender, ); - helper.apply_initial_refunds(round.cost_schedule); + helper.apply_initial_refunds(); let helper = match helper.validate(&call_context, &original, &round, round_limits) { Ok(helper) => helper, Err(result) => { diff --git a/rs/execution_environment/src/execution_environment.rs b/rs/execution_environment/src/execution_environment.rs index 5057c9c712b8..5a7e7af8925a 100644 --- a/rs/execution_environment/src/execution_environment.rs +++ b/rs/execution_environment/src/execution_environment.rs @@ -85,7 +85,7 @@ use ic_types::{ use ic_types::{messages::MessageId, methods::WasmMethod}; use ic_types_cycles::{ CanisterCyclesCostSchedule, CompoundCycles, Cycles, CyclesUseCase, ECDSAOutcalls, Instructions, - NominalCycles, NonConsumed, SchnorrOutcalls, VetKd, + NominalCycles, SchnorrOutcalls, VetKd, }; use ic_utils_thread::deallocator_thread::{DeallocationSender, DeallocatorThread}; use ic_wasm_types::WasmHash; @@ -2393,7 +2393,6 @@ impl ExecutionEnvironment { msg: &mut CanisterCall, state: &mut ReplicatedState, ) -> ExecuteSubnetMessageResult { - let cost_schedule = state.get_own_cost_schedule(); match state.canister_state_make_mut(&canister_id) { None => ExecuteSubnetMessageResult::Finished { response: Err(UserError::new( @@ -2405,9 +2404,7 @@ impl ExecutionEnvironment { Some(canister_state) => { let cycles = msg.take_cycles(); - canister_state - .system_state - .add_cycles(CompoundCycles::::new(cycles, cost_schedule)); + canister_state.system_state.add_cycles(cycles); if cycles.get() > LOG_CANISTER_OPERATION_CYCLES_THRESHOLD { info!( self.log, @@ -2527,16 +2524,9 @@ impl ExecutionEnvironment { state: &mut ReplicatedState, provisional_whitelist: &ProvisionalWhitelist, ) -> Result, UserError> { - let cost_schedule = state.get_own_cost_schedule(); let canister = canister_make_mut(canister_id, state)?; self.canister_manager - .add_cycles( - sender, - cycles, - canister, - provisional_whitelist, - cost_schedule, - ) + .add_cycles(sender, cycles, canister, provisional_whitelist) .map(|()| EmptyBlob.encode()) .map_err(|err| err.into()) } diff --git a/rs/execution_environment/src/query_handler/query_cache/tests.rs b/rs/execution_environment/src/query_handler/query_cache/tests.rs index ab1fa702e1ea..e2229daac93b 100644 --- a/rs/execution_environment/src/query_handler/query_cache/tests.rs +++ b/rs/execution_environment/src/query_handler/query_cache/tests.rs @@ -588,7 +588,7 @@ fn query_cache_ignores_balance_changes_when_query_does_not_read_balance() { // Change the canister balance. test.canister_state_mut(b_id) .system_state - .remove_cycles(CompoundCycles::::new( + .consume_cycles(CompoundCycles::::new( 1_u64.into(), CanisterCyclesCostSchedule::Normal, )); @@ -621,7 +621,7 @@ fn query_cache_ignores_balance_and_time_changes_when_query_is_static() { // Change the canister balance. test.canister_state_mut(b_id) .system_state - .remove_cycles(CompoundCycles::::new( + .consume_cycles(CompoundCycles::::new( 1_u64.into(), CanisterCyclesCostSchedule::Normal, )); @@ -777,7 +777,7 @@ fn query_cache_returns_different_results_for_different_canister_balances() { // Change the canister balance. test.canister_state_mut(b_id) .system_state - .remove_cycles(CompoundCycles::::new( + .consume_cycles(CompoundCycles::::new( 1_u64.into(), CanisterCyclesCostSchedule::Normal, )); @@ -807,7 +807,7 @@ fn query_cache_returns_different_results_for_different_canister_balance128s() { // Change the canister balance. test.canister_state_mut(b_id) .system_state - .remove_cycles(CompoundCycles::::new( + .consume_cycles(CompoundCycles::::new( 1_u64.into(), CanisterCyclesCostSchedule::Normal, )); @@ -848,7 +848,7 @@ fn query_cache_returns_different_results_on_combined_invalidation() { .bump_canister_version(); test.canister_state_mut(b_id) .system_state - .remove_cycles(CompoundCycles::::new( + .consume_cycles(CompoundCycles::::new( 1_u64.into(), CanisterCyclesCostSchedule::Normal, )); @@ -899,7 +899,7 @@ fn query_cache_frees_memory_after_invalidated_entries() { // Set the canister balance to 42, so the second reply will have just 42 bytes. test.canister_state_mut(id) .system_state - .remove_cycles(CompoundCycles::::new( + .consume_cycles(CompoundCycles::::new( ((BIG_RESPONSE_SIZE - SMALL_RESPONSE_SIZE) as u64).into(), CanisterCyclesCostSchedule::Normal, )); diff --git a/rs/execution_environment/src/scheduler.rs b/rs/execution_environment/src/scheduler.rs index 70e19c91341c..3507c877604e 100644 --- a/rs/execution_environment/src/scheduler.rs +++ b/rs/execution_environment/src/scheduler.rs @@ -956,14 +956,12 @@ impl SchedulerImpl { .system_state .output_queues_for_each(|canister_id, msg| { let own_subnet_type = state.metadata.own_subnet_type; - let own_cost_schedule = state.get_own_cost_schedule(); match state.canister_state_make_mut(canister_id) { Some(dest_canister) => dest_canister .push_input( (*msg).clone(), &mut subnet_available_guaranteed_response_memory, own_subnet_type, - own_cost_schedule, InputQueueType::LocalSubnet, ) .map(|_| ()) diff --git a/rs/execution_environment/src/scheduler/tests/charging.rs b/rs/execution_environment/src/scheduler/tests/charging.rs index 2fea36b09af1..83502ce24c80 100644 --- a/rs/execution_environment/src/scheduler/tests/charging.rs +++ b/rs/execution_environment/src/scheduler/tests/charging.rs @@ -14,7 +14,6 @@ use ic_replicated_state::canister_state::system_state::PausedExecutionId; use ic_replicated_state::testing::SystemStateTesting; use ic_types::messages::{CanisterMessageOrTask, CanisterTask}; use ic_types::time::UNIX_EPOCH; -use ic_types_cycles::{CanisterCyclesCostSchedule, CompoundCycles, NonConsumed}; use ic_types_test_utils::ids::canister_test_id; use std::time::Duration; @@ -403,10 +402,7 @@ fn snapshot_is_deleted_when_canister_is_out_of_cycles() { .canister_state_make_mut(&canister_id) .unwrap() .system_state - .add_cycles(CompoundCycles::::new( - expected_charge, - CanisterCyclesCostSchedule::Normal, - )); + .add_cycles(expected_charge); // Take a snapshot of the canister. let args: TakeCanisterSnapshotArgs = @@ -524,10 +520,7 @@ fn snapshot_is_deleted_when_uninstalled_canister_is_out_of_cycles() { .canister_state_make_mut(&canister_id) .unwrap() .system_state - .add_cycles(CompoundCycles::::new( - expected_charge, - CanisterCyclesCostSchedule::Normal, - )); + .add_cycles(expected_charge); // Take a snapshot of the canister. let args: TakeCanisterSnapshotArgs = diff --git a/rs/execution_environment/src/scheduler/tests/metrics.rs b/rs/execution_environment/src/scheduler/tests/metrics.rs index 09c03b912853..778693451383 100644 --- a/rs/execution_environment/src/scheduler/tests/metrics.rs +++ b/rs/execution_environment/src/scheduler/tests/metrics.rs @@ -1099,13 +1099,13 @@ fn consumed_cycles_are_updated_from_valid_canisters() { None, ); - let removed_cycles = CompoundCycles::::new( + let consumed_cycles = CompoundCycles::::new( Cycles::from(1000_u128), CanisterCyclesCostSchedule::Normal, ); test.canister_state_mut(canister_id) .system_state - .remove_cycles(removed_cycles); + .consume_cycles(consumed_cycles); test.scheduler().state_metrics.observe( test.scheduler().own_subnet_id, @@ -1121,7 +1121,7 @@ fn consumed_cycles_are_updated_from_valid_canisters() { ), metric_vec(&[( &[("use_case", "Instructions")], - removed_cycles.nominal().get() as f64 + consumed_cycles.nominal().get() as f64 ),]), ); } @@ -1139,13 +1139,13 @@ fn consumed_cycles_are_updated_from_deleted_canisters() { Some(CanisterStatusType::Stopped), ); - let removed_cycles = CompoundCycles::::new( + let consumed_cycles = CompoundCycles::::new( Cycles::from(1000_u128), CanisterCyclesCostSchedule::Normal, ); test.canister_state_mut(canister_id) .system_state - .remove_cycles(removed_cycles); + .consume_cycles(consumed_cycles); test.inject_call_to_ic00( Method::DeleteCanister, @@ -1171,11 +1171,11 @@ fn consumed_cycles_are_updated_from_deleted_canisters() { metric_vec(&[ ( &[("use_case", "Instructions")], - removed_cycles.nominal().get() as f64 + consumed_cycles.nominal().get() as f64 ), ( &[("use_case", "DeletedCanisters")], - (initial_balance.get() - removed_cycles.nominal().get()) as f64 + (initial_balance.get() - consumed_cycles.nominal().get()) as f64 ) ]), ); diff --git a/rs/messaging/src/routing/stream_builder/tests.rs b/rs/messaging/src/routing/stream_builder/tests.rs index 85e494ae03d5..25d4922d4678 100644 --- a/rs/messaging/src/routing/stream_builder/tests.rs +++ b/rs/messaging/src/routing/stream_builder/tests.rs @@ -30,7 +30,7 @@ use ic_types::messages::{ use ic_types::time::{CoarseTime, UNIX_EPOCH}; use ic_types::xnet::{StreamIndex, StreamIndexedQueue}; use ic_types::{CanisterId, SubnetId, Time}; -use ic_types_cycles::{CanisterCyclesCostSchedule, Cycles}; +use ic_types_cycles::Cycles; use lazy_static::lazy_static; use maplit::btreemap; use pretty_assertions::assert_eq; @@ -1669,7 +1669,6 @@ fn push_input(canister_state: &mut CanisterState, msg: RequestOrResponse) { msg, &mut subnet_available_memory, SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ) .unwrap() diff --git a/rs/messaging/src/routing/stream_handler.rs b/rs/messaging/src/routing/stream_handler.rs index 370f8daaaed4..6a5e5193d782 100644 --- a/rs/messaging/src/routing/stream_handler.rs +++ b/rs/messaging/src/routing/stream_handler.rs @@ -1002,7 +1002,7 @@ impl StreamHandlerImpl { Some(host_subnet) if host_subnet == self.subnet_id => { let cost_schedule = state.get_own_cost_schedule(); stream.push_accept_signal(); - if state.credit_refund(refund, cost_schedule) { + if state.credit_refund(refund) { self.observe_inducted_message_status( LABEL_VALUE_TYPE_REFUND, LABEL_VALUE_SUCCESS, diff --git a/rs/messaging/src/routing/stream_handler/tests.rs b/rs/messaging/src/routing/stream_handler/tests.rs index dfe7604bdb3e..caf0253bbe3f 100644 --- a/rs/messaging/src/routing/stream_handler/tests.rs +++ b/rs/messaging/src/routing/stream_handler/tests.rs @@ -1966,10 +1966,10 @@ fn duplicate_best_effort_response_is_dropped() { }); expected_state.with_streams(btreemap![LOCAL_SUBNET => loopback_stream]); // Cycles of the duplicate response are refunded. - expected_state.credit_refund( - &ic_types::messages::Refund::anonymous(*LOCAL_CANISTER, response.cycles()), - CanisterCyclesCostSchedule::Normal, - ); + expected_state.credit_refund(&ic_types::messages::Refund::anonymous( + *LOCAL_CANISTER, + response.cycles(), + )); // Push the clone of the best effort response onto the loopback stream. state.modify_streams(|streams| streams.get_mut(&LOCAL_SUBNET).unwrap().push(response)); @@ -2049,10 +2049,10 @@ fn inducting_best_effort_response_into_stopped_canister_does_not_raise_a_critica }, |expected_state, refund| { // Cycles attached to the late response are refunded. - expected_state.credit_refund( - &ic_types::messages::Refund::anonymous(*LOCAL_CANISTER, refund), - CanisterCyclesCostSchedule::Normal, - ); + expected_state.credit_refund(&ic_types::messages::Refund::anonymous( + *LOCAL_CANISTER, + refund, + )); }, ); } @@ -2744,7 +2744,7 @@ fn induct_stream_slices_with_refunds() { // Refund to `LOCAL_CANISTER` (@43) was applied. let refund43 = refund_in_slice(slices.get(&REMOTE_SUBNET), 43); - expected_state.credit_refund(refund43, CanisterCyclesCostSchedule::Normal); + expected_state.credit_refund(refund43); // Cycles in refund @44 are lost let refund44 = refund_in_slice(slices.get(&REMOTE_SUBNET), 44); @@ -3650,7 +3650,7 @@ fn push_input(state: &mut ReplicatedState, msg: StreamMessage) { .unwrap(); } StreamMessage::Refund(refund) => { - assert!(state.credit_refund(&refund, CanisterCyclesCostSchedule::Normal)); + assert!(state.credit_refund(&refund)); } } } diff --git a/rs/replicated_state/src/canister_state.rs b/rs/replicated_state/src/canister_state.rs index 21c7b91af1a8..bcb624133e34 100644 --- a/rs/replicated_state/src/canister_state.rs +++ b/rs/replicated_state/src/canister_state.rs @@ -28,7 +28,6 @@ use ic_types::{ CanisterId, CanisterLog, ComputeAllocation, MemoryAllocation, NumBytes, NumInstructions, PrincipalId, Time, }; -use ic_types_cycles::CanisterCyclesCostSchedule; use ic_validate_eq::ValidateEq; use ic_validate_eq_derive::ValidateEq; use phantom_newtype::AmountOf; @@ -143,14 +142,12 @@ impl CanisterState { msg: RequestOrResponse, subnet_available_guaranteed_response_memory: &mut i64, own_subnet_type: SubnetType, - own_cost_schedule: CanisterCyclesCostSchedule, input_queue_type: InputQueueType, ) -> Result { self.system_state.push_input( msg, subnet_available_guaranteed_response_memory, own_subnet_type, - own_cost_schedule, input_queue_type, ) } diff --git a/rs/replicated_state/src/canister_state/system_state.rs b/rs/replicated_state/src/canister_state/system_state.rs index 13d11c2dd822..491271431da8 100644 --- a/rs/replicated_state/src/canister_state/system_state.rs +++ b/rs/replicated_state/src/canister_state/system_state.rs @@ -46,7 +46,7 @@ use ic_types::{ }; use ic_types_cycles::{ CanisterCyclesCostSchedule, CompoundCycles, Cycles, CyclesUseCase, CyclesUseCaseKind, - IngressInduction, Instructions, NominalCycles, NonConsumed, Uninstall, + IngressInduction, Instructions, NominalCycles, Uninstall, }; use ic_validate_eq::ValidateEq; use ic_validate_eq_derive::ValidateEq; @@ -68,9 +68,10 @@ lazy_static! { /// Maximum number of canister changes stored in the canister history. pub const MAX_CANISTER_HISTORY_CHANGES: u64 = 20; +#[derive(PartialEq)] enum ConsumingCycles { - Yes, - No, + Prepayment, + Refund, } /// Keeps track of the types of messages executed by the canister. @@ -887,7 +888,7 @@ impl SystemState { // Continue the execution by dropping the remaining debit, which makes // some of the postponed charges free. } - self.remove_cycles(CompoundCycles::::new( + self.consume_cycles(CompoundCycles::::new( self.ingress_induction_cycles_debit, cost_schedule, )); @@ -1181,7 +1182,6 @@ impl SystemState { msg: RequestOrResponse, subnet_available_guaranteed_response_memory: &mut i64, own_subnet_type: SubnetType, - own_cost_schedule: CanisterCyclesCostSchedule, input_queue_type: InputQueueType, ) -> Result { #[cfg(debug_assertions)] @@ -1191,7 +1191,6 @@ impl SystemState { msg, subnet_available_guaranteed_response_memory, own_subnet_type, - own_cost_schedule, input_queue_type, ); @@ -1212,7 +1211,6 @@ impl SystemState { msg: RequestOrResponse, subnet_available_guaranteed_response_memory: &mut i64, own_subnet_type: SubnetType, - own_cost_schedule: CanisterCyclesCostSchedule, input_queue_type: InputQueueType, ) -> Result { assert_eq!( @@ -1228,7 +1226,7 @@ impl SystemState { (RequestOrResponse::Response(response), CanisterStatus::Stopped) if response.is_best_effort() => { - self.credit_refund(response, own_cost_schedule); + self.credit_refund(response); Ok(false) } @@ -1261,7 +1259,7 @@ impl SystemState { .map_err(|err| (err, msg.clone()))? { // Best effort response whose callback is gone. Silently drop it. - self.credit_refund(response, own_cost_schedule); + self.credit_refund(response); return Ok(false); } push_input( @@ -1274,7 +1272,7 @@ impl SystemState { .map(|dropped| { if let Some(response) = dropped { // Duplicate best-effort response that was silently dropped. - self.credit_refund(&response, own_cost_schedule); + self.credit_refund(&response); false } else { true @@ -1798,7 +1796,7 @@ impl SystemState { } /// Credits the canister with the refund in the inbound `Response`. - fn credit_refund(&mut self, response: &Response, cost_schedule: CanisterCyclesCostSchedule) { + fn credit_refund(&mut self, response: &Response) { debug_assert_eq!( self.canister_id, response.originator, "Can only credit refunds from `Responses` originating from self ({}), got {:?}", @@ -1807,28 +1805,62 @@ impl SystemState { debug_assert!(response.is_best_effort()); if !response.refund.is_zero() { - self.add_cycles(CompoundCycles::::new( - response.refund, - cost_schedule, - )); + self.add_cycles(response.refund); } } - /// Increments 'cycles_balance' and in case of refund for consumed cycles - /// decrements the metric `consumed_cycles`. - pub fn add_cycles(&mut self, amount: CompoundCycles) { - self.cycles_balance += amount.real(); + /// Increments the canister's cycles balance by the provided amount. + /// + /// If you need to have metrics about consumed cycles updated use + /// `refund_cycles` instead. + pub fn add_cycles(&mut self, amount: Cycles) { + self.cycles_balance += amount; + } + + /// Increments the canister's balance as well as the consumed metrics for the + /// respetive use case based on the amounts provided. + /// + /// Similar to `add_cycles` but additionally updates the metrics to match + /// the refund that was provided. Should be used after a prepayment has been + /// made. + pub fn refund_cycles( + &mut self, + prepayment: CompoundCycles, + refund: CompoundCycles, + ) { + debug_assert!( + prepayment.real() >= refund.real(), + "Expected real prepayment {prepayment:?} to be greater or equal to real refund {refund:?}" + ); + debug_assert!( + prepayment.nominal() >= refund.nominal(), + "Expected nominal prepayment {prepayment:?} to be greater or equal to nominal refund {refund:?}" + ); + + self.add_cycles(refund.real()); + + let use_case = T::cycles_use_case(); self.observe_consumed_cycles_with_use_case( - amount.nominal(), - T::cycles_use_case(), - ConsumingCycles::No, + prepayment.nominal(), + refund.nominal(), + use_case, + ConsumingCycles::Refund, ); } + /// Decreases 'cycles_balance' for 'requested_amount' from the canister's + /// main balance. + /// + /// If you need to have metrics about consumed cycles updated use + /// `consume_cycles` instead. + pub fn remove_cycles(&mut self, requested_amount: Cycles) { + self.cycles_balance -= requested_amount; + } + /// Decreases 'cycles_balance' for 'requested_amount'. /// The resource use cases first drain the `reserved_balance` and only after /// that drain the main `cycles_balance`. - pub fn remove_cycles(&mut self, requested_amount: CompoundCycles) { + pub fn consume_cycles(&mut self, requested_amount: CompoundCycles) { let requested_real = requested_amount.real(); let use_case = T::cycles_use_case(); let remaining_amount = match use_case { @@ -1853,8 +1885,9 @@ impl SystemState { self.cycles_balance -= remaining_amount; self.observe_consumed_cycles_with_use_case( requested_amount.nominal(), + NominalCycles::zero(), use_case, - ConsumingCycles::Yes, + ConsumingCycles::Prepayment, ); } @@ -1903,12 +1936,13 @@ impl SystemState { cost_schedule: CanisterCyclesCostSchedule, ) { let balance = self.cycles_balance + self.reserved_balance; - self.remove_cycles(CompoundCycles::::new(balance, cost_schedule)); + self.consume_cycles(CompoundCycles::::new(balance, cost_schedule)); } fn observe_consumed_cycles_with_use_case( &mut self, - amount: NominalCycles, + prepayment: NominalCycles, + refund: NominalCycles, use_case: CyclesUseCase, consuming_cycles: ConsumingCycles, ) { @@ -1919,7 +1953,30 @@ impl SystemState { debug_assert_ne!(use_case, CyclesUseCase::DeletedCanisters); debug_assert_ne!(use_case, CyclesUseCase::DroppedMessages); - if use_case == CyclesUseCase::NonConsumed || amount.is_zero() { + // CyclesUseCase::NonConsumed should never be sent to this function. + debug_assert_ne!( + use_case, + CyclesUseCase::NonConsumed, + "Non-consumed cycles should not be tracked in the canister metrics." + ); + + // `prepayment`` must be greater or equal to `refund`. + // `refund` must be 0 when we are handling a prepayment. + debug_assert!( + prepayment >= refund, + "Expected prepayment: {prepayment} to be greater or equal to refund {refund}." + ); + match consuming_cycles { + ConsumingCycles::Prepayment => { + debug_assert_eq!(refund, NominalCycles::zero()); + } + ConsumingCycles::Refund => {} + } + + // Skip if the amounts are zero and no metric updates are needed. + if (consuming_cycles == ConsumingCycles::Prepayment && prepayment.is_zero()) + || (consuming_cycles == ConsumingCycles::Refund && refund.is_zero()) + { return; } @@ -1929,13 +1986,13 @@ impl SystemState { let use_case_consumption = metric.entry(use_case).or_insert_with(NominalCycles::zero); match consuming_cycles { - ConsumingCycles::Yes => { - *use_case_consumption += amount; - self.canister_metrics.consumed_cycles += amount; + ConsumingCycles::Prepayment => { + *use_case_consumption += prepayment; + self.canister_metrics.consumed_cycles += prepayment; } - ConsumingCycles::No => { - *use_case_consumption -= amount; - self.canister_metrics.consumed_cycles -= amount; + ConsumingCycles::Refund => { + *use_case_consumption -= refund; + self.canister_metrics.consumed_cycles -= refund; } } } diff --git a/rs/replicated_state/src/canister_state/tests.rs b/rs/replicated_state/src/canister_state/tests.rs index 36df606c4e41..e685a5570428 100644 --- a/rs/replicated_state/src/canister_state/tests.rs +++ b/rs/replicated_state/src/canister_state/tests.rs @@ -146,14 +146,12 @@ impl CanisterStateFixture { &mut self, msg: RequestOrResponse, subnet_type: SubnetType, - cost_schedule: CanisterCyclesCostSchedule, input_queue_type: InputQueueType, ) -> Result { self.canister_state.push_input( msg, &mut SUBNET_AVAILABLE_MEMORY.clone(), subnet_type, - cost_schedule, input_queue_type, ) } @@ -181,7 +179,6 @@ impl CanisterStateFixture { self.push_input( message.clone(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ) .unwrap() @@ -213,7 +210,6 @@ fn canister_state_push_input_request_success() { .push_input( default_input_request(NO_DEADLINE), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ) .unwrap() @@ -234,7 +230,6 @@ fn canister_state_push_input_response_success() { .push_input( response, SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ) .unwrap() @@ -255,7 +250,6 @@ fn canister_state_push_input_guaranteed_response_no_reserved_slot() { fixture.push_input( response.into(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet ), ); @@ -275,7 +269,6 @@ fn canister_state_push_input_best_effort_response_no_reserved_slot() { .push_input( response.clone().into(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ) .unwrap() @@ -287,7 +280,6 @@ fn canister_state_push_input_best_effort_response_no_reserved_slot() { .push_input( response.clone().into(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ) .unwrap() @@ -320,7 +312,6 @@ fn canister_state_push_input_best_effort_response_canister_stopped() { fixture.push_input( response.into(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet ) ); @@ -337,7 +328,6 @@ fn canister_state_push_input_guaranteed_response_no_matching_callback() { fixture.push_input( response, SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet ), Err((StateError::NonMatchingResponse { .. }, _)) @@ -359,7 +349,6 @@ fn canister_state_push_input_best_effort_response_no_matching_callback() { .push_input( response, SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ) .unwrap() @@ -377,7 +366,6 @@ fn canister_state_push_input_guaranteed_response_mismatched_callback() { fixture.push_input( response.clone().into(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet ), Err(( @@ -397,7 +385,6 @@ fn canister_state_push_input_best_effort_response_mismatched_callback() { fixture.push_input( response.clone().into(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet ), Err(( @@ -420,7 +407,6 @@ fn canister_state_push_input_request_mismatched_receiver() { .build() .into(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ); } @@ -436,7 +422,6 @@ fn canister_state_push_input_response_mismatched_originator() { .build() .into(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ); } @@ -453,7 +438,6 @@ fn canister_state_push_input_guaranteed_response_duplicate_of_paused_response() fixture.push_input( response.clone(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ), Err((StateError::NonMatchingResponse { err_str, .. }, r)) @@ -476,7 +460,6 @@ fn canister_state_push_input_best_effort_response_duplicate_of_paused_response() .push_input( response.clone(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ) .unwrap() @@ -531,7 +514,6 @@ fn canister_state_induct_messages_to_self_duplicate_of_paused_response(deadline: .push_input( request.clone().into(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::LocalSubnet, ) .unwrap() @@ -544,7 +526,6 @@ fn canister_state_induct_messages_to_self_duplicate_of_paused_response(deadline: .push_input( response.clone().into(), SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::LocalSubnet, ) .unwrap() @@ -589,7 +570,6 @@ fn application_subnet_remote_push_input_request_not_enough_subnet_memory() { NO_DEADLINE, 13, SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, true, ); @@ -601,7 +581,6 @@ fn application_subnet_local_push_input_request_not_enough_subnet_memory() { NO_DEADLINE, 13, SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::LocalSubnet, true, ); @@ -613,7 +592,6 @@ fn system_subnet_remote_push_input_request_not_enough_subnet_memory() { NO_DEADLINE, 13, SubnetType::System, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, true, ); @@ -625,7 +603,6 @@ fn system_subnet_local_push_input_request_ignores_subnet_memory() { NO_DEADLINE, 13, SubnetType::System, - CanisterCyclesCostSchedule::Normal, InputQueueType::LocalSubnet, false, ); @@ -637,7 +614,6 @@ fn application_subnet_push_input_best_effort_request_ignores_subnet_memory() { SOME_DEADLINE, -13, SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, false, ); @@ -649,7 +625,6 @@ fn system_subnet_push_input_best_effort_request_ignores_subnet_memory() { SOME_DEADLINE, -13, SubnetType::System, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, false, ); @@ -666,7 +641,6 @@ fn canister_state_push_input_request_memory_limit_test_impl( deadline: CoarseTime, initial_subnet_available_memory: i64, own_subnet_type: SubnetType, - own_cost_schedule: CanisterCyclesCostSchedule, input_queue_type: InputQueueType, should_enforce_limit: bool, ) { @@ -679,7 +653,6 @@ fn canister_state_push_input_request_memory_limit_test_impl( request.clone(), &mut subnet_available_memory, own_subnet_type, - own_cost_schedule, input_queue_type, ); if should_enforce_limit { @@ -715,7 +688,6 @@ fn system_subnet_remote_push_input_request_ignores_memory_reservation_and_execut // Remote message inducted into system subnet. let own_subnet_type = SubnetType::System; - let own_cost_schedule = CanisterCyclesCostSchedule::Normal; let input_queue_type = InputQueueType::RemoteSubnet; // Tiny explicit allocation, not enough for a request. @@ -745,7 +717,6 @@ fn system_subnet_remote_push_input_request_ignores_memory_reservation_and_execut request, &mut subnet_available_memory, own_subnet_type, - own_cost_schedule, input_queue_type, ) .unwrap() @@ -768,7 +739,6 @@ fn system_subnet_remote_push_input_request_ignores_memory_reservation_and_execut fn application_subnet_remote_push_input_response_ignores_memory_limits() { canister_state_push_input_response_memory_limit_test_impl( SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ); } @@ -777,7 +747,6 @@ fn application_subnet_remote_push_input_response_ignores_memory_limits() { fn application_subnet_local_push_input_response_ignores_memory_limits() { canister_state_push_input_response_memory_limit_test_impl( SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::LocalSubnet, ); } @@ -786,7 +755,6 @@ fn application_subnet_local_push_input_response_ignores_memory_limits() { fn system_subnet_remote_push_input_response_ignores_memory_limits() { canister_state_push_input_response_memory_limit_test_impl( SubnetType::System, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ); } @@ -795,7 +763,6 @@ fn system_subnet_remote_push_input_response_ignores_memory_limits() { fn system_subnet_local_push_input_response_ignores_memory_limits() { canister_state_push_input_response_memory_limit_test_impl( SubnetType::System, - CanisterCyclesCostSchedule::Normal, InputQueueType::LocalSubnet, ); } @@ -809,7 +776,6 @@ fn system_subnet_local_push_input_response_ignores_memory_limits() { /// always return memory). fn canister_state_push_input_response_memory_limit_test_impl( own_subnet_type: SubnetType, - own_cost_schedule: CanisterCyclesCostSchedule, input_queue_type: InputQueueType, ) { let mut fixture = CanisterStateFixture::new(); @@ -827,7 +793,6 @@ fn canister_state_push_input_response_memory_limit_test_impl( response.clone(), &mut subnet_available_memory, own_subnet_type, - own_cost_schedule, input_queue_type, ) .unwrap() @@ -922,10 +887,9 @@ fn update_balance_and_consumed_cycles_correctly() { let mut system_state = CanisterStateFixture::new().canister_state.system_state; let initial_consumed_cycles = Cycles::new(1000); let cost_schedule = CanisterCyclesCostSchedule::Normal; - system_state.remove_cycles(CompoundCycles::::new( - initial_consumed_cycles, - cost_schedule, - )); + let prepaid_cycles = + CompoundCycles::::new(initial_consumed_cycles, cost_schedule); + system_state.consume_cycles(prepaid_cycles); assert_eq!( system_state.balance(), INITIAL_CYCLES - initial_consumed_cycles @@ -935,18 +899,15 @@ fn update_balance_and_consumed_cycles_correctly() { NominalCycles::new(initial_consumed_cycles.get()) ); - let cycles = Cycles::new(100); - system_state.add_cycles(CompoundCycles::::new( - cycles, - cost_schedule, - )); + let refund = CompoundCycles::::new(Cycles::new(100), cost_schedule); + system_state.refund_cycles(prepaid_cycles, refund); assert_eq!( system_state.balance(), - INITIAL_CYCLES - initial_consumed_cycles + cycles + INITIAL_CYCLES - initial_consumed_cycles + refund.real() ); assert_eq!( system_state.canister_metrics().consumed_cycles(), - NominalCycles::new((initial_consumed_cycles - cycles).get()) + NominalCycles::new((initial_consumed_cycles - refund.real()).get()) ); } @@ -955,19 +916,16 @@ fn update_balance_and_consumed_cycles_by_use_case_correctly() { let mut system_state = CanisterStateFixture::new().canister_state.system_state; let cycles_to_consume = Cycles::from(1000_u128); let cost_schedule = CanisterCyclesCostSchedule::Normal; - system_state.remove_cycles(CompoundCycles::::new( - cycles_to_consume, - cost_schedule, - )); + let prepaid_cycles = + CompoundCycles::::new(cycles_to_consume, cost_schedule); + system_state.consume_cycles(prepaid_cycles); - let cycles_to_add = Cycles::from(100_u128); - system_state.add_cycles(CompoundCycles::::new( - cycles_to_add, - cost_schedule, - )); + let refund = + CompoundCycles::::new(Cycles::from(100u128), cost_schedule); + system_state.refund_cycles(prepaid_cycles, refund); assert_eq!( system_state.balance(), - INITIAL_CYCLES - cycles_to_consume + cycles_to_add + INITIAL_CYCLES - cycles_to_consume + refund.real() ); assert_eq!( *system_state @@ -975,7 +933,7 @@ fn update_balance_and_consumed_cycles_by_use_case_correctly() { .consumed_cycles_by_use_cases() .get(&CyclesUseCase::Memory) .unwrap(), - NominalCycles::new((cycles_to_consume - cycles_to_add).get()) + NominalCycles::new((cycles_to_consume - refund.real()).get()) ); } diff --git a/rs/replicated_state/src/replicated_state.rs b/rs/replicated_state/src/replicated_state.rs index 95155d314be1..7fe3e8ac91e1 100644 --- a/rs/replicated_state/src/replicated_state.rs +++ b/rs/replicated_state/src/replicated_state.rs @@ -35,7 +35,7 @@ use ic_types::{ time::CoarseTime, }; use ic_types_cycles::{ - CanisterCyclesCostSchedule, CompoundCycles, CyclesUseCaseKind, DroppedMessages, NonConsumed, + CanisterCyclesCostSchedule, CompoundCycles, CyclesUseCaseKind, DroppedMessages, }; use ic_validate_eq::ValidateEq; use ic_validate_eq_derive::ValidateEq; @@ -1002,7 +1002,6 @@ impl ReplicatedState { msg, subnet_available_guaranteed_response_memory, own_subnet_type, - own_cost_schedule, input_queue_type, ), None => { @@ -1068,18 +1067,9 @@ impl ReplicatedState { /// /// Returns `true` if the recipient canister exists and was credited, `false` /// otherwise. - pub fn credit_refund( - &mut self, - refund: &Refund, - cost_schedule: CanisterCyclesCostSchedule, - ) -> bool { + pub fn credit_refund(&mut self, refund: &Refund) -> bool { if let Some(canister) = self.canister_state_make_mut(&refund.recipient()) { - canister - .system_state - .add_cycles(CompoundCycles::::new( - refund.amount(), - cost_schedule, - )); + canister.system_state.add_cycles(refund.amount()); true } else { false diff --git a/rs/replicated_state/tests/canister_state.rs b/rs/replicated_state/tests/canister_state.rs index 2d65b5095dc0..b7af85464164 100644 --- a/rs/replicated_state/tests/canister_state.rs +++ b/rs/replicated_state/tests/canister_state.rs @@ -14,7 +14,6 @@ use ic_types::{ messages::{CallbackId, NO_DEADLINE, Request, RequestOrResponse}, time::{CoarseTime, UNIX_EPOCH}, }; -use ic_types_cycles::CanisterCyclesCostSchedule; use std::sync::Arc; const CANISTER_ID: CanisterId = CanisterId::from_u64(0); @@ -85,7 +84,6 @@ impl CanisterFixture { msg, &mut subnet_available_memory, SubnetType::Application, - CanisterCyclesCostSchedule::Normal, InputQueueType::RemoteSubnet, ) } diff --git a/rs/replicated_state/tests/replicated_state.rs b/rs/replicated_state/tests/replicated_state.rs index e29629f5b739..8cabf62ba118 100644 --- a/rs/replicated_state/tests/replicated_state.rs +++ b/rs/replicated_state/tests/replicated_state.rs @@ -1505,10 +1505,9 @@ fn credit_refund() { let initial_balance = fixture.canister_balance(&CANISTER_ID).unwrap(); // Refund 10 cycles to `CANISTER_ID`. - let credited = fixture.state.credit_refund( - &Refund::anonymous(CANISTER_ID, Cycles::new(10)), - CanisterCyclesCostSchedule::Normal, - ); + let credited = fixture + .state + .credit_refund(&Refund::anonymous(CANISTER_ID, Cycles::new(10))); assert!(credited); assert_eq!( Some(initial_balance + Cycles::new(10)), @@ -1516,10 +1515,9 @@ fn credit_refund() { ); // Refunding to a non-existent canister is a no-op. - let credited = fixture.state.credit_refund( - &Refund::anonymous(OTHER_CANISTER_ID, Cycles::new(11)), - CanisterCyclesCostSchedule::Normal, - ); + let credited = fixture + .state + .credit_refund(&Refund::anonymous(OTHER_CANISTER_ID, Cycles::new(11))); assert!(!credited); assert_eq!( Some(initial_balance + Cycles::new(10)), diff --git a/rs/state_machine_tests/src/lib.rs b/rs/state_machine_tests/src/lib.rs index 66957fd03b82..9a319c981ee6 100644 --- a/rs/state_machine_tests/src/lib.rs +++ b/rs/state_machine_tests/src/lib.rs @@ -178,9 +178,7 @@ use ic_types::{ time::{GENESIS, Time}, xnet::{CertifiedStreamSlice, StreamIndex, StreamSlice}, }; -use ic_types_cycles::{ - CanisterCyclesCostSchedule, CompoundCycles, Cycles, CyclesUseCase, NominalCycles, NonConsumed, -}; +use ic_types_cycles::{CanisterCyclesCostSchedule, Cycles, CyclesUseCase, NominalCycles}; use ic_xnet_payload_builder::{ RefillTaskHandle, XNetPayloadBuilderImpl, XNetPayloadBuilderMetrics, XNetSlicePoolImpl, certified_slice_pool::CertifiedSlicePool, refill_stream_slice_indices, @@ -4911,12 +4909,7 @@ impl StateMachine { let canister_state = state .canister_state_make_mut(&canister_id) .unwrap_or_else(|| panic!("Canister {canister_id} not found")); - canister_state - .system_state - .add_cycles(CompoundCycles::::new( - Cycles::from(amount), - self.cost_schedule, - )); + canister_state.system_state.add_cycles(Cycles::from(amount)); let balance = canister_state.system_state.balance().get(); self.state_manager .commit_and_certify(state, CertificationScope::Metadata, None); diff --git a/rs/test_utilities/execution_environment/src/lib.rs b/rs/test_utilities/execution_environment/src/lib.rs index bdb3da85856e..990b47ae25bd 100644 --- a/rs/test_utilities/execution_environment/src/lib.rs +++ b/rs/test_utilities/execution_environment/src/lib.rs @@ -1386,7 +1386,6 @@ impl ExecutionTest { RequestOrResponse::Response(response_arc.clone()), &mut subnet_available_guaranteed_response_memory, state.metadata.own_subnet_type, - cost_schedule, InputQueueType::LocalSubnet, ) .unwrap(); @@ -2014,7 +2013,6 @@ impl ExecutionTest { /// `self.xnet_messages`. pub fn induct_messages(&mut self) { let mut state = self.state.take().unwrap(); - let cost_schedule = state.get_own_cost_schedule(); let mut subnet_available_guaranteed_response_memory = self .subnet_available_memory .get_guaranteed_response_message_memory(); @@ -2027,7 +2025,6 @@ impl ExecutionTest { message.clone(), &mut subnet_available_guaranteed_response_memory, state.metadata.own_subnet_type, - cost_schedule, InputQueueType::LocalSubnet, ); if result.is_err() { From cf8e86d9e0b6ab8dc8414670b19218b9d947e3f8 Mon Sep 17 00:00:00 2001 From: Dimitris Sarlis Date: Tue, 7 Apr 2026 09:04:23 +0000 Subject: [PATCH 2/3] cargo fix --- rs/replicated_state/src/canister_state/tests.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rs/replicated_state/src/canister_state/tests.rs b/rs/replicated_state/src/canister_state/tests.rs index e685a5570428..ce8aef2cd6ed 100644 --- a/rs/replicated_state/src/canister_state/tests.rs +++ b/rs/replicated_state/src/canister_state/tests.rs @@ -921,7 +921,7 @@ fn update_balance_and_consumed_cycles_by_use_case_correctly() { system_state.consume_cycles(prepaid_cycles); let refund = - CompoundCycles::::new(Cycles::from(100u128), cost_schedule); + CompoundCycles::::new(Cycles::from(100_u128), cost_schedule); system_state.refund_cycles(prepaid_cycles, refund); assert_eq!( system_state.balance(), From eb1e6f0c25c8a6adbe143c37a9f9e26fe588a9a7 Mon Sep 17 00:00:00 2001 From: Dimitris Sarlis Date: Tue, 7 Apr 2026 09:09:21 +0000 Subject: [PATCH 3/3] Review comments --- .../src/canister_state/system_state.rs | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/rs/replicated_state/src/canister_state/system_state.rs b/rs/replicated_state/src/canister_state/system_state.rs index 491271431da8..d1aed1e88f21 100644 --- a/rs/replicated_state/src/canister_state/system_state.rs +++ b/rs/replicated_state/src/canister_state/system_state.rs @@ -1848,8 +1848,7 @@ impl SystemState { ); } - /// Decreases 'cycles_balance' for 'requested_amount' from the canister's - /// main balance. + /// Decreases the canister's cycles balance by the provided amount. /// /// If you need to have metrics about consumed cycles updated use /// `consume_cycles` instead. @@ -1857,9 +1856,14 @@ impl SystemState { self.cycles_balance -= requested_amount; } - /// Decreases 'cycles_balance' for 'requested_amount'. + /// Decreases the canister's cycles balance by the provided amount. /// The resource use cases first drain the `reserved_balance` and only after /// that drain the main `cycles_balance`. + /// + /// Similar to `remove_cycles` but additionally updates the metrics to match + /// the consumed amount. Should be used either for cases where a prepayment + /// needs to be made (that will be refunded later with `refund_cycles`) or + /// a direct charge happens without a prepayment (e.g. when paying for memory). pub fn consume_cycles(&mut self, requested_amount: CompoundCycles) { let requested_real = requested_amount.real(); let use_case = T::cycles_use_case(); @@ -1960,7 +1964,7 @@ impl SystemState { "Non-consumed cycles should not be tracked in the canister metrics." ); - // `prepayment`` must be greater or equal to `refund`. + // `prepayment` must be greater or equal to `refund`. // `refund` must be 0 when we are handling a prepayment. debug_assert!( prepayment >= refund,