Fix DB manager downgrade

michaelsproul · michaelsproul · commit 064115c8dc8c · 2025-05-29T13:49:25.000+10:00
diff --git a/beacon_node/beacon_chain/src/schema_change/migration_schema_v24.rs b/beacon_node/beacon_chain/src/schema_change/migration_schema_v24.rs
@@ -3,14 +3,15 @@ use crate::{
     summaries_dag::{DAGStateSummary, DAGStateSummaryV22, StateSummariesDAG},
 };
 use ssz::{Decode, DecodeError, Encode};
-use ssz_derive::{Decode, Encode};
+use ssz_derive::Encode;
 use std::{
     sync::Arc,
     time::{Duration, Instant},
 };
 use store::{
-    hdiff::StorageStrategy, hot_cold_store::OptionalDiffBaseState, DBColumn, Error, HotColdDB,
-    HotStateSummary, KeyValueStore, KeyValueStoreOp, StoreItem,
+    hdiff::StorageStrategy,
+    hot_cold_store::{HotStateSummaryV22, OptionalDiffBaseState},
+    DBColumn, Error, HotColdDB, HotStateSummary, KeyValueStore, KeyValueStoreOp, StoreItem,
 };
 use tracing::{debug, info, warn};
 use types::{
@@ -158,27 +159,6 @@ impl StoreItem for PruningCheckpoint {
     }
 }
 
-#[derive(Debug, Clone, Copy, Encode, Decode)]
-pub struct HotStateSummaryV22 {
-    slot: Slot,
-    latest_block_root: Hash256,
-    epoch_boundary_state_root: Hash256,
-}
-
-impl StoreItem for HotStateSummaryV22 {
-    fn db_column() -> DBColumn {
-        DBColumn::BeaconStateSummary
-    }
-
-    fn as_store_bytes(&self) -> Vec<u8> {
-        self.as_ssz_bytes()
-    }
-
-    fn from_store_bytes(bytes: &[u8]) -> Result<Self, Error> {
-        Ok(Self::from_ssz_bytes(bytes)?)
-    }
-}
-
 pub fn upgrade_to_v24<T: BeaconChainTypes>(
     db: Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
 ) -> Result<Vec<KeyValueStoreOp>, Error> {
diff --git a/beacon_node/store/src/hot_cold_store.rs b/beacon_node/store/src/hot_cold_store.rs
@@ -289,15 +289,26 @@ impl<E: EthSpec> HotColdDB<E, BeaconNodeBackend<E>, BeaconNodeBackend<E>> {
         // stop and restart correctly. This needs to occur *before* running any migrations
         // because some migrations load states and depend on the split.
         //
-        // V24: `load_split` needs to load a hot state summary, which need to be migrated from V22
-        // to V24. Attempting to `load_split` here before the migration will trigger an SSZ decode
-        // error. Instead we load the partial split, and load the full split after the migration.
-        if let Some(split) = db.load_split_partial()? {
+        // We use a method that is ambivalent to the state summaries being V22 or V24, because
+        // we need to support several scenarios:
+        //
+        // - Migrating from V22 to V24: initially summaries are V22 , and we need
+        //   to be able to load a block root from them. Loading the split partially at first
+        //   (without reading a V24 summary) and then completing the full load after the migration
+        //   runs is possible in this case, but not in the next case.
+        // - Migrating from V24 to V22: initially summaries are V24, but after the migration runs
+        //   they will be V22. If we used the "load full split after migration" approach with strict
+        //   V24 summaries, it would break when trying to read V22 summaries after the migration.
+        //
+        // Therefore we take the most flexible approach of reading _either_ a V22 or V24 summary and
+        // using this to load the split correctly the first time.
+        if let Some(split) = db.load_split()? {
             *db.split.write() = split;
 
             info!(
                 %split.slot,
-                split_state = ?split.state_root,
+                ?split.state_root,
+                ?split.block_root,
                 "Hot-Cold DB initialized"
             );
         }
@@ -395,14 +406,6 @@ impl<E: EthSpec> HotColdDB<E, BeaconNodeBackend<E>, BeaconNodeBackend<E>> {
             db.store_schema_version(CURRENT_SCHEMA_VERSION)?;
         }
 
-        // Load the full split after the migration to set the `split.block_root` to the correct
-        // value.
-        if let Some(split) = db.load_split()? {
-            *db.split.write() = split;
-
-            debug!(?split, "Hot-Cold DB split initialized");
-        }
-
         // TODO(tree-states): Here we can choose to prune advanced states to reclaim disk space. As
         // it's a foreground task there's no risk of race condition that can corrupt the DB.
         // Advanced states for invalid blocks that were never written to the DB, or descendants of
@@ -2764,14 +2767,14 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
             Some(mut split) => {
                 debug!(?split, "Loaded split partial");
                 // Load the hot state summary to get the block root.
-                let summary = self
-                    .load_hot_state_summary(&split.state_root)
+                let latest_block_root = self
+                    .load_block_root_from_summary_any_version(&split.state_root)
                     .map_err(|e| Error::LoadHotStateSummaryForSplit(e.into()))?
                     .ok_or(HotColdDBError::MissingSplitState(
                         split.state_root,
                         split.slot,
                     ))?;
-                split.block_root = summary.latest_block_root;
+                split.block_root = latest_block_root;
                 Ok(Some(split))
             }
             None => Ok(None),
@@ -2801,6 +2804,31 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
             .map_err(|e| Error::LoadHotStateSummary(*state_root, e.into()))
     }
 
+    /// Load a hot state's summary in V22 format, given its root.
+    pub fn load_hot_state_summary_v22(
+        &self,
+        state_root: &Hash256,
+    ) -> Result<Option<HotStateSummaryV22>, Error> {
+        self.hot_db
+            .get(state_root)
+            .map_err(|e| Error::LoadHotStateSummary(*state_root, e.into()))
+    }
+
+    /// Load the latest block root for a hot state summary either in modern form, or V22 form.
+    ///
+    /// This function is required to open a V22 database for migration to V24, or vice versa.
+    pub fn load_block_root_from_summary_any_version(
+        &self,
+        state_root: &Hash256,
+    ) -> Result<Option<Hash256>, Error> {
+        self.load_hot_state_summary(state_root)
+            .map(|opt_summary| opt_summary.map(|summary| summary.latest_block_root))
+            .or_else(|_| {
+                self.load_hot_state_summary_v22(state_root)
+                    .map(|opt_summary| opt_summary.map(|summary| summary.latest_block_root))
+            })
+    }
+
     /// Load all hot state summaries present in the hot DB
     pub fn load_hot_state_summaries(&self) -> Result<Vec<(Hash256, HotStateSummary)>, Error> {
         self.hot_db
@@ -3689,6 +3717,30 @@ impl HotStateSummary {
     }
 }
 
+/// Legacy hot state summary used in schema V22 and before.
+///
+/// This can be deleted when we remove V22 support.
+#[derive(Debug, Clone, Copy, Encode, Decode)]
+pub struct HotStateSummaryV22 {
+    pub slot: Slot,
+    pub latest_block_root: Hash256,
+    pub epoch_boundary_state_root: Hash256,
+}
+
+impl StoreItem for HotStateSummaryV22 {
+    fn db_column() -> DBColumn {
+        DBColumn::BeaconStateSummary
+    }
+
+    fn as_store_bytes(&self) -> Vec<u8> {
+        self.as_ssz_bytes()
+    }
+
+    fn from_store_bytes(bytes: &[u8]) -> Result<Self, Error> {
+        Ok(Self::from_ssz_bytes(bytes)?)
+    }
+}
+
 /// Struct for summarising a state in the freezer database.
 #[derive(Debug, Clone, Copy, Default, Encode, Decode)]
 pub(crate) struct ColdStateSummary {
