feat: Add detailed memory usage metrics to CanisterStatusResultV2

q-uint · q-uint · commit 0907ae488ac1 · 2025-02-25T10:27:30.000Z
- Introduced new memory metrics fields in `CanisterStatusResultV2`, including: - `wasm_memory_size` - `stable_memory_size` - `global_memory_size` - `wasm_binary_size` - `custom_sections_size` - `canister_history_size` - `wasm_chunk_store_size` - `snapshots_size` - Updated `CanisterManager` and `CanisterState` to compute and track these new metrics. - Refactored `ExecutionState` to provide individual methods for retrieving specific memory usage values. - Modified tests to validate the new memory metrics in canister status responses. This change improves visibility into canister memory consumption, aiding debugging and resource management. Refer to [PR portal#5240](dfinity/portal#5240) for more details.
diff --git a/rs/execution_environment/src/canister_manager.rs b/rs/execution_environment/src/canister_manager.rs
@@ -1096,6 +1096,14 @@ impl CanisterManager {
             .collect::<Vec<PrincipalId>>();
 
         let canister_memory_usage = canister.memory_usage();
+        let canister_wasm_memory_usage = canister.wasm_memory_usage();
+        let canister_stable_memory_usage = canister.stable_memory_usage();
+        let canister_global_memory_usage = canister.global_memory_usage();
+        let canister_wasm_binary_memory_usage = canister.wasm_binary_memory_usage();
+        let canister_custom_sections_memory_usage = canister.wasm_custom_sections_memory_usage();
+        let canister_history_memory_usage = canister.canister_history_memory_usage();
+        let canister_wasm_chunk_store_memory_usage = canister.wasm_chunk_store_memory_usage();
+        let canister_snapshots_memory_usage = canister.snapshots_memory_usage();
         let canister_message_memory_usage = canister.message_memory_usage();
         let compute_allocation = canister.scheduler_state.compute_allocation;
         let memory_allocation = canister.memory_allocation();
@@ -1114,6 +1122,14 @@ impl CanisterManager {
             *controller,
             controllers,
             canister_memory_usage,
+            canister_wasm_memory_usage,
+            canister_stable_memory_usage,
+            canister_global_memory_usage,
+            canister_wasm_binary_memory_usage,
+            canister_custom_sections_memory_usage,
+            canister_history_memory_usage,
+            canister_wasm_chunk_store_memory_usage,
+            canister_snapshots_memory_usage,
             canister.system_state.balance().get(),
             compute_allocation.as_percent(),
             Some(memory_allocation.bytes().get()),
diff --git a/rs/execution_environment/src/execution_environment/tests.rs b/rs/execution_environment/src/execution_environment/tests.rs
@@ -35,12 +35,13 @@ use ic_types::{
     },
     nominal_cycles::NominalCycles,
     time::UNIX_EPOCH,
-    CanisterId, Cycles, PrincipalId, RegistryVersion,
+    CanisterId, Cycles, PrincipalId, RegistryVersion, CountBytes
 };
 use ic_types_test_utils::ids::{canister_test_id, node_test_id, subnet_test_id, user_test_id};
 use ic_universal_canister::{call_args, wasm, UNIVERSAL_CANISTER_WASM};
 use maplit::btreemap;
 use std::mem::size_of;
+use more_asserts::assert_gt;
 
 #[cfg(test)]
 mod canister_task;
@@ -885,6 +886,176 @@ fn get_canister_status_of_nonexisting_canister() {
     assert_eq!(ErrorCode::CanisterNotFound, err.code());
 }
 
+fn get_canister_status(test: &mut ExecutionTest, canister_id: CanisterId) -> CanisterStatusResultV2 {
+    match test.canister_status(canister_id).unwrap() {
+        WasmResult::Reply(reply) => CanisterStatusResultV2::decode(&reply).unwrap(),
+        WasmResult::Reject(msg) => panic!("{}", msg) ,
+    }
+}
+
+#[test]
+fn get_canister_status_memory_metrics() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.universal_canister().unwrap();
+    let csr: CanisterStatusResultV2 = get_canister_status(&mut test, canister_id);
+    
+    let wasm_memory_size = csr.wasm_memory_size();
+    let stable_memory_size = csr.stable_memory_size();
+    let global_memory_size = csr.global_memory_size();
+    let wasm_binary_size = csr.wasm_binary_size();
+    let custom_sections_size = csr.custom_sections_size();
+    
+    let execution_memory_size = wasm_memory_size + stable_memory_size + global_memory_size + wasm_binary_size + custom_sections_size;
+
+    let canister_history_size = csr.canister_history_size();
+    let wasm_chunk_store_size = csr.wasm_chunk_store_size();
+    let snapshots_size = csr.snapshots_size();
+
+    let system_memory_size = canister_history_size + wasm_chunk_store_size + snapshots_size;
+
+    let memory_size = csr.memory_size();
+    assert_eq!(memory_size, execution_memory_size + system_memory_size);
+}
+
+#[test]
+fn get_canister_status_memory_metrics_wasm_memory_size() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.universal_canister().unwrap();
+    let csr: CanisterStatusResultV2 = get_canister_status(&mut test, canister_id);
+
+    assert_eq!(
+        get_canister_status(&mut test, canister_id).wasm_memory_size(), 
+        csr.wasm_memory_size()
+    );
+
+    let canister_status_args = Encode!(&CanisterIdRecord::from(canister_id)).unwrap();
+    let get_canister_status = wasm()
+        .call_simple(
+            ic00::IC_00,
+            Method::CanisterStatus,
+            call_args().other_side(canister_status_args),
+        )
+        .build();
+    let result = test.ingress(canister_id, "update", get_canister_status);
+    let reply = get_reply(result);
+    let updated_csr = CanisterStatusResultV2::decode(&reply).unwrap();
+
+    assert_gt!(
+        updated_csr.wasm_memory_size(), 
+        csr.wasm_memory_size()
+    );
+}
+
+#[test]
+fn get_canister_status_memory_metrics_stable_memory_size() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.universal_canister().unwrap();
+    let csr: CanisterStatusResultV2 = get_canister_status(&mut test, canister_id);
+    test.ingress(
+        canister_id,
+        "update",
+        wasm()
+            .stable64_grow(1)
+            .stable64_read(WASM_PAGE_SIZE_IN_BYTES as u64 - 1, 1)
+            .push_bytes(&[])
+            .append_and_reply()
+            .build(),
+    ).unwrap();
+    assert_eq!(
+        get_canister_status(&mut test, canister_id).stable_memory_size(),
+        csr.stable_memory_size() + NumBytes::from(WASM_PAGE_SIZE_IN_BYTES as u64)
+    );
+}
+
+#[test]
+fn get_canister_status_memory_metrics_global_memory_size() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.universal_canister().unwrap();
+    let csr: CanisterStatusResultV2 = get_canister_status(&mut test, canister_id);
+    let exported_globals = test.execution_state(canister_id).exported_globals.clone();
+    assert_eq!(
+        csr.global_memory_size(), 
+        NumBytes::new(8 * exported_globals.len() as u64)
+    );
+}
+
+#[test]
+fn get_canister_status_memory_metrics_wasm_binary_size() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.universal_canister().unwrap();
+    let csr: CanisterStatusResultV2 = get_canister_status(&mut test, canister_id);
+    assert_eq!(
+        csr.wasm_binary_size(),
+        NumBytes::new(UNIVERSAL_CANISTER_WASM.len() as u64)
+    );
+}
+
+#[test]
+fn get_canister_status_memory_metrics_custom_sections_size() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.universal_canister().unwrap();
+    let csr: CanisterStatusResultV2 = get_canister_status(&mut test, canister_id);
+    let metadata = test.execution_state(canister_id).metadata.clone();
+    assert_eq!(
+        csr.custom_sections_size(),
+        NumBytes::new(
+            metadata.custom_sections()
+                .iter()
+                .map(|(k, v)| k.len() + v.count_bytes())
+                .sum::<usize>() as u64
+        ),
+    );
+}
+
+#[test]
+fn get_canister_status_memory_metrics_canister_history_size() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.universal_canister().unwrap();
+    let csr: CanisterStatusResultV2 = get_canister_status(&mut test, canister_id);
+    test.set_controller(canister_id, test.user_id().get()).unwrap();
+    let memory_difference = NumBytes::from((size_of::<CanisterChange>() + size_of::<PrincipalId>()) as u64);
+    assert_eq!(
+        get_canister_status(&mut test, canister_id).canister_history_size(), 
+        csr.canister_history_size() + memory_difference
+    );
+}
+
+#[test]
+fn get_canister_status_memory_metrics_wasm_chunk_store_size() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.universal_canister().unwrap();
+    let csr: CanisterStatusResultV2 = get_canister_status(&mut test, canister_id);
+    test.subnet_message(
+        "upload_chunk",
+        UploadChunkArgs {
+            canister_id: canister_id.into(),
+            chunk: vec![1, 2, 3, 4, 5],
+        }.encode(),
+    ).unwrap();
+    assert_gt!(
+        get_canister_status(&mut test, canister_id).wasm_chunk_store_size(),
+        csr.wasm_chunk_store_size()
+    );
+}
+
+#[test]
+fn get_canister_status_memory_metrics_snapshots_size() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let canister_id = test.universal_canister().unwrap();
+    let csr: CanisterStatusResultV2 = get_canister_status(&mut test, canister_id);
+    test.subnet_message(
+        "take_canister_snapshot",
+        TakeCanisterSnapshotArgs{
+            canister_id: canister_id.into(),
+            replace_snapshot: None,
+        }.encode()
+    ).unwrap();
+    assert_gt!(
+        get_canister_status(&mut test, canister_id).snapshots_size(), 
+        csr.snapshots_size()
+    );
+}
+
 #[test]
 fn deposit_cycles_to_non_existing_canister_fails() {
     let mut test = ExecutionTestBuilder::new().build();
diff --git a/rs/replica_tests/tests/canister_lifecycle.rs b/rs/replica_tests/tests/canister_lifecycle.rs
@@ -689,6 +689,7 @@ fn can_get_canister_information() {
             Ok(WasmResult::Reply(EmptyBlob.encode()))
         );
 
+        let canister_history_size = NumBytes::from((2 * size_of::<CanisterChange>() + 2 * size_of::<PrincipalId>()) as u64);
         // Request the status of canister_b.
         assert_matches!(
             test.ingress(
@@ -707,7 +708,15 @@ fn can_get_canister_information() {
                 None,
                 canister_a.get(),
                 vec![canister_a.get()],
-                NumBytes::from((2 * size_of::<CanisterChange>() + 2 * size_of::<PrincipalId>()) as u64),
+                canister_history_size,
+                NumBytes::from(0),
+                NumBytes::from(0),
+                NumBytes::from(0),
+                NumBytes::from(0),
+                NumBytes::from(0),
+                canister_history_size,
+                NumBytes::from(0),
+                NumBytes::from(0),
                 num_cycles.get(),
                 ComputeAllocation::default().as_percent(),
                 None,
@@ -767,6 +776,14 @@ fn can_get_canister_information() {
                     // We don't assert a specific memory size since the universal canister's
                     // size changes between updates.
                     NumBytes::from(0),
+                    NumBytes::from(0),
+                    NumBytes::from(0),
+                    NumBytes::from(0),
+                    NumBytes::from(0),
+                    NumBytes::from(0),
+                    NumBytes::from(0),
+                    NumBytes::from(0),
+                    NumBytes::from(0),
                     num_cycles.get(),
                     ComputeAllocation::default().as_percent(),
                     None,
diff --git a/rs/replicated_state/src/canister_state.rs b/rs/replicated_state/src/canister_state.rs
@@ -371,7 +371,7 @@ impl CanisterState {
         self.execution_memory_usage()
             + self.canister_history_memory_usage()
             + self.wasm_chunk_store_memory_usage()
-            + self.system_state.snapshots_memory_usage
+            + self.snapshots_memory_usage()
     }
 
     /// Returns the amount of Wasm memory currently used by the canister in bytes.
@@ -381,6 +381,27 @@ impl CanisterState {
             .map_or(NumBytes::new(0), |es| es.wasm_memory_usage())
     }
 
+    /// Returns the amount of stable memory currently used by the canister in bytes.
+    pub fn stable_memory_usage(&self) -> NumBytes {
+        self.execution_state
+            .as_ref()
+            .map_or(NumBytes::from(0), |es| es.stable_memory_usage())
+    }
+
+    /// Returns the amount of memory currently used by global variables of the canister in bytes.
+    pub fn global_memory_usage(&self) -> NumBytes {
+        self.execution_state
+            .as_ref()
+            .map_or(NumBytes::from(0), |es| es.global_memory_usage())
+    }
+
+    /// Returns the amount of memory currently used by the wasm binary.
+    pub fn wasm_binary_memory_usage(&self) -> NumBytes {
+        self.execution_state
+            .as_ref()
+            .map_or(NumBytes::from(0), |es| es.wasm_binary_memory_usage())
+    }
+
     /// Returns the amount of execution memory (heap, stable, globals, Wasm)
     /// currently used by the canister in bytes.
     pub fn execution_memory_usage(&self) -> NumBytes {
@@ -412,10 +433,14 @@ impl CanisterState {
     }
 
     /// Returns the memory usage of the wasm chunk store in bytes.
-    pub(super) fn wasm_chunk_store_memory_usage(&self) -> NumBytes {
+    pub fn wasm_chunk_store_memory_usage(&self) -> NumBytes {
         self.system_state.wasm_chunk_store.memory_usage()
     }
 
+    pub fn snapshots_memory_usage(&self) -> NumBytes {
+        self.system_state.snapshots_memory_usage
+    }
+
     /// Returns the snapshot size estimation in bytes based on the current canister's state.
     ///
     /// It represents the memory usage of a snapshot that would be created at the time of the call
diff --git a/rs/replicated_state/src/canister_state/execution_state.rs b/rs/replicated_state/src/canister_state/execution_state.rs
@@ -567,17 +567,37 @@ impl ExecutionState {
             .expect("could not convert from wasm memory number of pages to bytes")
     }
 
-    /// Returns the memory currently used by the `ExecutionState`.
-    pub fn memory_usage(&self) -> NumBytes {
+    /// Returns the stable memory currently used by the `ExecutionState`.
+    pub fn stable_memory_usage(&self) -> NumBytes {
+        num_bytes_try_from(self.stable_memory.size)
+            .expect("could not convert from stable memory number of pages to bytes")
+    }
+
+    // Returns the global memory currently used by the `ExecutionState`.
+    pub fn global_memory_usage(&self) -> NumBytes {
         // We use 8 bytes per global.
         let globals_size_bytes = 8 * self.exported_globals.len() as u64;
+        NumBytes::from(globals_size_bytes)
+    }
+
+    // Returns the memory size of the Wasm binary currently used by the `ExecutionState`.
+    pub fn wasm_binary_memory_usage(&self) -> NumBytes {
         let wasm_binary_size_bytes = self.wasm_binary.binary.len() as u64;
+        NumBytes::from(wasm_binary_size_bytes)
+    }
+
+    // Returns the memory size of the custom sections currently used by the `ExecutionState`.
+    pub fn custom_sections_memory_size(&self) -> NumBytes {
+        self.metadata.memory_usage()
+    } 
+
+    /// Returns the memory currently used by the `ExecutionState`.
+    pub fn memory_usage(&self) -> NumBytes {
         self.wasm_memory_usage()
-            + num_bytes_try_from(self.stable_memory.size)
-                .expect("could not convert from stable memory number of pages to bytes")
-            + NumBytes::from(globals_size_bytes)
-            + NumBytes::from(wasm_binary_size_bytes)
-            + self.metadata.memory_usage()
+            + self.stable_memory_usage()
+            + self.global_memory_usage()
+            + self.wasm_binary_memory_usage()
+            + self.custom_sections_memory_size()
     }
 
     /// Returns the number of global variables in the Wasm module.
diff --git a/rs/types/management_canister_types/src/lib.rs b/rs/types/management_canister_types/src/lib.rs
