Custody persist fix

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 CustodyContextSszV1 {
+    pub(crate) validator_custody_at_head: u64,
+    pub(crate) persisted_is_supernode: bool,
+}
+
+pub(crate) struct PersistedCustodyV1(CustodyContextSszV1);
+
+impl StoreItem for PersistedCustodyV1 {
+    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 = CustodyContextSszV1::from_ssz_bytes(bytes)?;
+        Ok(PersistedCustodyV1(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::<PersistedCustodyV1>(&CUSTODY_DB_KEY) {
+            Ok(Some(PersistedCustodyV1(v1))) => {
+                info!("Migrating `CustodyContext` to v26 schema");
+                let custody_context_v2 = CustodyContextSsz {
+                    validator_custody_at_head: v1.validator_custody_at_head,
+                    persisted_is_supernode: v1.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(v2))) => {
+            info!("Migrating `CustodyContext` back from v26 schema");
+            let custody_context_v1 = CustodyContextSszV1 {
+                validator_custody_at_head: v2.validator_custody_at_head,
+                persisted_is_supernode: v2.persisted_is_supernode,
+            };
+            vec![KeyValueStoreOp::PutKeyValue(
+                DBColumn::CustodyContext,
+                CUSTODY_DB_KEY.as_slice().to_vec(),
+                PersistedCustodyV1(custody_context_v1).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,15 +269,23 @@ 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 {
     fn from(context: &CustodyContext) -> Self {
         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/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.
 //