refactor: Separate adding/removing cycles from refunding/consuming cycles

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(())
     }
 

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::<NonConsumed>::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::<NonConsumed>::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::<NonConsumed>::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::<NonConsumed>::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::<NonConsumed>::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,

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
         }

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::<NonConsumed>::new(added, cost_schedule))
-            }
-            CyclesBalanceChange::Removed(removed) => {
-                state.remove_cycles(CompoundCycles::<NonConsumed>::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());
     }

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::<NonConsumed>::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);

rs/execution_environment/src/canister_manager.rs

@@ -70,7 +70,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};
@@ -1490,7 +1490,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,
@@ -1568,7 +1568,6 @@ impl CanisterManager {
         cycles_amount: Option<u128>,
         canister: &mut CanisterState,
         provisional_whitelist: &ProvisionalWhitelist,
-        cost_schedule: CanisterCyclesCostSchedule,
     ) -> Result<(), CanisterManagerError> {
         if !provisional_whitelist.contains(&sender) {
             return Err(CanisterManagerError::SenderNotInWhitelist(sender));
@@ -1579,12 +1578,7 @@ impl CanisterManager {
             None => self.config.default_provisional_cycles_balance,
         };
 
-        canister
-            .system_state
-            .add_cycles(CompoundCycles::<NonConsumed>::new(
-                cycles_amount,
-                cost_schedule,
-            ));
+        canister.system_state.add_cycles(cycles_amount);
 
         Ok(())
     }
@@ -1594,13 +1588,10 @@ impl CanisterManager {
         canister: &mut CanisterState,
         cycles: Cycles,
         sender: PrincipalId,
-        cost_schedule: CanisterCyclesCostSchedule,
     ) -> CanisterManagerResponse {
         let canister_id = canister.canister_id();
 
-        canister
-            .system_state
-            .add_cycles(CompoundCycles::<NonConsumed>::new(cycles, cost_schedule));
+        canister.system_state.add_cycles(cycles);
 
         if cycles.get() > LOG_CANISTER_OPERATION_CYCLES_THRESHOLD {
             info!(

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))
         );

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::<NonConsumed>::new(
-                balance + Cycles::new(1),
-                CanisterCyclesCostSchedule::Normal,
-            ));
+            .add_cycles(balance + Cycles::new(1));
 
         test.execute_slice(a_id);