From 0a11eb3a3d0525a6d7179224a04daa9fec00fb4f Mon Sep 17 00:00:00 2001
From: Andriy Berestovskyy <andriy.berestovskyy@dfinity.org>
Date: Mon, 11 Dec 2023 12:42:59 +0000
Subject: [PATCH] feat: RUN-847: Track time, balance and call System APIs

---
 .../src/wasmtime_embedder/system_api.rs       |   4 +-
 rs/interfaces/src/execution_environment.rs    |  22 ++-
 rs/system_api/BUILD.bazel                     |   9 ++
 rs/system_api/src/lib.rs                      |  41 ++++-
 rs/system_api/tests/system_api.rs             |   4 +-
 rs/system_api/tests/system_api_tracking.rs    | 149 ++++++++++++++++++
 6 files changed, 218 insertions(+), 11 deletions(-)
 create mode 100644 rs/system_api/tests/system_api_tracking.rs

diff --git a/rs/embedders/src/wasmtime_embedder/system_api.rs b/rs/embedders/src/wasmtime_embedder/system_api.rs
index 44fe546dc62..a33e1aeeb35 100644
--- a/rs/embedders/src/wasmtime_embedder/system_api.rs
+++ b/rs/embedders/src/wasmtime_embedder/system_api.rs
@@ -891,8 +891,8 @@ pub(crate) fn syscalls(
                     &mut caller,
                     overhead!(CANISTER_CYCLE_BALANCE128, metering_type),
                 )?;
-                with_memory_and_system_api(&mut caller, |system_api, memory| {
-                    system_api.ic0_canister_cycle_balance128(dst, memory)
+                with_memory_and_system_api(&mut caller, |s, memory| {
+                    s.ic0_canister_cycle_balance128(dst, memory)
                 })?;
                 if feature_flags.write_barrier == FlagStatus::Enabled {
                     mark_writes_on_bytemap(&mut caller, dst as usize, 16)
diff --git a/rs/interfaces/src/execution_environment.rs b/rs/interfaces/src/execution_environment.rs
index 93e7f3bd092..de0113d7ed1 100644
--- a/rs/interfaces/src/execution_environment.rs
+++ b/rs/interfaces/src/execution_environment.rs
@@ -85,6 +85,7 @@ pub enum PerformanceCounterType {
 }
 
 /// System API call ids to track their execution (in alphabetical order).
+#[derive(Debug)]
 pub enum SystemApiCallId {
     /// Tracker for `ic0.accept_message())`
     AcceptMessage,
@@ -194,6 +195,21 @@ pub enum SystemApiCallId {
     UpdateAvailableMemory,
 }
 
+/// System API call counters, i.e. how many times each tracked System API call
+/// was invoked.
+// #[derive(Serialize, Deserialize, Clone, Debug, Default)]
+#[derive(Clone, Default)]
+pub struct SystemApiCallCounters {
+    /// Counter for `ic0.call_perform()`
+    pub call_perform: usize,
+    /// Counter for `ic0.canister_cycle_balance()`
+    pub canister_cycle_balance: usize,
+    /// Counter for `ic0.canister_cycle_balance128()`
+    pub canister_cycle_balance128: usize,
+    /// Counter for `ic0.time()`
+    pub time: usize,
+}
+
 /// Tracks the available memory on a subnet. The main idea is to separately track
 /// the execution available memory, the message available memory and the wasm custom
 /// sections available memory. The different flavors of memory are independent of each
@@ -832,7 +848,7 @@ pub trait SystemApi {
     ) -> HypervisorResult<(NumPages, NumInstructions)>;
 
     /// The canister can query the IC for the current time.
-    fn ic0_time(&self) -> HypervisorResult<Time>;
+    fn ic0_time(&mut self) -> HypervisorResult<Time>;
 
     /// The canister can set a global one-off timer at the specific time.
     fn ic0_global_timer_set(&mut self, time: Time) -> HypervisorResult<Time>;
@@ -898,7 +914,7 @@ pub trait SystemApi {
     /// Returns the current balance in cycles.
     ///
     /// Traps if current canister balance cannot fit in a 64-bit value.
-    fn ic0_canister_cycle_balance(&self) -> HypervisorResult<u64>;
+    fn ic0_canister_cycle_balance(&mut self) -> HypervisorResult<u64>;
 
     /// This system call indicates the current cycle balance
     /// of the canister.
@@ -906,7 +922,7 @@ pub trait SystemApi {
     /// The amount of cycles is represented by a 128-bit value
     /// and is copied in the canister memory starting
     /// starting at the location `dst`.
-    fn ic0_canister_cycle_balance128(&self, dst: u32, heap: &mut [u8]) -> HypervisorResult<()>;
+    fn ic0_canister_cycle_balance128(&mut self, dst: u32, heap: &mut [u8]) -> HypervisorResult<()>;
 
     /// (deprecated) Please use `ic0_msg_cycles_available128` instead.
     /// This API supports only 64-bit values.
diff --git a/rs/system_api/BUILD.bazel b/rs/system_api/BUILD.bazel
index 6ac450a1ee8..22ae00a2686 100644
--- a/rs/system_api/BUILD.bazel
+++ b/rs/system_api/BUILD.bazel
@@ -83,3 +83,12 @@ rust_test(
         "@crate_index//:prometheus",
     ],
 )
+
+rust_test(
+    name = "system_api_tracking_test",
+    srcs = [
+        "tests/system_api_tracking.rs",
+    ],
+    crate_root = "tests/system_api_tracking.rs",
+    deps = INTEGRATION_TEST_DEPS,
+)
diff --git a/rs/system_api/src/lib.rs b/rs/system_api/src/lib.rs
index b135f3b3ad9..0d6cbd8b53c 100644
--- a/rs/system_api/src/lib.rs
+++ b/rs/system_api/src/lib.rs
@@ -12,7 +12,7 @@ use ic_interfaces::execution_environment::{
     ExecutionMode,
     HypervisorError::{self, *},
     HypervisorResult, OutOfInstructionsHandler, PerformanceCounterType, StableGrowOutcome,
-    StableMemoryApi, SubnetAvailableMemory, SystemApi,
+    StableMemoryApi, SubnetAvailableMemory, SystemApi, SystemApiCallCounters,
     TrapCode::{self, CyclesAmountTooBigFor64Bit},
 };
 use ic_logger::{error, ReplicaLogger};
@@ -833,6 +833,9 @@ pub struct SystemApiImpl {
     /// is initialized to 0 and updated after each out-of-instructions call that
     /// starts a new slice.
     instructions_executed_before_current_slice: i64,
+
+    /// How many times each tracked System API call was invoked.
+    call_counters: SystemApiCallCounters,
 }
 
 impl SystemApiImpl {
@@ -874,6 +877,7 @@ impl SystemApiImpl {
             log,
             current_slice_instruction_limit: i64::try_from(slice_limit).unwrap_or(i64::MAX),
             instructions_executed_before_current_slice: 0,
+            call_counters: SystemApiCallCounters::default(),
         }
     }
 
@@ -1324,6 +1328,31 @@ impl SystemApiImpl {
             panic!("{}", WASM_NATIVE_STABLE_MEMORY_ERROR)
         }
     }
+
+    /// Increase `ic0.call_perform()` system API call counter.
+    fn inc_call_perform_counter(&mut self) {
+        self.call_counters.call_perform += 1;
+    }
+
+    /// Increase `ic0.canister_cycle_balance()` system API call counter.
+    fn inc_canister_cycle_balance_counter(&mut self) {
+        self.call_counters.canister_cycle_balance += 1;
+    }
+
+    /// Increase `ic0.canister_cycle_balance128()` system API call counter.
+    fn inc_canister_cycle_balance128_counter(&mut self) {
+        self.call_counters.canister_cycle_balance128 += 1;
+    }
+
+    /// Increase `ic0.time()` system API call counter.
+    fn inc_time_counter(&mut self) {
+        self.call_counters.time += 1;
+    }
+
+    /// Return tracked System API call counters.
+    pub fn call_counters(&self) -> SystemApiCallCounters {
+        self.call_counters.clone()
+    }
 }
 
 impl SystemApi for SystemApiImpl {
@@ -2020,6 +2049,7 @@ impl SystemApi for SystemApiImpl {
     // or the output queues are full. In this case, we need to perform the
     // necessary cleanups.
     fn ic0_call_perform(&mut self) -> HypervisorResult<i32> {
+        self.inc_call_perform_counter();
         let result = match &mut self.api_type {
             ApiType::Start { .. }
             | ApiType::Init { .. }
@@ -2256,7 +2286,8 @@ impl SystemApi for SystemApiImpl {
         Ok((dirty_pages, cost))
     }
 
-    fn ic0_time(&self) -> HypervisorResult<Time> {
+    fn ic0_time(&mut self) -> HypervisorResult<Time> {
+        self.inc_time_counter();
         let result = match &self.api_type {
             ApiType::Start { .. } => Err(self.error_for("ic0_time")),
             ApiType::Init { time, .. }
@@ -2450,7 +2481,8 @@ impl SystemApi for SystemApiImpl {
         }
     }
 
-    fn ic0_canister_cycle_balance(&self) -> HypervisorResult<u64> {
+    fn ic0_canister_cycle_balance(&mut self) -> HypervisorResult<u64> {
+        self.inc_canister_cycle_balance_counter();
         let result = {
             let (high_amount, low_amount) = self
                 .ic0_canister_cycle_balance_helper("ic0_canister_cycle_balance")?
@@ -2464,7 +2496,8 @@ impl SystemApi for SystemApiImpl {
         result
     }
 
-    fn ic0_canister_cycle_balance128(&self, dst: u32, heap: &mut [u8]) -> HypervisorResult<()> {
+    fn ic0_canister_cycle_balance128(&mut self, dst: u32, heap: &mut [u8]) -> HypervisorResult<()> {
+        self.inc_canister_cycle_balance128_counter();
         let result = {
             let method_name = "ic0_canister_cycle_balance128";
             let cycles = self.ic0_canister_cycle_balance_helper(method_name)?;
diff --git a/rs/system_api/tests/system_api.rs b/rs/system_api/tests/system_api.rs
index 33b40e0b41a..e2296fc1cc4 100644
--- a/rs/system_api/tests/system_api.rs
+++ b/rs/system_api/tests/system_api.rs
@@ -1171,7 +1171,7 @@ fn test_canister_balance() {
             Time::from_nanos_since_unix_epoch(0),
         );
 
-    let api = get_system_api(
+    let mut api = get_system_api(
         ApiTypeBuilder::build_update_api(),
         &system_state,
         cycles_account_manager,
@@ -1199,7 +1199,7 @@ fn test_canister_cycle_balance() {
             Time::from_nanos_since_unix_epoch(0),
         );
 
-    let api = get_system_api(
+    let mut api = get_system_api(
         ApiTypeBuilder::build_update_api(),
         &system_state,
         cycles_account_manager,
diff --git a/rs/system_api/tests/system_api_tracking.rs b/rs/system_api/tests/system_api_tracking.rs
new file mode 100644
index 00000000000..76494630002
--- /dev/null
+++ b/rs/system_api/tests/system_api_tracking.rs
@@ -0,0 +1,149 @@
+use ic_interfaces::execution_environment::SystemApiCallCounters;
+use ic_system_api::NonReplicatedQueryKind;
+use ic_test_utilities::types::ids::subnet_test_id;
+use ic_test_utilities::{
+    mock_time, types::ids::user_test_id, wasmtime_instance::WasmtimeInstanceBuilder,
+};
+
+fn call_counters_on_ok_call(wat: &str) -> SystemApiCallCounters {
+    let mut instance = WasmtimeInstanceBuilder::new()
+        .with_wat(wat)
+        .with_api_type(ic_system_api::ApiType::non_replicated_query(
+            mock_time(),
+            user_test_id(0).get(),
+            subnet_test_id(1),
+            vec![0; 1024],
+            None,
+            NonReplicatedQueryKind::Stateful {
+                call_context_id: 0.into(),
+                outgoing_request: None,
+            },
+        ))
+        .build();
+    instance
+        .run(ic_types::methods::FuncRef::Method(
+            ic_types::methods::WasmMethod::CompositeQuery("call_system_api".into()),
+        ))
+        .unwrap();
+    let system_api = &instance.store_data().system_api().unwrap();
+    system_api.call_counters()
+}
+
+fn call_counters_on_err_call(wat: &str) -> SystemApiCallCounters {
+    let mut instance = WasmtimeInstanceBuilder::new()
+        .with_wat(wat)
+        .with_api_type(ic_system_api::ApiType::start(mock_time()))
+        .build();
+    instance
+        .run(ic_types::methods::FuncRef::Method(
+            ic_types::methods::WasmMethod::CompositeQuery("call_system_api".into()),
+        ))
+        .unwrap_err();
+    let system_api = &instance.store_data().system_api().unwrap();
+    system_api.call_counters()
+}
+
+#[test]
+fn track_call_perform() {
+    let wat = r#"(module
+            (import "ic0" "call_new"
+                (func $ic0_call_new
+                (param $callee_src i32)         (param $callee_size i32)
+                (param $name_src i32)           (param $name_size i32)
+                (param $reply_fun i32)          (param $reply_env i32)
+                (param $reject_fun i32)         (param $reject_env i32)
+            ))
+            (import "ic0" "call_perform" (func $ic0_call_perform (result i32)))
+            (memory 1)
+            (func (export "canister_composite_query call_system_api")
+                (call $ic0_call_new
+                    (i32.const 0)   (i32.const 10)
+                    (i32.const 100) (i32.const 18)
+                    (i32.const 11)  (i32.const 0) ;; non-existent function
+                    (i32.const 22)  (i32.const 0) ;; non-existent function
+                )
+                (drop (call $ic0_call_perform))
+            )
+        )"#;
+    let call_counters = call_counters_on_ok_call(wat);
+    assert_eq!(call_counters.call_perform, 1);
+    let wat = r#"(module
+            (import "ic0" "call_perform" (func $ic0_call_perform (result i32)))
+            (memory 1)
+            (func (export "canister_composite_query call_system_api")
+                (drop (call $ic0_call_perform))
+            )
+        )"#;
+    let call_counters = call_counters_on_err_call(wat);
+    assert_eq!(call_counters.call_perform, 1);
+}
+
+#[test]
+fn track_canister_cycle_balance() {
+    let wat = r#"(module
+            (import "ic0" "canister_cycle_balance"
+                (func $ic0_canister_cycle_balance (result i64))
+            )
+            (memory 1)
+            (func (export "canister_composite_query call_system_api")
+                (drop (call $ic0_canister_cycle_balance))
+            )
+        )"#;
+    let call_counters = call_counters_on_ok_call(wat);
+    assert_eq!(call_counters.canister_cycle_balance, 1);
+    let call_counters = call_counters_on_err_call(wat);
+    assert_eq!(call_counters.canister_cycle_balance, 1);
+}
+
+#[test]
+fn track_canister_cycle_balance128() {
+    let wat = r#"(module
+            (import "ic0" "canister_cycle_balance128"
+                (func $ic0_canister_cycle_balance128 (param i32))
+            )
+            (memory 1)
+            (func (export "canister_composite_query call_system_api")
+                (call $ic0_canister_cycle_balance128 (i32.const 0))
+            )
+        )"#;
+    let call_counters = call_counters_on_ok_call(wat);
+    assert_eq!(call_counters.canister_cycle_balance128, 1);
+    let call_counters = call_counters_on_err_call(wat);
+    assert_eq!(call_counters.canister_cycle_balance128, 1);
+}
+
+#[test]
+fn track_time() {
+    let wat = r#"(module
+                (import "ic0" "time" (func $ic0_time (result i64)))
+                (memory 1)
+                (func (export "canister_composite_query call_system_api")
+                    (drop (call $ic0_time))
+                )
+            )"#;
+    let call_counters = call_counters_on_ok_call(wat);
+    assert_eq!(call_counters.time, 1);
+    let call_counters = call_counters_on_err_call(wat);
+    assert_eq!(call_counters.time, 1);
+}
+
+#[test]
+fn track_other() {
+    let wat = r#"(module
+            (import "ic0" "msg_reply" (func $ic0_msg_reply))
+            (memory 1)
+            (func (export "canister_composite_query call_system_api")
+                (call $ic0_msg_reply)
+            )
+        )"#;
+    let call_counters = call_counters_on_ok_call(wat);
+    assert_eq!(call_counters.call_perform, 0);
+    assert_eq!(call_counters.canister_cycle_balance, 0);
+    assert_eq!(call_counters.canister_cycle_balance128, 0);
+    assert_eq!(call_counters.time, 0);
+    let call_counters = call_counters_on_err_call(wat);
+    assert_eq!(call_counters.call_perform, 0);
+    assert_eq!(call_counters.canister_cycle_balance, 0);
+    assert_eq!(call_counters.canister_cycle_balance128, 0);
+    assert_eq!(call_counters.time, 0);
+}