refactor: Consolidate how system state modifications are extracted fr…

…om the System API (#3706)

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.

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(

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,

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
     }

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,

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,