fix(execution): Canister snapshots should include Wasm globals

rs/execution_environment/src/canister_manager.rs

@@ -2140,6 +2140,7 @@ impl CanisterManager {
                 }
             };
 
+            new_execution_state.exported_globals = execution_snapshot.exported_globals.clone();
             new_execution_state.stable_memory = Memory::from(&execution_snapshot.stable_memory);
             new_execution_state.wasm_memory = Memory::from(&execution_snapshot.wasm_memory);
             (instructions_used, Some(new_execution_state))

rs/execution_environment/src/execution_environment/tests/canister_snapshots.rs

@@ -1714,3 +1714,83 @@ fn load_canister_snapshot_charges_canister_cycles() {
         test.canister_state(canister_id).system_state.balance() < initial_balance - expected_charge
     );
 }
+
+#[test]
+fn snapshot_must_include_globals() {
+    let wat = r#"
+    (module
+        (import "ic0" "msg_reply" (func $msg_reply))
+        (import "ic0" "msg_reply_data_append"
+          (func $msg_reply_data_append (param i32 i32)))
+
+        (func $read_global
+          (i32.store
+            (i32.const 0)
+            (global.get 0)
+          )
+          (call $msg_reply_data_append
+            (i32.const 0)
+            (i32.const 4))
+          (call $msg_reply)
+        )
+
+        (func $increase_global
+          (global.set 0
+            (i32.add
+              (global.get 0)
+              (i32.const 1)
+            )
+          )
+          (call $msg_reply)
+        )
+
+        (memory $memory 1)
+        (export "memory" (memory $memory))
+        (global (export "counter") (mut i32) (i32.const 0))
+        (export "canister_query read_global" (func $read_global))
+        (export "canister_update increase_global" (func $increase_global))
+      )"#;
+    let wasm = wat::parse_str(wat).unwrap();
+
+    let mut test = ExecutionTestBuilder::new()
+        .with_snapshots(FlagStatus::Enabled)
+        .build();
+
+    // Create canister.
+    let canister_id = test.canister_from_binary(wasm).unwrap();
+
+    // Check that global is initially 0
+    let result = test
+        .non_replicated_query(canister_id, "read_global", vec![])
+        .unwrap();
+    assert_eq!(result, WasmResult::Reply(vec![0, 0, 0, 0]));
+
+    // Increase global to 1
+    test.ingress(canister_id, "increase_global", vec![])
+        .unwrap();
+
+    // Check that global is now 1
+    let result = test
+        .non_replicated_query(canister_id, "read_global", vec![])
+        .unwrap();
+    assert_eq!(result, WasmResult::Reply(vec![1, 0, 0, 0]));
+
+    // Take a snapshot.
+    let args = TakeCanisterSnapshotArgs::new(canister_id, None);
+    let result = test
+        .subnet_message("take_canister_snapshot", args.encode())
+        .unwrap();
+    let response = CanisterSnapshotResponse::decode(&result.bytes()).unwrap();
+    let snapshot_id = response.snapshot_id();
+
+    // Load the snapshot.
+    let args = LoadCanisterSnapshotArgs::new(canister_id, snapshot_id, None);
+    test.subnet_message("load_canister_snapshot", args.encode())
+        .unwrap();
+
+    // Check that global is still 1
+    let result = test
+        .non_replicated_query(canister_id, "read_global", vec![])
+        .unwrap();
+    assert_eq!(result, WasmResult::Reply(vec![1, 0, 0, 0]));
+}

rs/protobuf/def/state/canister_snapshot_bits/v1/canister_snapshot_bits.proto

@@ -15,4 +15,5 @@ message CanisterSnapshotBits {
   uint64 stable_memory_size = 8;
   uint64 wasm_memory_size = 9;
   uint64 total_size = 10;
+  repeated canister_state_bits.v1.Global exported_globals = 11;
 }

rs/protobuf/src/gen/state/state.canister_snapshot_bits.v1.rs

@@ -22,4 +22,6 @@ pub struct CanisterSnapshotBits {
     pub wasm_memory_size: u64,
     #[prost(uint64, tag = "10")]
     pub total_size: u64,
+    #[prost(message, repeated, tag = "11")]
+    pub exported_globals: ::prost::alloc::vec::Vec<super::super::canister_state_bits::v1::Global>,
 }

rs/replicated_state/src/canister_snapshots.rs

@@ -1,7 +1,7 @@
 use crate::{
-    canister_state::execution_state::Memory,
-    canister_state::system_state::wasm_chunk_store::WasmChunkStore, CanisterState, NumWasmPages,
-    PageMap,
+    canister_state::execution_state::{Global, Memory},
+    canister_state::system_state::wasm_chunk_store::WasmChunkStore,
+    CanisterState, NumWasmPages, PageMap,
 };
 use ic_sys::PAGE_SIZE;
 use ic_types::{CanisterId, NumBytes, SnapshotId, Time};
@@ -284,6 +284,11 @@ pub struct ExecutionStateSnapshot {
     /// The raw canister module.
     #[validate_eq(Ignore)]
     pub wasm_binary: CanisterModule,
+    /// The Wasm global variables.
+    /// Note: The hypervisor instrumentations exports all global variables,
+    /// including originally internal global variables.
+    #[validate_eq(Ignore)]
+    pub exported_globals: Vec<Global>,
     /// Snapshot of stable memory.
     #[validate_eq(CompareWithValidateEq)]
     pub stable_memory: PageMemory,
@@ -345,6 +350,7 @@ impl CanisterSnapshot {
             .ok_or(CanisterSnapshotError::EmptyExecutionState(canister_id))?;
         let execution_snapshot = ExecutionStateSnapshot {
             wasm_binary: execution_state.wasm_binary.binary.clone(),
+            exported_globals: execution_state.exported_globals.clone(),
             stable_memory: PageMemory::from(&execution_state.stable_memory),
             wasm_memory: PageMemory::from(&execution_state.wasm_memory),
         };
@@ -392,6 +398,10 @@ impl CanisterSnapshot {
         &self.execution_snapshot.wasm_binary
     }
 
+    pub fn exported_globals(&self) -> &Vec<Global> {
+        &self.execution_snapshot.exported_globals
+    }
+
     pub fn chunk_store(&self) -> &WasmChunkStore {
         &self.chunk_store
     }
@@ -457,6 +467,7 @@ mod tests {
     ) -> (SnapshotId, CanisterSnapshot) {
         let execution_snapshot = ExecutionStateSnapshot {
             wasm_binary: CanisterModule::new(vec![1, 2, 3]),
+            exported_globals: vec![Global::I32(1), Global::I64(2), Global::F64(0.1)],
             stable_memory: PageMemory {
                 page_map: PageMap::new_for_testing(),
                 size: NumWasmPages::new(10),

rs/replicated_state/src/canister_state/execution_state.rs

@@ -106,6 +106,8 @@ impl PartialEq<Global> for Global {
     }
 }
 
+impl Eq for Global {}
+
 impl From<&Global> for pb::Global {
     fn from(item: &Global) -> Self {
         match item {

rs/state_layout/src/state_layout.rs

@@ -201,6 +201,8 @@ pub struct CanisterSnapshotBits {
     pub wasm_memory_size: NumWasmPages,
     /// The total size of the snapshot in bytes.
     pub total_size: NumBytes,
+    /// State of the exported Wasm globals.
+    pub exported_globals: Vec<Global>,
 }
 
 #[derive(Clone)]
@@ -2513,6 +2515,11 @@ impl From<CanisterSnapshotBits> for pb_canister_snapshot_bits::CanisterSnapshotB
             stable_memory_size: item.stable_memory_size.get() as u64,
             wasm_memory_size: item.wasm_memory_size.get() as u64,
             total_size: item.total_size.get(),
+            exported_globals: item
+                .exported_globals
+                .iter()
+                .map(|global| global.into())
+                .collect(),
         }
     }
 }
@@ -2537,6 +2544,12 @@ impl TryFrom<pb_canister_snapshot_bits::CanisterSnapshotBits> for CanisterSnapsh
             }
             None => None,
         };
+
+        let mut exported_globals = Vec::with_capacity(item.exported_globals.len());
+        for global in item.exported_globals.into_iter() {
+            exported_globals.push(global.try_into()?);
+        }
+
         Ok(Self {
             snapshot_id: SnapshotId::from((canister_id, item.snapshot_id)),
             canister_id,
@@ -2552,6 +2565,7 @@ impl TryFrom<pb_canister_snapshot_bits::CanisterSnapshotBits> for CanisterSnapsh
             stable_memory_size: NumWasmPages::from(item.stable_memory_size as usize),
             wasm_memory_size: NumWasmPages::from(item.wasm_memory_size as usize),
             total_size: NumBytes::from(item.total_size),
+            exported_globals,
         })
     }
 }

rs/state_layout/src/state_layout/tests.rs

@@ -212,6 +212,7 @@ fn test_canister_snapshots_decode() {
         stable_memory_size: NumWasmPages::new(10),
         wasm_memory_size: NumWasmPages::new(10),
         total_size: NumBytes::new(100),
+        exported_globals: vec![Global::I32(1), Global::I64(2), Global::F64(0.1)],
     };
 
     let pb_bits =

rs/state_manager/src/checkpoint.rs

@@ -579,8 +579,11 @@ pub fn load_snapshot<P: ReadPolicy>(
             .deserialize(canister_snapshot_bits.binary_hash)?;
         durations.insert("snapshot_canister_module", starting_time.elapsed());
 
+        let exported_globals = canister_snapshot_bits.exported_globals.clone();
+
         ExecutionStateSnapshot {
             wasm_binary,
+            exported_globals,
             stable_memory,
             wasm_memory,
         }

rs/state_manager/src/tip.rs

@@ -1076,6 +1076,7 @@ fn serialize_snapshot_to_tip(
             stable_memory_size: canister_snapshot.stable_memory().size,
             wasm_memory_size: canister_snapshot.wasm_memory().size,
             total_size: canister_snapshot.size(),
+            exported_globals: canister_snapshot.exported_globals().clone(),
         }
         .into(),
     )?;