Fix schema migrations

beacon_node/beacon_chain/src/persisted_fork_choice.rs

@@ -22,26 +22,37 @@ pub struct PersistedForkChoice {
     pub fork_choice_store: PersistedForkChoiceStoreV28,
 }
 
-macro_rules! impl_store_item {
+impl PersistedForkChoiceV28 {
-    ($type:ty) => {
+    pub fn from_bytes(bytes: &[u8], store_config: &StoreConfig) -> Result<Self, Error> {
-        impl store::StoreItem for $type {
+        let decompressed_bytes = store_config
-            fn db_column() -> DBColumn {
+            .decompress_bytes(bytes)
-                DBColumn::ForkChoice
+            .map_err(Error::Compression)?;
+        Self::from_ssz_bytes(&decompressed_bytes).map_err(Into::into)
     }
 
-            fn as_store_bytes(&self) -> Vec<u8> {
+    pub fn as_bytes(&self, store_config: &StoreConfig) -> Result<Vec<u8>, Error> {
-                self.as_ssz_bytes()
+        let encode_timer = metrics::start_timer(&metrics::FORK_CHOICE_ENCODE_TIMES);
+        let ssz_bytes = self.as_ssz_bytes();
+        drop(encode_timer);
+
+        let _compress_timer = metrics::start_timer(&metrics::FORK_CHOICE_COMPRESS_TIMES);
+        store_config
+            .compress_bytes(&ssz_bytes)
+            .map_err(Error::Compression)
     }
 
-            fn from_store_bytes(bytes: &[u8]) -> std::result::Result<Self, Error> {
+    pub fn as_kv_store_op(
-                Self::from_ssz_bytes(bytes).map_err(Into::into)
+        &self,
+        key: Hash256,
+        store_config: &StoreConfig,
+    ) -> Result<KeyValueStoreOp, Error> {
+        Ok(KeyValueStoreOp::PutKeyValue(
+            DBColumn::ForkChoice,
+            key.as_slice().to_vec(),
+            self.as_bytes(store_config)?,
+        ))
     }
 }
-    };
-}
-
-impl_store_item!(PersistedForkChoiceV28);
-impl_store_item!(PersistedForkChoiceV29);
 
 impl PersistedForkChoiceV29 {
     pub fn from_bytes(bytes: &[u8], store_config: &StoreConfig) -> Result<Self, Error> {
@@ -83,3 +94,12 @@ impl From<PersistedForkChoiceV28> for PersistedForkChoiceV29 {
         }
     }
 }
+
+impl From<PersistedForkChoiceV29> for PersistedForkChoiceV28 {
+    fn from(v29: PersistedForkChoiceV29) -> Self {
+        Self {
+            fork_choice_v28: v29.fork_choice.into(),
+            fork_choice_store: v29.fork_choice_store,
+        }
+    }
+}

beacon_node/beacon_chain/src/schema_change.rs

@@ -22,13 +22,13 @@ pub fn migrate_schema<T: BeaconChainTypes>(
         (_, _) if from == to && to == CURRENT_SCHEMA_VERSION => Ok(()),
         // Upgrade from v28 to v29.
         (SchemaVersion(28), SchemaVersion(29)) => {
-            upgrade_to_v29::<T>(&db)?;
+            let ops = upgrade_to_v29::<T>(&db)?;
-            db.store_schema_version_atomically(to, vec![])
+            db.store_schema_version_atomically(to, ops)
         }
         // Downgrade from v29 to v28.
         (SchemaVersion(29), SchemaVersion(28)) => {
-            downgrade_from_v29::<T>(&db)?;
+            let ops = downgrade_from_v29::<T>(&db)?;
-            db.store_schema_version_atomically(to, vec![])
+            db.store_schema_version_atomically(to, ops)
         }
         // Anything else is an error.
         (_, _) => Err(HotColdDBError::UnsupportedSchemaVersion {

beacon_node/beacon_chain/src/schema_change/migration_schema_v29.rs

@@ -1,12 +1,8 @@
-use crate::beacon_chain::BeaconChainTypes;
+use crate::beacon_chain::{BeaconChainTypes, FORK_CHOICE_DB_KEY};
 use crate::persisted_fork_choice::{PersistedForkChoiceV28, PersistedForkChoiceV29};
-use ssz::Decode;
 use store::hot_cold_store::HotColdDB;
 use store::{DBColumn, Error as StoreError, KeyValueStore, KeyValueStoreOp};
-use types::{EthSpec, Hash256};
+use types::EthSpec;
-
-/// The key used to store the fork choice in the database.
-const FORK_CHOICE_DB_KEY: Hash256 = Hash256::ZERO;
 
 /// Upgrade from schema v28 to v29.
 ///
@@ -14,24 +10,25 @@ const FORK_CHOICE_DB_KEY: Hash256 = Hash256::ZERO;
 ///   virtual tree walk).
 /// - Fails if the persisted fork choice contains any V17 (pre-Gloas) proto
 ///   nodes at or after the Gloas fork slot.
+///
+/// Returns a list of store ops to be applied atomically with the schema version write.
 pub fn upgrade_to_v29<T: BeaconChainTypes>(
     db: &HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>,
-) -> Result<(), StoreError> {
+) -> Result<Vec<KeyValueStoreOp>, StoreError> {
     let gloas_fork_slot = db
         .spec
         .gloas_fork_epoch
         .map(|epoch| epoch.start_slot(T::EthSpec::slots_per_epoch()));
 
-    // Load the persisted fork choice (v28 format, uncompressed SSZ).
+    // Load the persisted fork choice (v28 format).
     let Some(fc_bytes) = db
         .hot_db
         .get_bytes(DBColumn::ForkChoice, FORK_CHOICE_DB_KEY.as_slice())?
     else {
-        return Ok(());
+        return Ok(vec![]);
     };
 
-    let mut persisted_v28 =
+    let persisted_v28 = PersistedForkChoiceV28::from_bytes(&fc_bytes, db.get_config())?;
-        PersistedForkChoiceV28::from_ssz_bytes(&fc_bytes).map_err(StoreError::SszDecodeError)?;
 
     // Check for V17 nodes at/after the Gloas fork slot.
     if let Some(gloas_fork_slot) = gloas_fork_slot {
@@ -52,39 +49,30 @@ pub fn upgrade_to_v29<T: BeaconChainTypes>(
         }
     }
 
-    // Clear best_child/best_descendant — replaced by the virtual tree walk.
+    // Convert to v29 and encode.
-    for node in &mut persisted_v28.fork_choice_v28.proto_array_v28.nodes {
-        node.best_child = None;
-        node.best_descendant = None;
-    }
-
-    // Convert to v29 and write back.
     let persisted_v29 = PersistedForkChoiceV29::from(persisted_v28);
-    let fc_bytes = persisted_v29
-        .as_bytes(db.get_config())
-        .map_err(|e| StoreError::MigrationError(format!("failed to encode v29: {:?}", e)))?;
-    db.hot_db.do_atomically(vec![KeyValueStoreOp::PutKeyValue(
-        DBColumn::ForkChoice,
-        FORK_CHOICE_DB_KEY.as_slice().to_vec(),
-        fc_bytes,
-    )])?;
 
-    Ok(())
+    Ok(vec![
+        persisted_v29.as_kv_store_op(FORK_CHOICE_DB_KEY, db.get_config())?,
+    ])
 }
 
-/// Downgrade from schema v29 to v28 (no-op).
+/// Downgrade from schema v29 to v28.
 ///
+/// Converts the persisted fork choice from V29 format back to V28.
 /// Fails if the persisted fork choice contains any V29 proto nodes, as these contain
 /// payload-specific fields that cannot be losslessly converted back to V17 format.
+///
+/// Returns a list of store ops to be applied atomically with the schema version write.
 pub fn downgrade_from_v29<T: BeaconChainTypes>(
     db: &HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>,
-) -> Result<(), StoreError> {
+) -> Result<Vec<KeyValueStoreOp>, StoreError> {
     // Load the persisted fork choice (v29 format, compressed).
     let Some(fc_bytes) = db
         .hot_db
         .get_bytes(DBColumn::ForkChoice, FORK_CHOICE_DB_KEY.as_slice())?
     else {
-        return Ok(());
+        return Ok(vec![]);
     };
 
     let persisted_v29 =
@@ -111,5 +99,10 @@ pub fn downgrade_from_v29<T: BeaconChainTypes>(
         ));
     }
 
-    Ok(())
+    // Convert to v28 and encode.
+    let persisted_v28 = PersistedForkChoiceV28::from(persisted_v29);
+
+    Ok(vec![
+        persisted_v28.as_kv_store_op(FORK_CHOICE_DB_KEY, db.get_config())?,
+    ])
 }

consensus/fork_choice/src/fork_choice.rs

@@ -1840,6 +1840,15 @@ impl From<PersistedForkChoiceV28> for PersistedForkChoiceV29 {
     }
 }
 
+impl From<PersistedForkChoiceV29> for PersistedForkChoiceV28 {
+    fn from(v29: PersistedForkChoiceV29) -> Self {
+        Self {
+            proto_array_v28: v29.proto_array.into(),
+            queued_attestations: v29.queued_attestations,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use types::MainnetEthSpec;