feat(system-api): add call_cycles_add128_up_to (dfinity#1158)

Introduces `call_cycles_add128_up_to` to the system API to offer a way
of adding as many cycles as possible to a call without trapping in the
subsequent `call_perform`.
Corresponding spec MR:
dfinity/interface-spec#316
Closes
[SDK-1761](https://dfinity.atlassian.net/browse/SDK-1761?atlOrigin=eyJpIjoiY2QyZTJiZDRkNGZhNGZlMWI3NzRkNTBmZmVlNzNiZTkiLCJwIjoianN3LWdpdGxhYi1pbnQifQ)

[SDK-1761]:
https://dfinity.atlassian.net/browse/SDK-1761?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

rs/cycles_account_manager/src/lib.rs

@@ -25,7 +25,10 @@ use ic_replicated_state::{
 };
 use ic_types::{
     canister_http::MAX_CANISTER_HTTP_RESPONSE_BYTES,
-    messages::{Request, Response, SignedIngressContent, MAX_INTER_CANISTER_PAYLOAD_IN_BYTES},
+    messages::{
+        Request, Response, SignedIngressContent, MAX_INTER_CANISTER_PAYLOAD_IN_BYTES,
+        MAX_RESPONSE_COUNT_BYTES,
+    },
     CanisterId, ComputeAllocation, Cycles, MemoryAllocation, NumBytes, NumInstructions, SubnetId,
 };
 use prometheus::IntCounter;
@@ -374,6 +377,48 @@ impl CyclesAccountManager {
         )
     }
 
+    /// Withdraws up to `cycles` worth of cycles from the canister's balance
+    /// without putting the canister below its freezing threshold even if
+    /// the call currently under construction is performed.
+    ///
+    /// NOTE: This method is intended for use in inter-canister transfers.
+    ///       It doesn't report these cycles as consumed. To withdraw cycles
+    ///       and have them reported as consumed, use `consume_cycles`.
+    #[allow(clippy::too_many_arguments)]
+    pub fn withdraw_up_to_cycles_for_transfer(
+        &self,
+        freeze_threshold: NumSeconds,
+        memory_allocation: MemoryAllocation,
+        current_payload_size_bytes: NumBytes,
+        canister_current_memory_usage: NumBytes,
+        canister_current_message_memory_usage: NumBytes,
+        canister_compute_allocation: ComputeAllocation,
+        cycles_balance: &mut Cycles,
+        cycles: Cycles,
+        subnet_size: usize,
+        reserved_balance: Cycles,
+    ) -> Cycles {
+        let call_perform_cost = self.xnet_call_performed_fee(subnet_size)
+            + self.xnet_call_bytes_transmitted_fee(current_payload_size_bytes, subnet_size)
+            + self.prepayment_for_response_transmission(subnet_size)
+            + self.prepayment_for_response_execution(subnet_size);
+        let memory_used_to_enqueue_message =
+            current_payload_size_bytes.max((MAX_RESPONSE_COUNT_BYTES as u64).into());
+        let freeze_threshold = self.freeze_threshold_cycles(
+            freeze_threshold,
+            memory_allocation,
+            canister_current_memory_usage,
+            canister_current_message_memory_usage + memory_used_to_enqueue_message,
+            canister_compute_allocation,
+            subnet_size,
+            reserved_balance,
+        );
+        let available_for_withdrawal = *cycles_balance - freeze_threshold - call_perform_cost;
+        let withdrawn_cycles = available_for_withdrawal.min(cycles);
+        *cycles_balance -= withdrawn_cycles;
+        withdrawn_cycles
+    }
+
     /// Charges the canister for ingress induction cost.
     ///
     /// Note that this method reports the cycles withdrawn as consumed (i.e.

rs/cycles_account_manager/tests/cycles_account_manager.rs

@@ -17,7 +17,7 @@ use ic_test_utilities_types::{
     messages::SignedIngressBuilder,
 };
 use ic_types::{
-    messages::{extract_effective_canister_id, SignedIngressContent},
+    messages::{extract_effective_canister_id, SignedIngressContent, MAX_RESPONSE_COUNT_BYTES},
     nominal_cycles::NominalCycles,
     CanisterId, ComputeAllocation, Cycles, MemoryAllocation, NumBytes, NumInstructions,
 };
@@ -880,6 +880,75 @@ fn withdraw_cycles_for_transfer_checks_reserved_balance() {
     assert_eq!(Cycles::zero(), new_balance);
 }
 
+#[test]
+fn withdraw_up_to_respects_freezing_threshold() {
+    let cycles_account_manager = CyclesAccountManagerBuilder::new().build();
+    let initial_cycles = Cycles::new(200_000_000_000);
+    let mut system_state = SystemState::new_running_for_testing(
+        canister_test_id(1),
+        canister_test_id(2).get(),
+        initial_cycles,
+        NumSeconds::from(1_000),
+    );
+    let memory_usage = NumBytes::from(1_000_000);
+    let message_memory_usage = NumBytes::from(1_000);
+    let compute_allocation = ComputeAllocation::default();
+    let payload_size = NumBytes::from(0);
+    system_state.memory_allocation = MemoryAllocation::try_from(NumBytes::from(1 << 20)).unwrap();
+    let untouched_cycles = cycles_account_manager.freeze_threshold_cycles(
+        system_state.freeze_threshold,
+        system_state.memory_allocation,
+        memory_usage,
+        message_memory_usage + (MAX_RESPONSE_COUNT_BYTES as u64).into(),
+        compute_allocation,
+        SMALL_APP_SUBNET_MAX_SIZE,
+        system_state.reserved_balance(),
+    ) + cycles_account_manager
+        .xnet_call_performed_fee(SMALL_APP_SUBNET_MAX_SIZE)
+        + cycles_account_manager
+            .xnet_call_bytes_transmitted_fee(payload_size, SMALL_APP_SUBNET_MAX_SIZE)
+        + cycles_account_manager.prepayment_for_response_transmission(SMALL_APP_SUBNET_MAX_SIZE)
+        + cycles_account_manager.prepayment_for_response_execution(SMALL_APP_SUBNET_MAX_SIZE);
+
+    // full amount can be withdrawn
+    let mut new_balance = system_state.balance();
+    let withdraw_amount_1 = Cycles::new(1_000_000);
+    let withdrawn_amount_1 = cycles_account_manager.withdraw_up_to_cycles_for_transfer(
+        system_state.freeze_threshold,
+        system_state.memory_allocation,
+        payload_size,
+        memory_usage,
+        message_memory_usage,
+        compute_allocation,
+        &mut new_balance,
+        withdraw_amount_1,
+        SMALL_APP_SUBNET_MAX_SIZE,
+        system_state.reserved_balance(),
+    );
+    assert_eq!(withdraw_amount_1, withdrawn_amount_1);
+    assert_eq!(initial_cycles - withdraw_amount_1, new_balance);
+
+    // freezing threshold limits the amount that can be withdrawn
+    let withdraw_amount_2 = Cycles::new(u128::MAX);
+    let withdrawn_amount_2 = cycles_account_manager.withdraw_up_to_cycles_for_transfer(
+        system_state.freeze_threshold,
+        system_state.memory_allocation,
+        payload_size,
+        memory_usage,
+        message_memory_usage,
+        compute_allocation,
+        &mut new_balance,
+        withdraw_amount_2,
+        SMALL_APP_SUBNET_MAX_SIZE,
+        system_state.reserved_balance(),
+    );
+    assert_eq!(
+        initial_cycles - withdrawn_amount_1 - untouched_cycles,
+        withdrawn_amount_2
+    );
+    assert_eq!(untouched_cycles, new_balance);
+}
+
 #[test]
 fn freezing_threshold_uses_reserved_balance() {
     let cycles_account_manager = CyclesAccountManagerBuilder::new().build();

rs/embedders/src/wasm_utils/validation.rs

@@ -535,6 +535,16 @@ fn get_valid_system_apis_common(I: ValType) -> HashMap<String, HashMap<String, F
                 },
             )],
         ),
+        (
+            "call_cycles_add128_up_to",
+            vec![(
+                API_VERSION_IC0,
+                FunctionSignature {
+                    param_types: vec![ValType::I64, ValType::I64, ValType::I32],
+                    return_type: vec![],
+                },
+            )],
+        ),
         (
             "canister_cycle_balance128",
             vec![(

rs/embedders/src/wasmtime_embedder/system_api.rs

@@ -749,6 +749,21 @@ pub(crate) fn syscalls<
         })
         .unwrap();
 
+    linker
+        .func_wrap("ic0", "call_cycles_add128_up_to", {
+            move |mut caller: Caller<'_, StoreData>, amount_high: u64, amount_low: u64, dst: u32| {
+                charge_for_cpu(&mut caller, overhead::CALL_CYCLES_ADD128_UP_TO)?;
+                with_memory_and_system_api(&mut caller, |system, memory| {
+                    system.ic0_call_cycles_add128_up_to(
+                        Cycles::from_parts(amount_high, amount_low),
+                        dst as usize,
+                        memory,
+                    )
+                })
+            }
+        })
+        .unwrap();
+
     linker
         .func_wrap("ic0", "call_perform", {
             move |mut caller: Caller<'_, StoreData>| {

rs/embedders/src/wasmtime_embedder/system_api_complexity.rs

@@ -21,6 +21,7 @@ pub mod overhead {
     pub const ACCEPT_MESSAGE: NumInstructions = NumInstructions::new(500);
     pub const CALL_CYCLES_ADD: NumInstructions = NumInstructions::new(500);
     pub const CALL_CYCLES_ADD128: NumInstructions = NumInstructions::new(500);
+    pub const CALL_CYCLES_ADD128_UP_TO: NumInstructions = NumInstructions::new(500);
     pub const CALL_DATA_APPEND: NumInstructions = NumInstructions::new(500);
     pub const CALL_NEW: NumInstructions = NumInstructions::new(1_500);
     pub const CALL_ON_CLEANUP: NumInstructions = NumInstructions::new(500);

rs/embedders/tests/validation.rs

@@ -926,6 +926,7 @@ fn can_validate_module_cycles_u128_related_imports() {
     let wasm = wat2wasm(
         r#"(module
         (import "ic0" "call_cycles_add128" (func $ic0_call_cycles_add128 (param i64 i64)))
+        (import "ic0" "call_cycles_add128_up_to" (func $ic0_call_cycles_add128_up_to (param i64 i64 i32)))
         (import "ic0" "canister_cycle_balance128" (func $ic0_canister_cycle_balance128 (param i32)))
         (import "ic0" "msg_cycles_available128" (func $ic0_msg_cycles_available128 (param i32)))
         (import "ic0" "msg_cycles_refunded128" (func $ic0_msg_cycles_refunded128 (param i32)))

rs/execution_environment/benches/system_api/execute_update.rs

@@ -155,6 +155,15 @@ pub fn execute_update_bench(c: &mut Criterion) {
             Module::CallNewLoop.from_ic0("call_cycles_add128", Params2(0_i64, 100_i64), Result::No),
             2059000006,
         ),
+        common::Benchmark(
+            "call_new+ic0_call_cycles_add128_up_to()".into(),
+            Module::CallNewLoop.from_ic0(
+                "call_cycles_add128_up_to",
+                Params3(0_i64, 100_i64, 0_i32),
+                Result::No,
+            ),
+            2059000006,
+        ),
         common::Benchmark(
             "call_new+ic0_call_perform()".into(),
             Module::CallNewLoop.from_ic0("call_perform", NoParams, Result::I32),