refactor: Consolidate how system state modifications are extracted from the System API (#3706)

dsarlis · web-flow · commit cde707113a91 · 2025-01-31T13:03:37.000Z
Currently, there are two similar functions in the `SystemApiImpl` to extract system state modifications. One of them `take_system_state_modifications` is used in the sandbox when preparing the state changes to be transmitted back to the replica. Then, the replica after deserializing the result that it receives from the sandbox calls `into_system_state_changes` on the reconstructed `system_api` to extract the changes again. In a recent [PR](#363), `into_system_state_modifications` was changed (to make it more clear which changes are relevant per message type) but `take_system_state_modifications` wasn't. This works correctly because `into_system_state_modifications` is the last one that's called before applying the state changes back to the canister state. However, it's also a very clear example of how the code can easily diverge and go unnoticed, potentially with more severe implications in the future. This PR proposes to use one function which seems to provide the usual benefits of consistent way of applying the changes and "having one way" of doing the same task. We keep `take_system_state_modifications` (this allows us to get rid of some `clone`s but more importantly it's needed in the sandbox, see comment in the existing code) and change the call sites respectively.
diff --git a/rs/canister_sandbox/src/sandbox_manager.rs b/rs/canister_sandbox/src/sandbox_manager.rs
@@ -173,7 +173,7 @@ impl Execution {
                             .system_api_mut()
                             .expect("System api not present in the wasmtime instance")
                             .take_system_state_modifications(),
-                        Err(system_api) => system_api.into_system_state_modifications(),
+                        Err(mut system_api) => system_api.take_system_state_modifications(),
                     };
 
                     let execution_state_modifications = deltas.map(
diff --git a/rs/embedders/src/wasm_executor.rs b/rs/embedders/src/wasm_executor.rs
@@ -258,11 +258,11 @@ impl WasmExecutor for WasmExecutorImpl {
             }),
             None => None,
         };
-        let system_api = match instance_or_system_api {
+        let mut system_api = match instance_or_system_api {
             Ok(instance) => instance.into_store_data().system_api.unwrap(),
             Err(system_api) => system_api,
         };
-        let system_state_modifications = system_api.into_system_state_modifications();
+        let system_state_modifications = system_api.take_system_state_modifications();
 
         (
             compilation_result,
diff --git a/rs/system_api/src/lib.rs b/rs/system_api/src/lib.rs
@@ -1494,7 +1494,8 @@ impl SystemApiImpl {
         }
     }
 
-    pub fn into_system_state_modifications(self) -> SystemStateModifications {
+    pub fn take_system_state_modifications(&mut self) -> SystemStateModifications {
+        let system_state_modifications = self.sandbox_safe_system_state.take_changes();
         match self.api_type {
             // List all fields of `SystemStateModifications` so that
             // there's an explicit decision that needs to be made
@@ -1514,53 +1515,28 @@ impl SystemApiImpl {
             },
             ApiType::NonReplicatedQuery { .. } => SystemStateModifications {
                 new_certified_data: None,
-                callback_updates: self
-                    .sandbox_safe_system_state
-                    .system_state_modifications
-                    .callback_updates
-                    .clone(),
+                callback_updates: system_state_modifications.callback_updates,
                 cycles_balance_change: CyclesBalanceChange::zero(),
                 reserved_cycles: Cycles::zero(),
                 consumed_cycles_by_use_case: BTreeMap::new(),
                 call_context_balance_taken: None,
-                request_slots_used: self
-                    .sandbox_safe_system_state
-                    .system_state_modifications
-                    .request_slots_used
-                    .clone(),
-                requests: self
-                    .sandbox_safe_system_state
-                    .system_state_modifications
-                    .requests
-                    .clone(),
+                request_slots_used: system_state_modifications.request_slots_used,
+                requests: system_state_modifications.requests,
                 new_global_timer: None,
                 canister_log: Default::default(),
                 on_low_wasm_memory_hook_condition_check_result: None,
             },
             ApiType::ReplicatedQuery { .. } => SystemStateModifications {
                 new_certified_data: None,
                 callback_updates: vec![],
-                cycles_balance_change: self
-                    .sandbox_safe_system_state
-                    .system_state_modifications
-                    .cycles_balance_change,
+                cycles_balance_change: system_state_modifications.cycles_balance_change,
                 reserved_cycles: Cycles::zero(),
-                consumed_cycles_by_use_case: self
-                    .sandbox_safe_system_state
-                    .system_state_modifications
-                    .consumed_cycles_by_use_case,
-                call_context_balance_taken: self
-                    .sandbox_safe_system_state
-                    .system_state_modifications
-                    .call_context_balance_taken,
+                consumed_cycles_by_use_case: system_state_modifications.consumed_cycles_by_use_case,
+                call_context_balance_taken: system_state_modifications.call_context_balance_taken,
                 request_slots_used: BTreeMap::new(),
                 requests: vec![],
                 new_global_timer: None,
-                canister_log: self
-                    .sandbox_safe_system_state
-                    .system_state_modifications
-                    .canister_log
-                    .clone(),
+                canister_log: system_state_modifications.canister_log,
                 on_low_wasm_memory_hook_condition_check_result: None,
             },
             ApiType::Start { .. }
@@ -1570,16 +1546,10 @@ impl SystemApiImpl {
             | ApiType::Update { .. }
             | ApiType::Cleanup { .. }
             | ApiType::ReplyCallback { .. }
-            | ApiType::RejectCallback { .. } => {
-                self.sandbox_safe_system_state.system_state_modifications
-            }
+            | ApiType::RejectCallback { .. } => system_state_modifications,
         }
     }
 
-    pub fn take_system_state_modifications(&mut self) -> SystemStateModifications {
-        self.sandbox_safe_system_state.take_changes()
-    }
-
     pub fn stable_memory_size(&self) -> NumWasmPages {
         self.stable_memory().stable_memory_size
     }
diff --git a/rs/system_api/tests/sandbox_safe_system_state.rs b/rs/system_api/tests/sandbox_safe_system_state.rs
@@ -443,7 +443,7 @@ fn call_increases_cycles_consumed_metric() {
     api.ic0_call_new(0, 0, 0, 0, 0, 0, 0, 0, &[]).unwrap();
     api.ic0_call_perform().unwrap();
 
-    let system_state_modifications = api.into_system_state_modifications();
+    let system_state_modifications = api.take_system_state_modifications();
     system_state_modifications
         .apply_changes(
             UNIX_EPOCH,
diff --git a/rs/system_api/tests/system_api.rs b/rs/system_api/tests/system_api.rs
@@ -1164,7 +1164,7 @@ fn certified_data_set() {
     // Copy the certified data into the system state.
     api.ic0_certified_data_set(0, 32, &heap).unwrap();
 
-    let system_state_modifications = api.into_system_state_modifications();
+    let system_state_modifications = api.take_system_state_modifications();
     system_state_modifications
         .apply_changes(
             UNIX_EPOCH,
@@ -1337,7 +1337,7 @@ fn call_perform_not_enough_cycles_does_not_trap() {
             res
         ),
     }
-    let system_state_modifications = api.into_system_state_modifications();
+    let system_state_modifications = api.take_system_state_modifications();
     system_state_modifications
         .apply_changes(
             UNIX_EPOCH,
@@ -1481,7 +1481,7 @@ fn helper_test_on_low_wasm_memory(
             .unwrap();
     }
 
-    let system_state_modifications = api.into_system_state_modifications();
+    let system_state_modifications = api.take_system_state_modifications();
     system_state_modifications
         .apply_changes(
             UNIX_EPOCH,
@@ -1750,7 +1750,7 @@ fn push_output_request_respects_memory_limits() {
     );
 
     // Ensure that exactly one output request was pushed.
-    let system_state_modifications = api.into_system_state_modifications();
+    let system_state_modifications = api.take_system_state_modifications();
     system_state_modifications
         .apply_changes(
             UNIX_EPOCH,
@@ -1866,7 +1866,7 @@ fn push_output_request_oversized_request_memory_limits() {
     );
 
     // Ensure that exactly one output request was pushed.
-    let system_state_modifications = api.into_system_state_modifications();
+    let system_state_modifications = api.take_system_state_modifications();
     system_state_modifications
         .apply_changes(
             UNIX_EPOCH,
@@ -1902,7 +1902,7 @@ fn ic0_global_timer_set_is_propagated_from_sandbox() {
 
     // Propagate system state changes
     assert_eq!(system_state.global_timer, CanisterTimer::Inactive);
-    let system_state_modifications = api.into_system_state_modifications();
+    let system_state_modifications = api.take_system_state_modifications();
     system_state_modifications
         .apply_changes(
             UNIX_EPOCH,
