From 3025158422126955a43f6487418363f8bb345d5e Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Tue, 10 Dec 2024 20:21:44 +0800
Subject: [PATCH] Sketch finalization migration

---
 beacon_node/beacon_chain/src/migrate.rs     | 10 +++--
 beacon_node/store/src/garbage_collection.rs |  7 +++-
 beacon_node/store/src/hot_cold_store.rs     | 46 +++++++++++++++++----
 beacon_node/store/src/lib.rs                |  6 ++-
 4 files changed, 54 insertions(+), 15 deletions(-)

diff --git a/beacon_node/beacon_chain/src/migrate.rs b/beacon_node/beacon_chain/src/migrate.rs
index 2546447c37e..67d54f7ef90 100644
--- a/beacon_node/beacon_chain/src/migrate.rs
+++ b/beacon_node/beacon_chain/src/migrate.rs
@@ -709,10 +709,12 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> BackgroundMigrator<E, Ho
                 ]
             })
             .chain(abandoned_states.into_iter().flat_map(|(slot, state_hash)| {
-                [
-                    StoreOp::DeleteState(state_hash.into(), Some(slot)),
-                    StoreOp::DeleteStateHotDiff(state_hash.into()),
-                ]
+                [StoreOp::DeleteState {
+                    state_root: state_hash.into(),
+                    state_slot: Some(slot),
+                    // Abandoned forks will never be used, so we can safely delete the diffs
+                    prune_hot_diff: true,
+                }]
             }))
             .collect();
 
diff --git a/beacon_node/store/src/garbage_collection.rs b/beacon_node/store/src/garbage_collection.rs
index 5f8ed8f5e73..165272fd05c 100644
--- a/beacon_node/store/src/garbage_collection.rs
+++ b/beacon_node/store/src/garbage_collection.rs
@@ -20,7 +20,12 @@ where
             self.iter_temporary_state_roots()
                 .try_fold(vec![], |mut ops, state_root| {
                     let state_root = state_root?;
-                    ops.push(StoreOp::DeleteState(state_root, None));
+                    ops.push(StoreOp::DeleteState {
+                        state_root,
+                        state_slot: None,
+                        // This states will never be used, safe to delete the hot diffs
+                        prune_hot_diff: true,
+                    });
                     Result::<_, Error>::Ok(ops)
                 })?;
 
diff --git a/beacon_node/store/src/hot_cold_store.rs b/beacon_node/store/src/hot_cold_store.rs
index b1c3ca30ff7..1be8b55406a 100644
--- a/beacon_node/store/src/hot_cold_store.rs
+++ b/beacon_node/store/src/hot_cold_store.rs
@@ -1051,11 +1051,17 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
     /// than the split point. You shouldn't delete states from the finalized portion of the chain
     /// (which are frozen, and won't be deleted), or valid descendents of the finalized checkpoint
     /// (which will be deleted by this function but shouldn't be).
-    pub fn delete_state(&self, state_root: &Hash256, slot: Slot) -> Result<(), Error> {
-        self.do_atomically_with_block_and_blobs_cache(vec![StoreOp::DeleteState(
-            *state_root,
-            Some(slot),
-        )])
+    pub fn delete_state(
+        &self,
+        state_root: &Hash256,
+        slot: Slot,
+        prune_hot_diff: bool,
+    ) -> Result<(), Error> {
+        self.do_atomically_with_block_and_blobs_cache(vec![StoreOp::DeleteState {
+            state_root: *state_root,
+            state_slot: Some(slot),
+            prune_hot_diff,
+        }])
     }
 
     pub fn forwards_block_roots_iterator(
@@ -1196,7 +1202,11 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
                     }
                 }
 
-                StoreOp::DeleteState(state_root, slot) => {
+                StoreOp::DeleteState {
+                    state_root,
+                    state_slot: slot,
+                    prune_hot_diff,
+                } => {
                     // Delete the hot state summary.
                     let state_summary_key =
                         get_key_for_col(DBColumn::BeaconStateSummary.into(), state_root.as_slice());
@@ -1211,6 +1221,12 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
                     key_value_batch.push(KeyValueStoreOp::DeleteKey(state_temp_key));
 
                     // TODO(hdiff): Should delete the diff under this state root if any
+                    if prune_hot_diff {
+                        key_value_batch.push(KeyValueStoreOp::DeleteKey(get_key_for_col(
+                            DBColumn::BeaconStateHotDiff.into(),
+                            state_root.as_slice(),
+                        )));
+                    }
 
                     // TODO(hdiff): Review under HDiff
                     if slot.map_or(true, |slot| slot % E::slots_per_epoch() == 0) {
@@ -1372,7 +1388,7 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
                     self.state_cache.lock().delete_block_states(&block_root);
                 }
 
-                StoreOp::DeleteState(state_root, _) => {
+                StoreOp::DeleteState { state_root, .. } => {
                     self.state_cache.lock().delete_state(&state_root)
                 }
 
@@ -3070,7 +3086,12 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
                         "slot" => summary.slot,
                         "reason" => reason,
                     );
-                    state_delete_batch.push(StoreOp::DeleteState(state_root, Some(summary.slot)));
+                    state_delete_batch.push(StoreOp::DeleteState {
+                        state_root,
+                        state_slot: Some(summary.slot),
+                        // TODO(hdiff): The logic here is duplicated from `migrate_database`
+                        prune_hot_diff: false,
+                    });
                 }
             }
         }
@@ -3172,7 +3193,14 @@ pub fn migrate_database<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>(
 
         // Delete the old summary, and the full state if we lie on an epoch boundary.
         if !required_finalized_diff_state_roots.contains(&state_root) {
-            hot_db_ops.push(StoreOp::DeleteState(state_root, Some(slot)));
+            hot_db_ops.push(StoreOp::DeleteState {
+                state_root,
+                state_slot: Some(slot),
+                // TODO(hdiff): must not delete *all* finalized hdiffs. Instead, keep
+                // the more recent diff of each layer including the snapshot.
+                // Implement a routine somewhere to figure out which diffs should be kept
+                prune_hot_diff: false,
+            });
         }
 
         // Do not try to store states if a restore point is yet to be stored, or will never be
diff --git a/beacon_node/store/src/lib.rs b/beacon_node/store/src/lib.rs
index 699ffdd7159..0149697a09d 100644
--- a/beacon_node/store/src/lib.rs
+++ b/beacon_node/store/src/lib.rs
@@ -241,7 +241,11 @@ pub enum StoreOp<'a, E: EthSpec> {
     DeleteBlock(Hash256),
     DeleteBlobs(Hash256),
     DeleteDataColumns(Hash256, Vec<ColumnIndex>),
-    DeleteState(Hash256, Option<Slot>),
+    DeleteState {
+        state_root: Hash256,
+        state_slot: Option<Slot>,
+        prune_hot_diff: bool,
+    },
     DeleteStateHotDiff(Hash256),
     DeleteExecutionPayload(Hash256),
     DeleteSyncCommitteeBranch(Hash256),