DB migration for fork choice cleanup (#4265)

## Issue Addressed

#4233

## Proposed Changes

Remove the `best_justified_checkpoint` from the `PersistedForkChoiceStore` type as it is now unused.
Additionally, remove the `Option`'s wrapping the `justified_checkpoint` and `finalized_checkpoint` fields on `ProtoNode` which were only present to facilitate a previous migration.

Include the necessary code to facilitate the migration to a new DB schema.

beacon_node/beacon_chain/src/beacon_fork_choice_store.rs

@@ -218,7 +218,6 @@ where
             finalized_checkpoint: self.finalized_checkpoint,
             justified_checkpoint: self.justified_checkpoint,
             justified_balances: self.justified_balances.effective_balances.clone(),
-            best_justified_checkpoint: JUNK_BEST_JUSTIFIED_CHECKPOINT,
             unrealized_justified_checkpoint: self.unrealized_justified_checkpoint,
             unrealized_finalized_checkpoint: self.unrealized_finalized_checkpoint,
             proposer_boost_root: self.proposer_boost_root,
@@ -355,24 +354,62 @@ where
     }
 }
 
+pub type PersistedForkChoiceStore = PersistedForkChoiceStoreV17;
+
 /// A container which allows persisting the `BeaconForkChoiceStore` to the on-disk database.
-#[superstruct(variants(V11), variant_attributes(derive(Encode, Decode)), no_enum)]
+#[superstruct(
+    variants(V11, V17),
+    variant_attributes(derive(Encode, Decode)),
+    no_enum
+)]
 pub struct PersistedForkChoiceStore {
-    #[superstruct(only(V11))]
+    #[superstruct(only(V11, V17))]
     pub balances_cache: BalancesCacheV8,
     pub time: Slot,
     pub finalized_checkpoint: Checkpoint,
     pub justified_checkpoint: Checkpoint,
     pub justified_balances: Vec<u64>,
+    #[superstruct(only(V11))]
     pub best_justified_checkpoint: Checkpoint,
-    #[superstruct(only(V11))]
+    #[superstruct(only(V11, V17))]
     pub unrealized_justified_checkpoint: Checkpoint,
-    #[superstruct(only(V11))]
+    #[superstruct(only(V11, V17))]
     pub unrealized_finalized_checkpoint: Checkpoint,
-    #[superstruct(only(V11))]
+    #[superstruct(only(V11, V17))]
     pub proposer_boost_root: Hash256,
-    #[superstruct(only(V11))]
+    #[superstruct(only(V11, V17))]
     pub equivocating_indices: BTreeSet<u64>,
 }
 
-pub type PersistedForkChoiceStore = PersistedForkChoiceStoreV11;
+impl Into<PersistedForkChoiceStore> for PersistedForkChoiceStoreV11 {
+    fn into(self) -> PersistedForkChoiceStore {
+        PersistedForkChoiceStore {
+            balances_cache: self.balances_cache,
+            time: self.time,
+            finalized_checkpoint: self.finalized_checkpoint,
+            justified_checkpoint: self.justified_checkpoint,
+            justified_balances: self.justified_balances,
+            unrealized_justified_checkpoint: self.unrealized_justified_checkpoint,
+            unrealized_finalized_checkpoint: self.unrealized_finalized_checkpoint,
+            proposer_boost_root: self.proposer_boost_root,
+            equivocating_indices: self.equivocating_indices,
+        }
+    }
+}
+
+impl Into<PersistedForkChoiceStoreV11> for PersistedForkChoiceStore {
+    fn into(self) -> PersistedForkChoiceStoreV11 {
+        PersistedForkChoiceStoreV11 {
+            balances_cache: self.balances_cache,
+            time: self.time,
+            finalized_checkpoint: self.finalized_checkpoint,
+            justified_checkpoint: self.justified_checkpoint,
+            justified_balances: self.justified_balances,
+            best_justified_checkpoint: JUNK_BEST_JUSTIFIED_CHECKPOINT,
+            unrealized_justified_checkpoint: self.unrealized_justified_checkpoint,
+            unrealized_finalized_checkpoint: self.unrealized_finalized_checkpoint,
+            proposer_boost_root: self.proposer_boost_root,
+            equivocating_indices: self.equivocating_indices,
+        }
+    }
+}

beacon_node/beacon_chain/src/persisted_fork_choice.rs

@@ -1,17 +1,41 @@
-use crate::beacon_fork_choice_store::PersistedForkChoiceStoreV11;
+use crate::beacon_fork_choice_store::{PersistedForkChoiceStoreV11, PersistedForkChoiceStoreV17};
 use ssz::{Decode, Encode};
 use ssz_derive::{Decode, Encode};
 use store::{DBColumn, Error, StoreItem};
 use superstruct::superstruct;
 
 // If adding a new version you should update this type alias and fix the breakages.
-pub type PersistedForkChoice = PersistedForkChoiceV11;
+pub type PersistedForkChoice = PersistedForkChoiceV17;
 
-#[superstruct(variants(V11), variant_attributes(derive(Encode, Decode)), no_enum)]
+#[superstruct(
+    variants(V11, V17),
+    variant_attributes(derive(Encode, Decode)),
+    no_enum
+)]
 pub struct PersistedForkChoice {
     pub fork_choice: fork_choice::PersistedForkChoice,
     #[superstruct(only(V11))]
     pub fork_choice_store: PersistedForkChoiceStoreV11,
+    #[superstruct(only(V17))]
+    pub fork_choice_store: PersistedForkChoiceStoreV17,
+}
+
+impl Into<PersistedForkChoice> for PersistedForkChoiceV11 {
+    fn into(self) -> PersistedForkChoice {
+        PersistedForkChoice {
+            fork_choice: self.fork_choice,
+            fork_choice_store: self.fork_choice_store.into(),
+        }
+    }
+}
+
+impl Into<PersistedForkChoiceV11> for PersistedForkChoice {
+    fn into(self) -> PersistedForkChoiceV11 {
+        PersistedForkChoiceV11 {
+            fork_choice: self.fork_choice,
+            fork_choice_store: self.fork_choice_store.into(),
+        }
+    }
 }
 
 macro_rules! impl_store_item {
@@ -33,3 +57,4 @@ macro_rules! impl_store_item {
 }
 
 impl_store_item!(PersistedForkChoiceV11);
+impl_store_item!(PersistedForkChoiceV17);

beacon_node/beacon_chain/src/schema_change.rs

@@ -4,6 +4,7 @@ mod migration_schema_v13;
 mod migration_schema_v14;
 mod migration_schema_v15;
 mod migration_schema_v16;
+mod migration_schema_v17;
 
 use crate::beacon_chain::{BeaconChainTypes, ETH1_CACHE_DB_KEY};
 use crate::eth1_chain::SszEth1;
@@ -141,6 +142,14 @@ pub fn migrate_schema<T: BeaconChainTypes>(
             let ops = migration_schema_v16::downgrade_from_v16::<T>(db.clone(), log)?;
             db.store_schema_version_atomically(to, ops)
         }
+        (SchemaVersion(16), SchemaVersion(17)) => {
+            let ops = migration_schema_v17::upgrade_to_v17::<T>(db.clone(), log)?;
+            db.store_schema_version_atomically(to, ops)
+        }
+        (SchemaVersion(17), SchemaVersion(16)) => {
+            let ops = migration_schema_v17::downgrade_from_v17::<T>(db.clone(), log)?;
+            db.store_schema_version_atomically(to, ops)
+        }
         // Anything else is an error.
         (_, _) => Err(HotColdDBError::UnsupportedSchemaVersion {
             target_version: to,

beacon_node/beacon_chain/src/schema_change/migration_schema_v17.rs

@@ -0,0 +1,88 @@
+use crate::beacon_chain::{BeaconChainTypes, FORK_CHOICE_DB_KEY};
+use crate::persisted_fork_choice::{PersistedForkChoiceV11, PersistedForkChoiceV17};
+use proto_array::core::{SszContainerV16, SszContainerV17};
+use slog::{debug, Logger};
+use ssz::{Decode, Encode};
+use std::sync::Arc;
+use store::{Error, HotColdDB, KeyValueStoreOp, StoreItem};
+
+pub fn upgrade_fork_choice(
+    mut fork_choice: PersistedForkChoiceV11,
+) -> Result<PersistedForkChoiceV17, Error> {
+    let ssz_container_v16 = SszContainerV16::from_ssz_bytes(
+        &fork_choice.fork_choice.proto_array_bytes,
+    )
+    .map_err(|e| {
+        Error::SchemaMigrationError(format!(
+            "Failed to decode ProtoArrayForkChoice during schema migration: {:?}",
+            e
+        ))
+    })?;
+
+    let ssz_container_v17: SszContainerV17 = ssz_container_v16.try_into().map_err(|e| {
+        Error::SchemaMigrationError(format!(
+            "Missing checkpoint during schema migration: {:?}",
+            e
+        ))
+    })?;
+    fork_choice.fork_choice.proto_array_bytes = ssz_container_v17.as_ssz_bytes();
+
+    Ok(fork_choice.into())
+}
+
+pub fn downgrade_fork_choice(
+    mut fork_choice: PersistedForkChoiceV17,
+) -> Result<PersistedForkChoiceV11, Error> {
+    let ssz_container_v17 = SszContainerV17::from_ssz_bytes(
+        &fork_choice.fork_choice.proto_array_bytes,
+    )
+    .map_err(|e| {
+        Error::SchemaMigrationError(format!(
+            "Failed to decode ProtoArrayForkChoice during schema migration: {:?}",
+            e
+        ))
+    })?;
+
+    let ssz_container_v16: SszContainerV16 = ssz_container_v17.into();
+    fork_choice.fork_choice.proto_array_bytes = ssz_container_v16.as_ssz_bytes();
+
+    Ok(fork_choice.into())
+}
+
+pub fn upgrade_to_v17<T: BeaconChainTypes>(
+    db: Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
+    log: Logger,
+) -> Result<Vec<KeyValueStoreOp>, Error> {
+    // Get persisted_fork_choice.
+    let v11 = db
+        .get_item::<PersistedForkChoiceV11>(&FORK_CHOICE_DB_KEY)?
+        .ok_or_else(|| Error::SchemaMigrationError("fork choice missing from database".into()))?;
+
+    let v17 = upgrade_fork_choice(v11)?;
+
+    debug!(
+        log,
+        "Removing unused best_justified_checkpoint from fork choice store."
+    );
+
+    Ok(vec![v17.as_kv_store_op(FORK_CHOICE_DB_KEY)])
+}
+
+pub fn downgrade_from_v17<T: BeaconChainTypes>(
+    db: Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
+    log: Logger,
+) -> Result<Vec<KeyValueStoreOp>, Error> {
+    // Get persisted_fork_choice.
+    let v17 = db
+        .get_item::<PersistedForkChoiceV17>(&FORK_CHOICE_DB_KEY)?
+        .ok_or_else(|| Error::SchemaMigrationError("fork choice missing from database".into()))?;
+
+    let v11 = downgrade_fork_choice(v17)?;
+
+    debug!(
+        log,
+        "Adding junk best_justified_checkpoint to fork choice store."
+    );
+
+    Ok(vec![v11.as_kv_store_op(FORK_CHOICE_DB_KEY)])
+}