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>,
-    pub best_justified_checkpoint: Checkpoint,
     #[superstruct(only(V11))]
+    pub best_justified_checkpoint: Checkpoint,
+    #[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)])
+}

beacon_node/http_api/src/lib.rs

@@ -2202,12 +2202,8 @@ pub fn serve<T: BeaconChainTypes>(
                                 .parent
                                 .and_then(|index| proto_array.nodes.get(index))
                                 .map(|parent| parent.root),
-                            justified_epoch: node
-                                .justified_checkpoint
-                                .map(|checkpoint| checkpoint.epoch),
-                            finalized_epoch: node
-                                .finalized_checkpoint
-                                .map(|checkpoint| checkpoint.epoch),
+                            justified_epoch: node.justified_checkpoint.epoch,
+                            finalized_epoch: node.finalized_checkpoint.epoch,
                             weight: node.weight,
                             validity: execution_status,
                             execution_block_hash: node

beacon_node/http_api/tests/tests.rs

@@ -1956,8 +1956,8 @@ impl ApiTester {
                         .parent
                         .and_then(|index| expected_proto_array.nodes.get(index))
                         .map(|parent| parent.root),
-                    justified_epoch: node.justified_checkpoint.map(|checkpoint| checkpoint.epoch),
-                    finalized_epoch: node.finalized_checkpoint.map(|checkpoint| checkpoint.epoch),
+                    justified_epoch: node.justified_checkpoint.epoch,
+                    finalized_epoch: node.finalized_checkpoint.epoch,
                     weight: node.weight,
                     validity: execution_status,
                     execution_block_hash: node

beacon_node/store/src/metadata.rs

@@ -4,7 +4,7 @@ use ssz::{Decode, Encode};
 use ssz_derive::{Decode, Encode};
 use types::{Checkpoint, Hash256, Slot};
 
-pub const CURRENT_SCHEMA_VERSION: SchemaVersion = SchemaVersion(16);
+pub const CURRENT_SCHEMA_VERSION: SchemaVersion = SchemaVersion(17);
 
 // All the keys that get stored under the `BeaconMeta` column.
 //

common/eth2/src/types.rs

@@ -1229,8 +1229,8 @@ pub struct ForkChoiceNode {
     pub slot: Slot,
     pub block_root: Hash256,
     pub parent_root: Option<Hash256>,
-    pub justified_epoch: Option<Epoch>,
-    pub finalized_epoch: Option<Epoch>,
+    pub justified_epoch: Epoch,
+    pub finalized_epoch: Epoch,
     #[serde(with = "serde_utils::quoted_u64")]
     pub weight: u64,
     pub validity: Option<String>,

consensus/proto_array/Cargo.toml

@@ -16,3 +16,4 @@ serde = "1.0.116"
 serde_derive = "1.0.116"
 serde_yaml = "0.8.13"
 safe_arith = { path = "../safe_arith" }
+superstruct = "0.5.0"

consensus/proto_array/src/error.rs

@@ -14,6 +14,8 @@ pub enum Error {
     InvalidBestDescendant(usize),
     InvalidParentDelta(usize),
     InvalidNodeDelta(usize),
+    MissingJustifiedCheckpoint,
+    MissingFinalizedCheckpoint,
     DeltaOverflow(usize),
     ProposerBoostOverflow(usize),
     ReOrgThresholdOverflow,
@@ -67,6 +69,6 @@ pub struct InvalidBestNodeInfo {
     pub justified_checkpoint: Checkpoint,
     pub finalized_checkpoint: Checkpoint,
     pub head_root: Hash256,
-    pub head_justified_checkpoint: Option<Checkpoint>,
-    pub head_finalized_checkpoint: Option<Checkpoint>,
+    pub head_justified_checkpoint: Checkpoint,
+    pub head_finalized_checkpoint: Checkpoint,
 }

consensus/proto_array/src/lib.rs

@@ -16,5 +16,5 @@ pub use error::Error;
 pub mod core {
     pub use super::proto_array::{ProposerBoost, ProtoArray, ProtoNode};
     pub use super::proto_array_fork_choice::VoteTracker;
-    pub use super::ssz_container::SszContainer;
+    pub use super::ssz_container::{SszContainer, SszContainerV16, SszContainerV17};
 }