chore: [MR-615] Refactor list_state_heights and make it an associated…

… method (dfinity#1690)

Close [MR-615](https://dfinity.atlassian.net/browse/MR-615)

This PR (1) excludes checkpoint_heights from the result of
list_state_heights and (2) removes it from the StateReader trait and
make it an associated method.

[MR-615]:
https://dfinity.atlassian.net/browse/MR-615?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

rs/interfaces/state_manager/mocks/src/lib.rs

@@ -1,7 +1,7 @@
 use ic_crypto_tree_hash::{LabeledTree, MixedHashTree};
 use ic_interfaces_state_manager::{
-    CertificationMask, CertificationScope, CertifiedStateSnapshot, Labeled, StateHashError,
-    StateManager, StateManagerResult, StateReader,
+    CertificationScope, CertifiedStateSnapshot, Labeled, StateHashError, StateManager,
+    StateManagerResult, StateReader,
 };
 use ic_replicated_state::ReplicatedState;
 use ic_types::{
@@ -46,11 +46,6 @@ mock! {
 
         fn deliver_state_certification(&self, certification: Certification);
 
-        fn list_state_heights(
-            &self,
-            cert_mask: CertificationMask,
-        ) -> Vec<Height>;
-
         fn remove_states_below(&self, height: Height);
 
         fn remove_inmemory_states_below(&self, height: Height);

rs/interfaces/state_manager/src/lib.rs

@@ -222,23 +222,6 @@ pub trait StateManager: StateReader {
         cup_interval_length: Height,
     );
 
-    /// Returns the list of heights corresponding to accessible states matching
-    /// the mask.  E.g. `list_state_heights(CERT_ANY)` will return all
-    /// accessible states.
-    ///
-    /// Note that the initial state at height 0 is considered uncertified from
-    /// the State Manager point of view.  This is because the protocol requires
-    /// each replica to individually obtain the initial state using some
-    /// out-of-band mechanism (i.e., not state sync).  Also note that the
-    /// authenticity of this initial state will be verified by some protocol
-    /// external to this component.
-    ///
-    /// The list of heights is guaranteed to be
-    /// * Non-empty if `cert_mask = CERT_ANY` as it will contain at least height
-    ///   0 even if no states were committed yet.
-    /// * Sorted in ascending order.
-    fn list_state_heights(&self, cert_mask: CertificationMask) -> Vec<Height>;
-
     /// Notify this state manager that states with heights strictly less than
     /// the specified `height` can be removed.
     ///

rs/state_manager/src/lib.rs

@@ -32,9 +32,9 @@ use ic_interfaces_certified_stream_store::{
     CertifiedStreamStore, DecodeStreamError, EncodeStreamError,
 };
 use ic_interfaces_state_manager::{
-    CertificationMask, CertificationScope, CertifiedStateSnapshot, Labeled,
-    PermanentStateHashError::*, StateHashError, StateManager, StateManagerError,
-    StateManagerResult, StateReader, TransientStateHashError::*, CERT_CERTIFIED, CERT_UNCERTIFIED,
+    CertificationScope, CertifiedStateSnapshot, Labeled, PermanentStateHashError::*,
+    StateHashError, StateManager, StateManagerError, StateManagerResult, StateReader,
+    TransientStateHashError::*,
 };
 use ic_logger::{debug, error, fatal, info, warn, ReplicaLogger};
 use ic_metrics::{buckets::decimal_buckets, MetricsRegistry};
@@ -2392,6 +2392,62 @@ impl StateManagerImpl {
         result
     }
 
+    /// Returns the list of heights corresponding to snapshots matching
+    /// the mask. E.g. `list_state_heights(CERT_ANY)` will return all snapshots.
+    ///
+    /// Note that the initial state at height 0 is considered uncertified from
+    /// the State Manager point of view.  This is because the protocol requires
+    /// each replica to individually obtain the initial state using some
+    /// out-of-band mechanism (i.e., not state sync).  Also note that the
+    /// authenticity of this initial state will be verified by some protocol
+    /// external to this component.
+    ///
+    /// The list of heights is guaranteed to be
+    /// * Non-empty if `cert_mask = CERT_ANY` as it will contain at least height
+    ///   0 even if no states were committed yet.
+    /// * Sorted in ascending order.
+    #[allow(dead_code)]
+    pub fn list_state_heights(
+        &self,
+        cert_mask: ic_interfaces_state_manager::CertificationMask,
+    ) -> Vec<Height> {
+        let _timer = self
+            .metrics
+            .api_call_duration
+            .with_label_values(&["list_state_heights"])
+            .start_timer();
+
+        fn matches(
+            cert: Option<&Certification>,
+            mask: ic_interfaces_state_manager::CertificationMask,
+        ) -> bool {
+            match cert {
+                Some(_) => mask.is_set(ic_interfaces_state_manager::CERT_CERTIFIED),
+                None => mask.is_set(ic_interfaces_state_manager::CERT_UNCERTIFIED),
+            }
+        }
+
+        let states = self.states.read();
+
+        let heights: BTreeSet<_> = states
+            .snapshots
+            .iter()
+            .map(|snapshot| snapshot.height)
+            .filter(|h| {
+                matches(
+                    states
+                        .certifications_metadata
+                        .get(h)
+                        .and_then(|metadata| metadata.certification.as_ref()),
+                    cert_mask,
+                )
+            })
+            .collect();
+
+        // convert the b-tree into a vector
+        heights.into_iter().collect()
+    }
+
     // Creates a checkpoint and switches state to it.
     fn create_checkpoint_and_switch(
         &self,
@@ -3159,45 +3215,6 @@ impl StateManager for StateManagerImpl {
         }
     }
 
-    /// # Panics
-    ///
-    /// This method panics if checkpoint labels can not be retrieved
-    /// from the disk.
-    fn list_state_heights(&self, cert_mask: CertificationMask) -> Vec<Height> {
-        let _timer = self
-            .metrics
-            .api_call_duration
-            .with_label_values(&["list_state_heights"])
-            .start_timer();
-
-        fn matches(cert: Option<&Certification>, mask: CertificationMask) -> bool {
-            match cert {
-                Some(_) => mask.is_set(CERT_CERTIFIED),
-                None => mask.is_set(CERT_UNCERTIFIED),
-            }
-        }
-
-        let states = self.states.read();
-
-        let heights: BTreeSet<_> = self
-            .checkpoint_heights()
-            .into_iter()
-            .chain(states.snapshots.iter().map(|snapshot| snapshot.height))
-            .filter(|h| {
-                matches(
-                    states
-                        .certifications_metadata
-                        .get(h)
-                        .and_then(|metadata| metadata.certification.as_ref()),
-                    cert_mask,
-                )
-            })
-            .collect();
-
-        // convert the b-tree into a vector
-        heights.into_iter().collect()
-    }
-
     /// This method instructs the state manager that Consensus doesn't need
     /// any states strictly lower than the specified `height`.  The
     /// implementation purges some of these states using the heuristic

rs/test_utilities/src/state_manager.rs

@@ -4,10 +4,9 @@ use ic_interfaces_certified_stream_store::{
     CertifiedStreamStore, DecodeStreamError, EncodeStreamError,
 };
 use ic_interfaces_state_manager::{
-    CertificationMask, CertificationScope, CertifiedStateSnapshot, Labeled,
-    PermanentStateHashError::*, StateHashError, StateManager, StateManagerError,
-    StateManagerResult, StateReader, TransientStateHashError::*, CERT_ANY, CERT_CERTIFIED,
-    CERT_UNCERTIFIED,
+    CertificationScope, CertifiedStateSnapshot, Labeled, PermanentStateHashError::*,
+    StateHashError, StateManager, StateManagerError, StateManagerResult, StateReader,
+    TransientStateHashError::*,
 };
 use ic_interfaces_state_manager_mocks::MockStateManager;
 use ic_registry_subnet_type::SubnetType;
@@ -42,13 +41,6 @@ impl Snapshot {
     fn make_labeled_state(&self) -> Labeled<Arc<ReplicatedState>> {
         Labeled::new(self.height, self.state.clone())
     }
-
-    fn certification_mask(&self) -> CertificationMask {
-        match self.certification {
-            Some(_) => CERT_CERTIFIED,
-            None => CERT_UNCERTIFIED,
-        }
-    }
 }
 
 /// A fake implementation of the `StateManager` interface.
@@ -96,9 +88,10 @@ impl FakeStateManager {
     }
 }
 
+const INITIAL_STATE_HEIGHT: Height = Height::new(0);
 fn initial_state() -> Labeled<Arc<ReplicatedState>> {
     Labeled::new(
-        Height::new(0),
+        INITIAL_STATE_HEIGHT,
         Arc::new(ReplicatedState::new(
             subnet_test_id(1),
             SubnetType::Application,
@@ -196,21 +189,6 @@ impl StateManager for FakeStateManager {
         }
     }
 
-    fn list_state_heights(&self, cert_mask: CertificationMask) -> Vec<Height> {
-        self.states
-            .read()
-            .unwrap()
-            .iter()
-            .filter_map(|snapshot| {
-                if cert_mask.is_set(snapshot.certification_mask()) {
-                    Some(snapshot.height)
-                } else {
-                    None
-                }
-            })
-            .collect()
-    }
-
     fn remove_states_below(&self, height: Height) {
         self.states
             .write()
@@ -259,9 +237,11 @@ impl StateReader for FakeStateManager {
     type State = ReplicatedState;
 
     fn latest_state_height(&self) -> Height {
-        *StateManager::list_state_heights(self, CERT_ANY)
-            .last()
+        self.states
+            .read()
             .unwrap()
+            .last()
+            .map_or(INITIAL_STATE_HEIGHT, |snap| snap.height)
     }
 
     // No certification support in FakeStateManager
@@ -682,10 +662,6 @@ impl StateManager for RefMockStateManager {
             .deliver_state_certification(certification)
     }
 
-    fn list_state_heights(&self, cert_mask: CertificationMask) -> Vec<Height> {
-        self.mock.read().unwrap().list_state_heights(cert_mask)
-    }
-
     fn remove_states_below(&self, height: Height) {
         self.mock.read().unwrap().remove_states_below(height)
     }