feat: Allow accepting and burning cycles in replicated queries

hs/spec_compliance/src/IC/Test/Spec/CanisterVersion.hs

@@ -44,9 +44,9 @@ canister_version_tests ecid =
                               ctr <- callToQuery'' cid (replyData canister_version) >>= is2xx >>= isReply >>= asWord64
                               ctr @?= 1
                               ctr <- callToQuery'' cid (replyData canister_version) >>= is2xx >>= isReply >>= asWord64
-                              ctr @?= 1
+                              ctr @?= 2
                               ctr <- callToQuery'' cid (replyData canister_version) >>= is2xx >>= isReply >>= asWord64
-                              ctr @?= 1,
+                              ctr @?= 3,
                             simpleTestCase "in update" ecid $ \cid -> do
                               ctr <- call cid (replyData canister_version) >>= asWord64
                               ctr @?= 1

rs/canister_sandbox/src/protocol/ctlsvc.rs

@@ -74,7 +74,9 @@ mod tests {
         ctlsvc::{ExecutionFinishedReply, ExecutionPausedReply, ExecutionPausedRequest, Reply},
         id::ExecId,
         logging::{LogLevel, LogRequest},
-        structs::{MemoryModifications, SandboxExecOutput, StateModifications},
+        structs::{
+            ExecutionStateModifications, MemoryModifications, SandboxExecOutput, StateModifications,
+        },
     };
 
     use super::{ExecutionFinishedRequest, Request};
@@ -105,24 +107,26 @@ mod tests {
                 system_api_call_counters: SystemApiCallCounters::default(),
                 canister_log: CanisterLog::default(),
             },
-            state: Some(StateModifications {
-                globals: vec![
-                    Global::I32(10),
-                    Global::I64(32),
-                    Global::F32(10.5),
-                    Global::F64(1.1),
-                    Global::V128(123),
-                ],
-                wasm_memory: MemoryModifications {
-                    page_delta: PageMap::new_for_testing().serialize_delta(&[]),
-                    size: NumWasmPages::new(10),
-                },
-                stable_memory: MemoryModifications {
-                    page_delta: PageMap::new_for_testing().serialize_delta(&[]),
-                    size: NumWasmPages::new(42),
-                },
+            state: StateModifications {
+                execution_state_modifications: Some(ExecutionStateModifications {
+                    globals: vec![
+                        Global::I32(10),
+                        Global::I64(32),
+                        Global::F32(10.5),
+                        Global::F64(1.1),
+                        Global::V128(123),
+                    ],
+                    wasm_memory: MemoryModifications {
+                        page_delta: PageMap::new_for_testing().serialize_delta(&[]),
+                        size: NumWasmPages::new(10),
+                    },
+                    stable_memory: MemoryModifications {
+                        page_delta: PageMap::new_for_testing().serialize_delta(&[]),
+                        size: NumWasmPages::new(42),
+                    },
+                }),
                 system_state_changes: SystemStateChanges::default(),
-            }),
+            },
             execute_total_duration: Duration::from_secs(10),
             execute_run_duration: Duration::from_secs(1),
         };

rs/canister_sandbox/src/protocol/structs.rs

@@ -36,20 +36,42 @@ pub struct SandboxExecInput {
 pub struct SandboxExecOutput {
     pub slice: SliceExecutionOutput,
     pub wasm: WasmExecutionOutput,
-    pub state: Option<StateModifications>,
+    pub state: StateModifications,
     pub execute_total_duration: std::time::Duration,
     pub execute_run_duration: std::time::Duration,
 }
 
+impl SandboxExecOutput {
+    pub fn take_state_modifications(&mut self) -> StateModifications {
+        std::mem::take(&mut self.state)
+    }
+}
+
 /// Describes the memory changes performed by execution.
 #[derive(Clone, PartialEq, Debug, Deserialize, Serialize)]
 pub struct MemoryModifications {
     pub page_delta: PageDeltaSerialization,
     pub size: NumWasmPages,
 }
 
-#[derive(Clone, PartialEq, Debug, Deserialize, Serialize)]
+#[derive(Serialize, Default, Debug, Deserialize, Clone, PartialEq)]
 pub struct StateModifications {
+    /// Modifications in the execution state of the canister.
+    ///
+    /// This field is optional because the state changes might or might not
+    /// be applied depending on the method executed.
+    pub execution_state_modifications: Option<ExecutionStateModifications>,
+
+    /// Modifications in the system state of the canister.
+    ///
+    /// The system state changes contain parts that are always applied
+    /// and parts that are only applied depending on the method executed
+    /// (similarly to `execution_state_modifications`).
+    pub system_state_changes: SystemStateChanges,
+}
+
+#[derive(Serialize, Debug, Deserialize, Clone, PartialEq)]
+pub struct ExecutionStateModifications {
     /// The state of the global variables after execution.
     pub globals: Vec<Global>,
 
@@ -58,19 +80,15 @@ pub struct StateModifications {
 
     /// Modifications in the stable memory.
     pub stable_memory: MemoryModifications,
-
-    /// Modifications in the system state.
-    pub system_state_changes: SystemStateChanges,
 }
 
-impl StateModifications {
+impl ExecutionStateModifications {
     pub fn new(
         globals: Vec<Global>,
         wasm_memory: &Memory,
         stable_memory: &Memory,
         wasm_memory_delta: &[PageIndex],
         stable_memory_delta: &[PageIndex],
-        system_state_changes: SystemStateChanges,
     ) -> Self {
         let wasm_memory = MemoryModifications {
             page_delta: wasm_memory.page_map.serialize_delta(wasm_memory_delta),
@@ -82,11 +100,10 @@ impl StateModifications {
             size: stable_memory.size,
         };
 
-        StateModifications {
+        ExecutionStateModifications {
             globals,
             wasm_memory,
             stable_memory,
-            system_state_changes,
         }
     }
 }

rs/canister_sandbox/src/replica_controller/sandboxed_execution_controller.rs

@@ -3,14 +3,14 @@ use crate::controller_launcher_service::ControllerLauncherService;
 use crate::launcher_service::LauncherService;
 use crate::protocol::id::{ExecId, MemoryId, WasmId};
 use crate::protocol::sbxsvc::MemorySerialization;
-use crate::protocol::structs::{SandboxExecInput, SandboxExecOutput};
+use crate::protocol::structs::{SandboxExecInput, SandboxExecOutput, StateModifications};
 use crate::sandbox_service::SandboxService;
 use crate::{protocol, rpc};
 use ic_config::embedders::Config as EmbeddersConfig;
 use ic_config::flag_status::FlagStatus;
 use ic_embedders::wasm_executor::{
-    get_wasm_reserved_pages, wasm_execution_error, CanisterStateChanges, PausedWasmExecution,
-    SliceExecutionOutput, WasmExecutionResult, WasmExecutor,
+    get_wasm_reserved_pages, wasm_execution_error, CanisterStateChanges, ExecutionStateChanges,
+    PausedWasmExecution, SliceExecutionOutput, WasmExecutionResult, WasmExecutor,
 };
 use ic_embedders::{
     wasm_utils::WasmImportsDetails, CompilationCache, CompilationResult, StoredCompilation,
@@ -1575,22 +1575,31 @@ impl SandboxedExecutionController {
         next_stable_memory_id: MemoryId,
         canister_id: CanisterId,
         sandbox_process: Arc<SandboxProcess>,
-    ) -> Option<CanisterStateChanges> {
+    ) -> CanisterStateChanges {
         // If the execution has failed, then we don't apply any changes.
         if exec_output.wasm.wasm_result.is_err() {
-            return None;
+            return CanisterStateChanges::default();
         }
-        match exec_output.state.take() {
-            None => None,
-            Some(state_modifications) => {
+
+        let StateModifications {
+            execution_state_modifications,
+            system_state_changes,
+        } = exec_output.take_state_modifications();
+
+        match execution_state_modifications {
+            None => CanisterStateChanges {
+                execution_state_changes: None,
+                system_state_changes,
+            },
+            Some(execution_state_modifications) => {
                 // TODO: If a canister has broken out of wasm then it might have allocated more
                 // wasm or stable memory then allowed. We should add an additional check here
                 // that thet canister is still within it's allowed memory usage.
                 let mut wasm_memory = execution_state.wasm_memory.clone();
                 wasm_memory
                     .page_map
-                    .deserialize_delta(state_modifications.wasm_memory.page_delta);
-                wasm_memory.size = state_modifications.wasm_memory.size;
+                    .deserialize_delta(execution_state_modifications.wasm_memory.page_delta);
+                wasm_memory.size = execution_state_modifications.wasm_memory.size;
                 wasm_memory.sandbox_memory = SandboxMemory::synced(wrap_remote_memory(
                     &sandbox_process,
                     next_wasm_memory_id,
@@ -1610,8 +1619,8 @@ impl SandboxedExecutionController {
                 let mut stable_memory = execution_state.stable_memory.clone();
                 stable_memory
                     .page_map
-                    .deserialize_delta(state_modifications.stable_memory.page_delta);
-                stable_memory.size = state_modifications.stable_memory.size;
+                    .deserialize_delta(execution_state_modifications.stable_memory.page_delta);
+                stable_memory.size = execution_state_modifications.stable_memory.size;
                 stable_memory.sandbox_memory = SandboxMemory::synced(wrap_remote_memory(
                     &sandbox_process,
                     next_stable_memory_id,
@@ -1628,12 +1637,14 @@ impl SandboxedExecutionController {
                         .sandboxed_execution_critical_error_invalid_memory_size
                         .inc();
                 }
-                Some(CanisterStateChanges {
-                    globals: state_modifications.globals,
-                    wasm_memory,
-                    stable_memory,
-                    system_state_changes: state_modifications.system_state_changes,
-                })
+                CanisterStateChanges {
+                    execution_state_changes: Some(ExecutionStateChanges {
+                        globals: execution_state_modifications.globals,
+                        wasm_memory,
+                        stable_memory,
+                    }),
+                    system_state_changes,
+                }
             }
         }
     }

rs/canister_sandbox/src/sandbox_manager.rs

@@ -21,7 +21,8 @@ use std::time::{Duration, Instant};
 use crate::protocol::id::{ExecId, MemoryId, WasmId};
 use crate::protocol::sbxsvc::{CreateExecutionStateSerializedSuccessReply, OpenMemoryRequest};
 use crate::protocol::structs::{
-    MemoryModifications, SandboxExecInput, SandboxExecOutput, StateModifications,
+    ExecutionStateModifications, MemoryModifications, SandboxExecInput, SandboxExecOutput,
+    StateModifications,
 };
 use crate::{controller_service::ControllerService, protocol};
 use ic_config::embedders::Config as EmbeddersConfig;
@@ -158,37 +159,44 @@ impl Execution {
 
         match wasm_result {
             Ok(_) => {
-                let state_modifications = deltas.map(
-                    |WasmStateChanges {
-                         dirty_page_indices,
-                         globals,
-                     }| {
-                        let system_state_changes = match instance_or_system_api {
-                            // Here we use `store_data_mut` instead of
-                            // `into_store_data` because the later will drop the
-                            // wasmtime Instance which can be an expensive
-                            // operation. Mutating the store instead allows us
-                            // to delay the drop until after the execution
-                            // completed message is sent back to the main
-                            // process.
-                            Ok(mut instance) => instance
-                                .store_data_mut()
-                                .system_api_mut()
-                                .expect("System api not present in the wasmtime instance")
-                                .take_system_state_changes(),
-                            Err(system_api) => system_api.into_system_state_changes(),
-                        };
-                        StateModifications::new(
-                            globals,
-                            &wasm_memory,
-                            &stable_memory,
-                            &dirty_page_indices.wasm_memory_delta,
-                            &dirty_page_indices.stable_memory_delta,
-                            system_state_changes,
-                        )
-                    },
-                );
-                if state_modifications.is_some() {
+                let state_modifications = {
+                    let system_state_changes = match instance_or_system_api {
+                        // Here we use `store_data_mut` instead of
+                        // `into_store_data` because the later will drop the
+                        // wasmtime Instance which can be an expensive
+                        // operation. Mutating the store instead allows us
+                        // to delay the drop until after the execution
+                        // completed message is sent back to the main
+                        // process.
+                        Ok(mut instance) => instance
+                            .store_data_mut()
+                            .system_api_mut()
+                            .expect("System api not present in the wasmtime instance")
+                            .take_system_state_changes(),
+                        Err(system_api) => system_api.into_system_state_changes(),
+                    };
+
+                    let execution_state_modifications = deltas.map(
+                        |WasmStateChanges {
+                             dirty_page_indices,
+                             globals,
+                         }| {
+                            ExecutionStateModifications::new(
+                                globals,
+                                &wasm_memory,
+                                &stable_memory,
+                                &dirty_page_indices.wasm_memory_delta,
+                                &dirty_page_indices.stable_memory_delta,
+                            )
+                        },
+                    );
+
+                    StateModifications {
+                        execution_state_modifications,
+                        system_state_changes,
+                    }
+                };
+                if state_modifications.execution_state_modifications.is_some() {
                     self.sandbox_manager
                         .add_memory(exec_input.next_wasm_memory_id, wasm_memory);
                     self.sandbox_manager
@@ -237,7 +245,7 @@ impl Execution {
                         exec_output: SandboxExecOutput {
                             slice,
                             wasm: wasm_output,
-                            state: None,
+                            state: StateModifications::default(),
                             execute_total_duration: total_timer.elapsed(),
                             execute_run_duration: run_timer.elapsed(),
                         },