feat: EXC-1527: Add canister snapshots to ReplicatedState

Author: AlexandraZapuc

rs/http_endpoints/public/src/read_state/canister.rs

@@ -430,7 +430,9 @@ mod test {
     use hyper::StatusCode;
     use ic_crypto_tree_hash::{Digest, Label, MixedHashTree, Path};
     use ic_registry_subnet_type::SubnetType;
-    use ic_replicated_state::{CanisterQueues, ReplicatedState, SystemMetadata};
+    use ic_replicated_state::{
+        canister_snapshots::CanisterSnapshots, CanisterQueues, ReplicatedState, SystemMetadata,
+    };
     use ic_test_utilities::{
         state::insert_dummy_canister,
         types::ids::{canister_test_id, subnet_test_id, user_test_id},
@@ -558,6 +560,7 @@ mod test {
             metadata,
             CanisterQueues::default(),
             RawQueryStats::default(),
+            CanisterSnapshots::default(),
         );
         assert_eq!(
             verify_paths(

rs/http_endpoints/public/src/state_reader_executor.rs

@@ -80,7 +80,9 @@ mod tests {
     use ic_crypto_tree_hash::{flatmap, Label, LabeledTree};
     use ic_interfaces_state_manager_mocks::MockStateManager;
     use ic_registry_subnet_type::SubnetType;
-    use ic_replicated_state::{CanisterQueues, ReplicatedState, SystemMetadata};
+    use ic_replicated_state::{
+        canister_snapshots::CanisterSnapshots, CanisterQueues, ReplicatedState, SystemMetadata,
+    };
     use ic_test_utilities::{state::ReplicatedStateBuilder, types::ids::subnet_test_id};
     use ic_test_utilities_time::mock_time;
     use ic_types::{
@@ -111,6 +113,7 @@ mod tests {
                         metadata,
                         CanisterQueues::default(),
                         RawQueryStats::default(),
+                        CanisterSnapshots::default(),
                     )),
                 )
             });

rs/http_endpoints/public/tests/common/mod.rs

@@ -34,7 +34,10 @@ use ic_registry_keys::{
 use ic_registry_provisional_whitelist::ProvisionalWhitelist;
 use ic_registry_routing_table::{CanisterMigrations, RoutingTable};
 use ic_registry_subnet_type::SubnetType;
-use ic_replicated_state::{CanisterQueues, NetworkTopology, ReplicatedState, SystemMetadata};
+use ic_replicated_state::{
+    canister_snapshots::CanisterSnapshots, CanisterQueues, NetworkTopology, ReplicatedState,
+    SystemMetadata,
+};
 use ic_test_utilities::{
     crypto::{temp_crypto_component_with_fake_registry, CryptoReturningOk},
     state::ReplicatedStateBuilder,
@@ -211,6 +214,7 @@ pub fn default_get_latest_state() -> Labeled<Arc<ReplicatedState>> {
             metadata,
             CanisterQueues::default(),
             RawQueryStats::default(),
+            CanisterSnapshots::default(),
         )),
     )
 }

rs/ingress_manager/benches/handle_ingress.rs

@@ -32,7 +32,9 @@ use ic_registry_client::client::RegistryClientImpl;
 use ic_registry_keys::make_subnet_record_key;
 use ic_registry_proto_data_provider::ProtoRegistryDataProvider;
 use ic_registry_subnet_type::SubnetType;
-use ic_replicated_state::{CanisterQueues, ReplicatedState, SystemMetadata};
+use ic_replicated_state::{
+    canister_snapshots::CanisterSnapshots, CanisterQueues, ReplicatedState, SystemMetadata,
+};
 use ic_test_utilities::{
     crypto::temp_crypto_component_with_fake_registry,
     cycles_account_manager::CyclesAccountManagerBuilder,
@@ -189,6 +191,7 @@ where
                         metadata,
                         CanisterQueues::default(),
                         RawQueryStats::default(),
+                        CanisterSnapshots::default(),
                     )),
                 )
             });

rs/replicated_state/src/canister_snapshots.rs

@@ -80,6 +80,11 @@ impl CanisterSnapshots {
     pub fn take_unflushed_changes(&mut self) -> Vec<SnapshotOperation> {
         std::mem::take(&mut self.unflushed_changes)
     }
+
+    /// Returns true if unflushed changes list is empty.
+    pub fn is_unflushed_changes_empty(&self) -> bool {
+        self.unflushed_changes.is_empty()
+    }
 }
 
 /// Contains all information related to a canister snapshot.

rs/replicated_state/src/replicated_state.rs

@@ -3,6 +3,7 @@ use super::{
     metadata_state::{IngressHistoryState, Stream, Streams, SystemMetadata},
 };
 use crate::{
+    canister_snapshots::CanisterSnapshots,
     canister_state::queues::CanisterQueuesLoopDetector,
     canister_state::system_state::{push_input, CanisterOutputQueuesIterator},
     metadata_state::{subnet_call_context_manager::SignWithEcdsaContext, StreamMap},
@@ -406,6 +407,9 @@ pub struct ReplicatedState {
     /// Temporary query stats received during the current epoch.
     /// Reset during the start of each epoch.
     pub epoch_query_stats: RawQueryStats,
+
+    /// Manages the canister snapshots.
+    pub canister_snapshots: CanisterSnapshots,
 }
 
 impl ReplicatedState {
@@ -417,6 +421,7 @@ impl ReplicatedState {
             subnet_queues: CanisterQueues::default(),
             consensus_queue: Vec::new(),
             epoch_query_stats: RawQueryStats::default(),
+            canister_snapshots: CanisterSnapshots::default(),
         }
     }
 
@@ -426,13 +431,15 @@ impl ReplicatedState {
         metadata: SystemMetadata,
         subnet_queues: CanisterQueues,
         epoch_query_stats: RawQueryStats,
+        canister_snapshots: CanisterSnapshots,
     ) -> Self {
         let mut res = Self {
             canister_states,
             metadata,
             subnet_queues,
             consensus_queue: Vec::new(),
             epoch_query_stats,
+            canister_snapshots,
         };
         res.update_stream_responses_size_bytes();
         res
@@ -923,6 +930,7 @@ impl ReplicatedState {
             mut subnet_queues,
             consensus_queue,
             epoch_query_stats: _,
+            canister_snapshots,
         } = self;
 
         // Consensus queue is always empty at the end of the round.
@@ -958,6 +966,7 @@ impl ReplicatedState {
             subnet_queues,
             consensus_queue,
             epoch_query_stats: RawQueryStats::default(), // Don't preserve query stats during subnet splitting.
+            canister_snapshots,
         })
     }
 
@@ -980,6 +989,7 @@ impl ReplicatedState {
             ref mut subnet_queues,
             consensus_queue: _,
             epoch_query_stats: _,
+            canister_snapshots: _,
         } = self;
 
         // Reset query stats after subnet split
@@ -1140,6 +1150,8 @@ pub mod testing {
             subnet_queues: Default::default(),
             consensus_queue: Default::default(),
             epoch_query_stats: Default::default(),
+            // TODO(EXC-1527): Handle canister snapshots during a subnet split.
+            canister_snapshots: CanisterSnapshots::default(),
         };
     }
 }

rs/state_manager/src/checkpoint.rs

@@ -7,6 +7,7 @@ use ic_base_types::{subnet_id_try_from_protobuf, CanisterId};
 use ic_config::flag_status::FlagStatus;
 use ic_logger::error;
 use ic_registry_subnet_type::SubnetType;
+use ic_replicated_state::canister_snapshots::CanisterSnapshots;
 use ic_replicated_state::page_map::PageAllocatorFileDescriptor;
 use ic_replicated_state::{
     canister_state::execution_state::WasmBinary, page_map::PageMap, CanisterMetrics, CanisterState,
@@ -258,8 +259,15 @@ pub fn load_checkpoint<P: ReadPolicy + Send + Sync>(
         canister_states
     };
 
-    let state =
-        ReplicatedState::new_from_checkpoint(canister_states, metadata, subnet_queues, query_stats);
+    // TODO(EXC-1539): Load canister snapshots from checkpoint.
+    let canister_snapshots = CanisterSnapshots::default();
+    let state = ReplicatedState::new_from_checkpoint(
+        canister_states,
+        metadata,
+        subnet_queues,
+        query_stats,
+        canister_snapshots,
+    );
 
     Ok(state)
 }

rs/state_manager/src/tip.rs

@@ -788,6 +788,9 @@ fn serialize_to_tip(
     metrics: &StorageMetrics,
     lsmt_storage: FlagStatus,
 ) -> Result<(), CheckpointError> {
+    //TODO(MR-530): Implement serializing canister snapshots.
+    debug_assert!(state.canister_snapshots.is_unflushed_changes_empty());
+
     // Serialize ingress history separately. The `SystemMetadata` proto does not
     // encode it.
     //