Custody persist fix (#7661)

N/A


  Persist the epoch -> cgc values. This is to ensure that `ValidatorRegistrations::latest_validator_custody_requirement` always returns a `Some` value post restart assuming the `epoch_validator_custody_requirements` map has been updated in the previous runs.

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -654,6 +654,10 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 
     /// Persists the custody information to disk.
     pub fn persist_custody_context(&self) -> Result<(), Error> {
+        if !self.spec.is_peer_das_scheduled() {
+            return Ok(());
+        }
+
         let custody_context: CustodyContextSsz = self
             .data_availability_checker
             .custody_context()

beacon_node/beacon_chain/src/persisted_custody.rs

@@ -7,7 +7,7 @@ use types::{EthSpec, Hash256};
 /// 32-byte key for accessing the `CustodyContext`. All zero because `CustodyContext` has its own column.
 pub const CUSTODY_DB_KEY: Hash256 = Hash256::ZERO;
 
-pub struct PersistedCustody(CustodyContextSsz);
+pub struct PersistedCustody(pub CustodyContextSsz);
 
 pub fn load_custody_context<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>(
     store: Arc<HotColdDB<E, Hot, Cold>>,

beacon_node/beacon_chain/src/schema_change.rs

@@ -2,6 +2,7 @@
 mod migration_schema_v23;
 mod migration_schema_v24;
 mod migration_schema_v25;
+mod migration_schema_v26;
 
 use crate::beacon_chain::BeaconChainTypes;
 use std::sync::Arc;
@@ -58,6 +59,14 @@ pub fn migrate_schema<T: BeaconChainTypes>(
             let ops = migration_schema_v25::downgrade_from_v25()?;
             db.store_schema_version_atomically(to, ops)
         }
+        (SchemaVersion(25), SchemaVersion(26)) => {
+            let ops = migration_schema_v26::upgrade_to_v26::<T>(db.clone())?;
+            db.store_schema_version_atomically(to, ops)
+        }
+        (SchemaVersion(26), SchemaVersion(25)) => {
+            let ops = migration_schema_v26::downgrade_from_v26::<T>(db.clone())?;
+            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_v26.rs

@@ -0,0 +1,91 @@
+use crate::persisted_custody::{PersistedCustody, CUSTODY_DB_KEY};
+use crate::validator_custody::CustodyContextSsz;
+use crate::BeaconChainTypes;
+use ssz::{Decode, Encode};
+use ssz_derive::{Decode, Encode};
+use std::sync::Arc;
+use store::{DBColumn, Error, HotColdDB, KeyValueStoreOp, StoreItem};
+use tracing::info;
+
+#[derive(Debug, Encode, Decode, Clone)]
+pub(crate) struct CustodyContextSszV24 {
+    pub(crate) validator_custody_at_head: u64,
+    pub(crate) persisted_is_supernode: bool,
+}
+
+pub(crate) struct PersistedCustodyV24(CustodyContextSszV24);
+
+impl StoreItem for PersistedCustodyV24 {
+    fn db_column() -> DBColumn {
+        DBColumn::CustodyContext
+    }
+
+    fn as_store_bytes(&self) -> Vec<u8> {
+        self.0.as_ssz_bytes()
+    }
+
+    fn from_store_bytes(bytes: &[u8]) -> Result<Self, Error> {
+        let custody_context = CustodyContextSszV24::from_ssz_bytes(bytes)?;
+        Ok(PersistedCustodyV24(custody_context))
+    }
+}
+
+/// Upgrade the `CustodyContext` entry to v26.
+pub fn upgrade_to_v26<T: BeaconChainTypes>(
+    db: Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
+) -> Result<Vec<KeyValueStoreOp>, Error> {
+    let ops = if db.spec.is_peer_das_scheduled() {
+        match db.get_item::<PersistedCustodyV24>(&CUSTODY_DB_KEY) {
+            Ok(Some(PersistedCustodyV24(ssz_v24))) => {
+                info!("Migrating `CustodyContext` to v26 schema");
+                let custody_context_v2 = CustodyContextSsz {
+                    validator_custody_at_head: ssz_v24.validator_custody_at_head,
+                    persisted_is_supernode: ssz_v24.persisted_is_supernode,
+                    epoch_validator_custody_requirements: vec![],
+                };
+                vec![KeyValueStoreOp::PutKeyValue(
+                    DBColumn::CustodyContext,
+                    CUSTODY_DB_KEY.as_slice().to_vec(),
+                    PersistedCustody(custody_context_v2).as_store_bytes(),
+                )]
+            }
+            _ => {
+                vec![]
+            }
+        }
+    } else {
+        // Delete it from db if PeerDAS hasn't been scheduled
+        vec![KeyValueStoreOp::DeleteKey(
+            DBColumn::CustodyContext,
+            CUSTODY_DB_KEY.as_slice().to_vec(),
+        )]
+    };
+
+    Ok(ops)
+}
+
+pub fn downgrade_from_v26<T: BeaconChainTypes>(
+    db: Arc<HotColdDB<T::EthSpec, T::HotStore, T::ColdStore>>,
+) -> Result<Vec<KeyValueStoreOp>, Error> {
+    let res = db.get_item::<PersistedCustody>(&CUSTODY_DB_KEY);
+    let ops = match res {
+        Ok(Some(PersistedCustody(ssz_v26))) => {
+            info!("Migrating `CustodyContext` back from v26 schema");
+            let custody_context_v24 = CustodyContextSszV24 {
+                validator_custody_at_head: ssz_v26.validator_custody_at_head,
+                persisted_is_supernode: ssz_v26.persisted_is_supernode,
+            };
+            vec![KeyValueStoreOp::PutKeyValue(
+                DBColumn::CustodyContext,
+                CUSTODY_DB_KEY.as_slice().to_vec(),
+                PersistedCustodyV24(custody_context_v24).as_store_bytes(),
+            )]
+        }
+        _ => {
+            // no op if it's not on the db, as previous versions gracefully handle data missing from disk.
+            vec![]
+        }
+    };
+
+    Ok(ops)
+}

beacon_node/beacon_chain/src/validator_custody.rs

@@ -163,7 +163,13 @@ impl CustodyContext {
             validator_custody_count: AtomicU64::new(ssz_context.validator_custody_at_head),
             current_is_supernode: is_supernode,
             persisted_is_supernode: ssz_context.persisted_is_supernode,
-            validator_registrations: Default::default(),
+            validator_registrations: RwLock::new(ValidatorRegistrations {
+                validators: Default::default(),
+                epoch_validator_custody_requirements: ssz_context
+                    .epoch_validator_custody_requirements
+                    .into_iter()
+                    .collect(),
+            }),
         }
     }
 
@@ -263,8 +269,9 @@ pub struct CustodyCountChanged {
 /// The custody information that gets persisted across runs.
 #[derive(Debug, Encode, Decode, Clone)]
 pub struct CustodyContextSsz {
-    validator_custody_at_head: u64,
-    persisted_is_supernode: bool,
+    pub validator_custody_at_head: u64,
+    pub persisted_is_supernode: bool,
+    pub epoch_validator_custody_requirements: Vec<(Epoch, u64)>,
 }
 
 impl From<&CustodyContext> for CustodyContextSsz {
@@ -272,6 +279,13 @@ impl From<&CustodyContext> for CustodyContextSsz {
         CustodyContextSsz {
             validator_custody_at_head: context.validator_custody_count.load(Ordering::Relaxed),
             persisted_is_supernode: context.persisted_is_supernode,
+            epoch_validator_custody_requirements: context
+                .validator_registrations
+                .read()
+                .epoch_validator_custody_requirements
+                .iter()
+                .map(|(epoch, count)| (*epoch, *count))
+                .collect(),
         }
     }
 }

beacon_node/beacon_chain/tests/schema_stability.rs

@@ -88,7 +88,7 @@ async fn schema_stability() {
     check_db_columns();
     check_metadata_sizes(&store);
     check_op_pool(&store);
-    check_custody_context(&store);
+    check_custody_context(&store, &harness.spec);
     check_persisted_chain(&store);
 
     // Not covered here:
@@ -134,12 +134,13 @@ fn check_op_pool(store: &Store<E>) {
     assert_eq!(op_pool.as_store_bytes().len(), 28);
 }
 
-fn check_custody_context(store: &Store<E>) {
-    let custody_context = store
-        .get_item::<PersistedCustody>(&Hash256::ZERO)
-        .unwrap()
-        .unwrap();
-    assert_eq!(custody_context.as_store_bytes().len(), 9);
+fn check_custody_context(store: &Store<E>, spec: &ChainSpec) {
+    let custody_context_opt = store.get_item::<PersistedCustody>(&Hash256::ZERO).unwrap();
+    if spec.is_peer_das_scheduled() {
+        assert_eq!(custody_context_opt.unwrap().as_store_bytes().len(), 13);
+    } else {
+        assert!(custody_context_opt.is_none());
+    }
 }
 
 fn check_persisted_chain(store: &Store<E>) {

beacon_node/store/src/metadata.rs

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