separate checkpoints_to_keep and inmemory_states_to_keep

rs/state_manager/src/lib.rs

@@ -2200,17 +2200,19 @@ impl StateManagerImpl {
                     state_metadata.bundled_manifest.as_ref().map(|_| *height)
                 });
 
-        let heights_to_keep: BTreeSet<Height> = states
+        let checkpoint_heights_to_keep: BTreeSet<Height> = states
             .states_metadata
             .keys()
             .copied()
             .filter(|height| {
                 *height == Self::INITIAL_STATE_HEIGHT || *height >= last_checkpoint_to_keep
             })
-            .chain(std::iter::once(latest_certified_height))
             .chain(latest_manifest_height)
             .collect();
 
+        let extra_inmemory_heights_to_keep: BTreeSet<Height> =
+            std::iter::once(latest_certified_height).collect();
+
         // Send object to deallocation thread if it has capacity.
         let deallocate = |x| {
             if self.deallocation_sender.len() < DEALLOCATION_BACKLOG_THRESHOLD {
@@ -2224,7 +2226,7 @@ impl StateManagerImpl {
 
         let (removed, retained) = states.snapshots.drain(0..).partition(|snapshot| {
             heights_to_remove.contains(&snapshot.height)
-                && !heights_to_keep.contains(&snapshot.height)
+                && !extra_inmemory_heights_to_keep.contains(&snapshot.height)
         });
         states.snapshots = retained;
 
@@ -2262,7 +2264,7 @@ impl StateManagerImpl {
         deallocate(Box::new(removed));
 
         for (height, metadata) in states.states_metadata.range(heights_to_remove) {
-            if heights_to_keep.contains(height) {
+            if checkpoint_heights_to_keep.contains(height) {
                 continue;
             }
             if let Some(ref checkpoint_layout) = metadata.checkpoint_layout {
@@ -2275,7 +2277,7 @@ impl StateManagerImpl {
             .certifications_metadata
             .split_off(&last_height_to_keep);
 
-        for h in heights_to_keep.iter() {
+        for h in extra_inmemory_heights_to_keep.iter() {
             if let Some(cert_metadata) = states.certifications_metadata.remove(h) {
                 certifications_metadata.insert(*h, cert_metadata);
             }
@@ -2305,7 +2307,7 @@ impl StateManagerImpl {
 
         let mut metadata_to_keep = states.states_metadata.split_off(&last_height_to_keep);
 
-        for h in heights_to_keep.iter() {
+        for h in checkpoint_heights_to_keep.iter() {
             if let Some(metadata) = states.states_metadata.remove(h) {
                 metadata_to_keep.insert(*h, metadata);
             }
@@ -2345,10 +2347,15 @@ impl StateManagerImpl {
                 });
 
             let state_heights = self.list_state_heights(CERT_ANY);
+            let certified_state_heights = self.list_state_heights(CERT_CERTIFIED);
+
+            debug_assert!(checkpoint_heights_to_keep
+                .iter()
+                .all(|h| unfiltered_checkpoint_heights.contains(h)));
 
-            debug_assert!(heights_to_keep.iter().all(|h| unfiltered_checkpoint_heights
-                .contains(h)
-                || *h == latest_certified_height));
+            debug_assert!(extra_inmemory_heights_to_keep
+                .iter()
+                .all(|h| certified_state_heights.contains(h)));
 
             debug_assert!(state_heights.contains(&latest_state_height));
             debug_assert!(state_heights.contains(&latest_certified_height));