From 8762d82adfd9316adf787df7dfc6941f8c76b49e Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Thu, 23 May 2024 10:17:53 +1000
Subject: [PATCH] Fix hot state disk leak (#5768)

* Fix hot state leak

* Don't delete the genesis state when split is 0x0!
---
 .../beacon_chain/src/block_verification.rs    | 22 ++++----
 beacon_node/beacon_chain/src/migrate.rs       |  5 ++
 beacon_node/store/src/hot_cold_store.rs       | 51 +++++++++++++++++++
 3 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/beacon_node/beacon_chain/src/block_verification.rs b/beacon_node/beacon_chain/src/block_verification.rs
index 866dde5a76..f4f6526a56 100644
--- a/beacon_node/beacon_chain/src/block_verification.rs
+++ b/beacon_node/beacon_chain/src/block_verification.rs
@@ -1382,18 +1382,20 @@ impl<T: BeaconChainTypes> ExecutionPendingBlock<T> {
         let catchup_timer = metrics::start_timer(&metrics::BLOCK_PROCESSING_CATCHUP_STATE);
 
         // Stage a batch of operations to be completed atomically if this block is imported
-        // successfully. We include the state root of the pre-state, which may be an advanced state
-        // that was stored in the DB with a `temporary` flag.
+        // successfully. If there is a skipped slot, we include the state root of the pre-state,
+        // which may be an advanced state that was stored in the DB with a `temporary` flag.
         let mut state = parent.pre_state;
 
-        let mut confirmed_state_roots = if state.slot() > parent.beacon_block.slot() {
-            // Advanced pre-state. Delete its temporary flag.
-            let pre_state_root = state.update_tree_hash_cache()?;
-            vec![pre_state_root]
-        } else {
-            // Pre state is parent state. It is already stored in the DB without temporary status.
-            vec![]
-        };
+        let mut confirmed_state_roots =
+            if block.slot() > state.slot() && state.slot() > parent.beacon_block.slot() {
+                // Advanced pre-state. Delete its temporary flag.
+                let pre_state_root = state.update_tree_hash_cache()?;
+                vec![pre_state_root]
+            } else {
+                // Pre state is either unadvanced, or should not be stored long-term because there
+                // is no skipped slot between `parent` and `block`.
+                vec![]
+            };
 
         // The block must have a higher slot than its parent.
         if block.slot() <= parent.beacon_block.slot() {
diff --git a/beacon_node/beacon_chain/src/migrate.rs b/beacon_node/beacon_chain/src/migrate.rs
index ad597bf92a..08b2a51720 100644
--- a/beacon_node/beacon_chain/src/migrate.rs
+++ b/beacon_node/beacon_chain/src/migrate.rs
@@ -703,6 +703,11 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
         ));
 
         store.do_atomically_with_block_and_blobs_cache(batch)?;
+
+        // Do a quick separate pass to delete obsoleted hot states, usually pre-states from the state
+        // advance which are not canonical due to blocks being applied on top.
+        store.prune_old_hot_states()?;
+
         debug!(log, "Database pruning complete");
 
         Ok(PruningOutcome::Successful {
diff --git a/beacon_node/store/src/hot_cold_store.rs b/beacon_node/store/src/hot_cold_store.rs
index 484a1139bf..9c247c983a 100644
--- a/beacon_node/store/src/hot_cold_store.rs
+++ b/beacon_node/store/src/hot_cold_store.rs
@@ -2485,6 +2485,57 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
 
         Ok(())
     }
+
+    /// Prune states from the hot database which are prior to the split.
+    ///
+    /// This routine is important for cleaning up advanced states which are stored in the database
+    /// with a temporary flag.
+    pub fn prune_old_hot_states(&self) -> Result<(), Error> {
+        let split = self.get_split_info();
+        debug!(
+            self.log,
+            "Database state pruning started";
+            "split_slot" => split.slot,
+        );
+        let mut state_delete_batch = vec![];
+        for res in self
+            .hot_db
+            .iter_column::<Hash256>(DBColumn::BeaconStateSummary)
+        {
+            let (state_root, summary_bytes) = res?;
+            let summary = HotStateSummary::from_ssz_bytes(&summary_bytes)?;
+
+            if summary.slot <= split.slot {
+                let old = summary.slot < split.slot;
+                let non_canonical = summary.slot == split.slot
+                    && state_root != split.state_root
+                    && !split.state_root.is_zero();
+                if old || non_canonical {
+                    let reason = if old {
+                        "old dangling state"
+                    } else {
+                        "non-canonical"
+                    };
+                    debug!(
+                        self.log,
+                        "Deleting state";
+                        "state_root" => ?state_root,
+                        "slot" => summary.slot,
+                        "reason" => reason,
+                    );
+                    state_delete_batch.push(StoreOp::DeleteState(state_root, Some(summary.slot)));
+                }
+            }
+        }
+        let num_deleted_states = state_delete_batch.len();
+        self.do_atomically_with_block_and_blobs_cache(state_delete_batch)?;
+        debug!(
+            self.log,
+            "Database state pruning complete";
+            "num_deleted_states" => num_deleted_states,
+        );
+        Ok(())
+    }
 }
 
 /// Advance the split point of the store, moving new finalized states to the freezer.