feat: no state cloning during catch-up (#8411)

This PR optimizes the state manager to not clone the state while the
node is catching up.

In more detail, this PR performs the following changes:
- stores the latest certified height of the subnet (passed from
consensus via calls to `StateManagerImpl::remove_inmemory_states_below`)
in `StateManagerImpl`;
- introduces new state manager metrics (the latest certified height of
the subnet and the number of times
`StateManagerImpl::commmit_and_certify` did not clone state);
- does not clone state and does not compute certification metadata in
`StateManagerImpl::commmit_and_certify` if the committed height is below
the latest certified height of the subnet, the certification scope for
the committed state is set to partial certification (the "metadata"
variant), and the committed height is not divisible by 10 (i.e., we skip
the optimization every 10th round; to have a "recent" certified state
snapshot stored and available throughout the process of catching up).

Author: mraszyk

rs/consensus/src/consensus/purger.rs

@@ -316,11 +316,7 @@ impl Purger {
 
     /// Ask state manager to purge all states below the given height
     fn purge_replicated_state_by_finalized_certified_height(&self, pool: &PoolReader<'_>) {
-        let height = pool
-            .get_finalized_tip()
-            .context
-            .certified_height
-            .min(self.state_manager.latest_state_height());
+        let height = pool.get_finalized_tip().context.certified_height;
 
         let extra_heights_to_keep = get_pending_idkg_cup_heights(pool);
         self.state_manager
@@ -340,7 +336,8 @@ impl Purger {
     /// the given height can be removed.
     ///
     /// Note from the [`StateManager::remove_states_below`] docs:
-    ///  * The initial state (height = 0) cannot be removed.
+    ///  * The initial state (height = 0) is not removed.
+    ///  * The latest state is not removed.
     ///  * Some states matching the removal criteria might be kept alive.  For
     ///    example, the last fully persisted state might be preserved to
     ///    optimize future operations.

rs/interfaces/state_manager/src/lib.rs

@@ -226,9 +226,15 @@ pub trait StateManager: StateReader {
     /// Notify the state manager that states committed with partial certification
     /// state and heights strictly less than the specified `height` can be removed, except
     /// for any heights provided in `extra_heights_to_keep`, which will still be retained.
+    /// The specified `height` is expected to be the *latest certified height*
+    /// of the subnet, i.e., the latest height for which a valid certification is available
+    /// to consensus.
+    /// There are no guarantees for future heights in `extra_heights_to_keep`
+    /// w.r.t. the height of the latest state snapshot stored by the state manager.
     ///
     /// Note that:
-    ///  * The initial state (height = 0) cannot be removed.
+    ///  * The initial state (height = 0) is not removed.
+    ///  * The latest state is not removed.
     ///  * Some states matching the removal criteria might be kept alive.  For
     ///    example, the last fully persisted state might be preserved to
     ///    optimize future operations.

rs/state_manager/src/lib.rs

@@ -69,6 +69,7 @@ use ic_utils_thread::{JoinOnDrop, deallocator_thread::DeallocatorThread};
 use ic_wasm_types::ModuleLoadingStatus;
 use prometheus::{Histogram, HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec};
 use prost::Message;
+use std::cmp::min;
 use std::convert::{From, TryFrom};
 use std::fs::File;
 use std::fs::OpenOptions;
@@ -116,6 +117,10 @@ const CRITICAL_ERROR_REPLICATED_STATE_ALTERED_AFTER_CHECKPOINT: &str =
 /// How long to keep archived and diverged states.
 const ARCHIVED_DIVERGED_CHECKPOINT_MAX_AGE: Duration = Duration::from_secs(30 * 24 * 60 * 60); // 30 days
 
+/// The maximum number of consecutive rounds for which the optimization of
+/// skipping state cloning and certification metadata computation triggers.
+const MAX_CONSECUTIVE_ROUNDS_WITHOUT_STATE_CLONING: u64 = 10;
+
 /// Write an overlay file this many rounds before each checkpoint.
 pub const NUM_ROUNDS_BEFORE_CHECKPOINT_TO_WRITE_OVERLAY: u64 = 50;
 
@@ -173,6 +178,8 @@ pub struct StateManagerMetrics {
     merge_metrics: MergeMetrics,
     latest_hash_tree_size: IntGauge,
     latest_hash_tree_max_index: IntGauge,
+    latest_subnet_certified_height: IntGauge,
+    no_state_clone_count: IntCounter,
     tip_hash_count: IntCounter,
 }
 
@@ -459,6 +466,16 @@ impl StateManagerMetrics {
             "Largest index in the latest hash tree.",
         );
 
+        let latest_subnet_certified_height = metrics_registry.int_gauge(
+            "state_manager_latest_subnet_certified_height",
+            "Height of the latest validated certification.",
+        );
+
+        let no_state_clone_count = metrics_registry.int_counter(
+            "state_manager_no_state_clone_count",
+            "Number of heights whose states were not cloned and not stored by this node.",
+        );
+
         let tip_hash_count = metrics_registry.int_counter(
             "state_manager_tip_hash_count",
             "Number of tip heights whose state snapshot was not stored by this node and whose state hash was computed by this node.",
@@ -489,6 +506,8 @@ impl StateManagerMetrics {
             merge_metrics: MergeMetrics::new(metrics_registry),
             latest_hash_tree_size,
             latest_hash_tree_max_index,
+            latest_subnet_certified_height,
+            no_state_clone_count,
             tip_hash_count,
         }
     }
@@ -916,6 +935,7 @@ pub struct StateManagerImpl {
     // requested quite often and this causes high contention on the lock.
     latest_state_height: AtomicU64,
     latest_certified_height: AtomicU64,
+    latest_subnet_certified_height: AtomicU64,
     persist_metadata_guard: Arc<Mutex<()>>,
     tip_channel: Sender<TipRequest>,
     _tip_thread_handle: JoinOnDrop<()>,
@@ -1478,6 +1498,7 @@ impl StateManagerImpl {
 
         let latest_state_height = AtomicU64::new(0);
         let latest_certified_height = AtomicU64::new(0);
+        let latest_subnet_certified_height = AtomicU64::new(0);
 
         let initial_snapshot = Snapshot {
             height: Self::INITIAL_STATE_HEIGHT,
@@ -1592,6 +1613,7 @@ impl StateManagerImpl {
             deallocator_thread,
             latest_state_height,
             latest_certified_height,
+            latest_subnet_certified_height,
             persist_metadata_guard,
             tip_channel,
             _tip_thread_handle,
@@ -1939,8 +1961,8 @@ impl StateManagerImpl {
         } else {
             info!(
                 self.log,
-                "The previous certification metadata at height {} has been removed. This can happen when the replica \
-                syncs a newer state concurrently and removes the states below.",
+                "The previous certification metadata at height {} are not available. This can happen when the replica \
+                (i) catches up or (ii) syncs a newer state concurrently and removes the states below.",
                 prev_height,
             );
         }
@@ -2688,10 +2710,6 @@ impl StateManager for StateManagerImpl {
                 } else {
                     std::mem::drop(states);
 
-                    if tip_height != Self::INITIAL_STATE_HEIGHT {
-                        fatal!(self.log, "Bug: missing tip metadata @{}", tip_height);
-                    }
-
                     let mut tip_certification_metadata =
                         Self::compute_certification_metadata(&self.metrics, &self.log, &tip)
                             .unwrap_or_else(|err| {
@@ -2987,9 +3005,9 @@ impl StateManager for StateManagerImpl {
                     hash, certification.signed.content.hash
                 );
             }
+
             let latest_certified =
                 update_latest_height(&self.latest_certified_height, certification.height);
-
             self.metrics
                 .latest_certified_height
                 .set(latest_certified as i64);
@@ -3107,9 +3125,15 @@ impl StateManager for StateManagerImpl {
 
     /// Variant of `remove_states_below()` that only removes states committed with
     /// partial certification scope.
+    /// The specified `requested_height` is expected to be the *latest certified height*
+    /// of the subnet, i.e., the latest height for which a valid certification is available
+    /// to consensus.
+    /// There are no guarantees for future heights in `extra_heights_to_keep`
+    /// w.r.t. the height of the latest state snapshot stored by the state manager, i.e.,
+    /// for heights greater than `self.latest_state_height`.
     ///
     /// The following states are NOT removed:
-    /// * Any state with height >= requested_height
+    /// * Any state with height >= min(requested_height, latest state height)
     /// * Checkpoint heights
     /// * The latest state
     /// * The latest certified state
@@ -3126,6 +3150,14 @@ impl StateManager for StateManagerImpl {
             .with_label_values(&["remove_inmemory_states_below"])
             .start_timer();
 
+        let latest_subnet_certified_height =
+            update_latest_height(&self.latest_subnet_certified_height, requested_height);
+        self.metrics
+            .latest_subnet_certified_height
+            .set(latest_subnet_certified_height as i64);
+
+        let requested_height = min(requested_height, self.latest_state_height());
+
         // The latest state must be kept.
         let latest_state_height = self.latest_state_height();
         let oldest_height_to_keep = latest_state_height
@@ -3184,10 +3216,6 @@ impl StateManager for StateManagerImpl {
             .with_label_values(&["commit_and_certify"])
             .start_timer();
 
-        self.metrics
-            .tip_handler_queue_length
-            .set(self.tip_channel.len() as i64);
-
         self.populate_extra_metadata(&mut state, height);
 
         if let CertificationScope::Metadata = scope {
@@ -3206,6 +3234,50 @@ impl StateManager for StateManagerImpl {
             }
         }
 
+        let assert_tip_is_none = |states: &SharedState| {
+            // The following assert validates that we don't have two clients
+            // modifying TIP at the same time and that each commit_and_certify()
+            // is preceded by a call to take_tip().
+            if let Some((tip_height, _)) = &states.tip {
+                fatal!(
+                    self.log,
+                    "Attempt to commit state not borrowed from this StateManager, height = {}, tip_height = {}",
+                    height,
+                    tip_height,
+                );
+            }
+        };
+
+        // If the node is catching up (`height.get() < latest_subnet_certified_height`)
+        // and this is not a checkpoint height (`matches!(scope, CertificationScope::Metadata)`),
+        // then we do not clone, do not hash, and do not store the state and certification metadata.
+        // This optimization is skipped every `MAX_CONSECUTIVE_ROUNDS_WITHOUT_STATE_CLONING` heights
+        // so that we always have a reasonably "recent" state snapshot and
+        // its certification metadata available.
+        let latest_subnet_certified_height =
+            self.latest_subnet_certified_height.load(Ordering::Relaxed);
+        if matches!(scope, CertificationScope::Metadata)
+            && height.get() < latest_subnet_certified_height
+            && !height
+                .get()
+                .is_multiple_of(MAX_CONSECUTIVE_ROUNDS_WITHOUT_STATE_CLONING)
+        {
+            let mut states = self.states.write();
+            #[cfg(debug_assertions)]
+            check_certifications_metadata_snapshots_and_states_metadata_are_consistent(&states);
+
+            assert_tip_is_none(&states);
+
+            self.metrics.no_state_clone_count.inc();
+
+            states.tip = Some((height, state));
+            return;
+        }
+
+        self.metrics
+            .tip_handler_queue_length
+            .set(self.tip_channel.len() as i64);
+
         let mut state_metadata_and_compute_manifest_request: Option<(StateMetadata, TipRequest)> =
             None;
         let mut follow_up_tip_requests = Vec::new();
@@ -3254,17 +3326,7 @@ impl StateManager for StateManagerImpl {
         #[cfg(debug_assertions)]
         check_certifications_metadata_snapshots_and_states_metadata_are_consistent(&states);
 
-        // The following assert validates that we don't have two clients
-        // modifying TIP at the same time and that each commit_and_certify()
-        // is preceded by a call to take_tip().
-        if let Some((tip_height, _)) = &states.tip {
-            fatal!(
-                self.log,
-                "Attempt to commit state not borrowed from this StateManager, height = {}, tip_height = {}",
-                height,
-                tip_height,
-            );
-        }
+        assert_tip_is_none(&states);
 
         // It's possible that we already computed this state before.  We
         // validate that hashes agree to spot bugs causing non-determinism as
@@ -3996,6 +4058,27 @@ pub mod testing {
 
         /// Testing only: Wait till deallocation queue is empty.
         fn flush_deallocation_channel(&self);
+
+        /// Testing only: Returns heights in `states.snapshots`.
+        fn state_snapshot_heights(&self) -> Vec<Height>;
+
+        /// Testing only: Returns state at a given height in `states.snapshots`.
+        fn state_snapshot(&self, height: Height) -> Arc<ReplicatedState>;
+
+        /// Testing only: Returns heights in `states.certifications_metadata`.
+        fn certifications_metadata_heights(&self) -> Vec<Height>;
+
+        /// Testing only: Returns hash tree at a given height in `states.certifications_metadata`.
+        fn certifications_metadata_hash_tree(&self, height: Height) -> Option<Arc<HashTree>>;
+
+        /// Testing only: Returns state hash at a given height in `states.certifications_metadata`.
+        fn certifications_metadata_state_hash(&self, height: Height) -> CryptoHash;
+
+        /// Testing only: Returns certification at a given height in `states.certifications_metadata`.
+        fn certifications_metadata_certification(&self, height: Height) -> Option<Certification>;
+
+        /// Testing only: Returns `latest_subnet_certified_height`.
+        fn latest_subnet_certified_height(&self) -> u64;
     }
 
     impl StateManagerTesting for StateManagerImpl {
@@ -4023,5 +4106,61 @@ pub mod testing {
         fn flush_deallocation_channel(&self) {
             self.deallocator_thread.flush_deallocation_channel();
         }
+
+        fn state_snapshot_heights(&self) -> Vec<Height> {
+            let states = self.states.read();
+            states.snapshots.iter().map(|s| s.height).collect()
+        }
+
+        fn state_snapshot(&self, height: Height) -> Arc<ReplicatedState> {
+            let states = self.states.read();
+            states
+                .snapshots
+                .iter()
+                .find(|s| s.height == height)
+                .expect("Did not find state at given height")
+                .state
+                .clone()
+        }
+
+        fn certifications_metadata_heights(&self) -> Vec<Height> {
+            let states = self.states.read();
+            states.certifications_metadata.keys().cloned().collect()
+        }
+
+        fn certifications_metadata_hash_tree(&self, height: Height) -> Option<Arc<HashTree>> {
+            let states = self.states.read();
+            states
+                .certifications_metadata
+                .get(&height)
+                .expect("Did not find metadata at given height")
+                .hash_tree
+                .clone()
+                .map(|(hash_tree, _)| hash_tree)
+        }
+
+        fn certifications_metadata_state_hash(&self, height: Height) -> CryptoHash {
+            let states = self.states.read();
+            states
+                .certifications_metadata
+                .get(&height)
+                .expect("Did not find metadata at given height")
+                .certified_state_hash
+                .clone()
+        }
+
+        fn certifications_metadata_certification(&self, height: Height) -> Option<Certification> {
+            let states = self.states.read();
+            states
+                .certifications_metadata
+                .get(&height)
+                .expect("Did not find metadata at given height")
+                .certification
+                .clone()
+        }
+
+        fn latest_subnet_certified_height(&self) -> u64 {
+            self.latest_subnet_certified_height.load(Ordering::Relaxed)
+        }
     }
 }