Implement basic validator custody framework (no backfill)

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -58,12 +58,14 @@ use crate::observed_data_sidecars::ObservedDataSidecars;
 use crate::observed_operations::{ObservationOutcome, ObservedOperations};
 use crate::observed_slashable::ObservedSlashable;
 use crate::persisted_beacon_chain::PersistedBeaconChain;
+use crate::persisted_custody::{clear_custody_context, persist_custody_context};
 use crate::persisted_fork_choice::PersistedForkChoice;
 use crate::pre_finalization_cache::PreFinalizationBlockCache;
 use crate::shuffling_cache::{BlockShufflingIds, ShufflingCache};
 use crate::sync_committee_verification::{
     Error as SyncCommitteeError, VerifiedSyncCommitteeMessage, VerifiedSyncContribution,
 };
+use crate::validator_custody::CustodyContextSsz;
 use crate::validator_monitor::{
     get_slot_delay_ms, timestamp_now, ValidatorMonitor,
     HISTORIC_EPOCHS as VALIDATOR_MONITOR_HISTORIC_EPOCHS,
@@ -668,6 +670,34 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         Ok(())
     }
 
+    /// Persists the custody information to disk.
+    pub fn persist_custody_context(&self) -> Result<(), Error> {
+        let custody_context: CustodyContextSsz = self
+            .data_availability_checker
+            .custody_context()
+            .as_ref()
+            .into();
+        debug!(?custody_context, "Persisting custody context to store");
+
+        if let Err(e) =
+            clear_custody_context::<T::EthSpec, T::HotStore, T::ColdStore>(self.store.clone())
+        {
+            error!(error = ?e, "Failed to clear old custody context");
+        }
+
+        match persist_custody_context::<T::EthSpec, T::HotStore, T::ColdStore>(
+            self.store.clone(),
+            custody_context,
+        ) {
+            Ok(_) => info!("Saved custody state"),
+            Err(e) => error!(
+                error = ?e,
+                "Failed to persist custody context on drop"
+            ),
+        }
+        Ok(())
+    }
+
     /// Returns the slot _right now_ according to `self.slot_clock`. Returns `Err` if the slot is
     /// unavailable.
     ///
@@ -7187,6 +7217,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
 impl<T: BeaconChainTypes> Drop for BeaconChain<T> {
     fn drop(&mut self) {
         let drop = || -> Result<(), Error> {
+            self.persist_custody_context()?;
             self.persist_fork_choice()?;
             self.persist_op_pool()?;
             self.persist_eth1_cache()

beacon_node/beacon_chain/src/builder.rs

@@ -13,10 +13,12 @@ use crate::light_client_server_cache::LightClientServerCache;
 use crate::migrate::{BackgroundMigrator, MigratorConfig};
 use crate::observed_data_sidecars::ObservedDataSidecars;
 use crate::persisted_beacon_chain::PersistedBeaconChain;
+use crate::persisted_custody::load_custody_context;
 use crate::shuffling_cache::{BlockShufflingIds, ShufflingCache};
 use crate::validator_monitor::{ValidatorMonitor, ValidatorMonitorConfig};
 use crate::validator_pubkey_cache::ValidatorPubkeyCache;
 use crate::ChainConfig;
+use crate::CustodyContext;
 use crate::{
     BeaconChain, BeaconChainTypes, BeaconForkChoiceStore, BeaconSnapshot, Eth1Chain,
     Eth1ChainBackend, ServerSentEventHandler,
@@ -926,6 +928,20 @@ where
             }
         };
 
+        // Load the persisted custody context from the db and initialize
+        // the context for this run
+        let custody_context = if let Some(custody) =
+            load_custody_context::<E, THotStore, TColdStore>(store.clone())
+        {
+            Arc::new(CustodyContext::new_from_persisted_custody_context(
+                custody,
+                self.import_all_data_columns,
+            ))
+        } else {
+            Arc::new(CustodyContext::new(self.import_all_data_columns))
+        };
+        debug!(?custody_context, "Loading persisted custody context");
+
         let beacon_chain = BeaconChain {
             spec: self.spec.clone(),
             config: self.chain_config,
@@ -999,8 +1015,14 @@ where
             validator_monitor: RwLock::new(validator_monitor),
             genesis_backfill_slot,
             data_availability_checker: Arc::new(
-                DataAvailabilityChecker::new(slot_clock, self.kzg.clone(), store, self.spec)
-                    .map_err(|e| format!("Error initializing DataAvailabilityChecker: {:?}", e))?,
+                DataAvailabilityChecker::new(
+                    slot_clock,
+                    self.kzg.clone(),
+                    store,
+                    custody_context,
+                    self.spec,
+                )
+                .map_err(|e| format!("Error initializing DataAvailabilityChecker: {:?}", e))?,
             ),
             kzg: self.kzg.clone(),
             rng: Arc::new(Mutex::new(rng)),

beacon_node/beacon_chain/src/data_availability_checker.rs

@@ -5,7 +5,7 @@ use crate::block_verification_types::{
 use crate::data_availability_checker::overflow_lru_cache::{
     DataAvailabilityCheckerInner, ReconstructColumnsDecision,
 };
-use crate::{metrics, BeaconChain, BeaconChainTypes, BeaconStore};
+use crate::{metrics, BeaconChain, BeaconChainTypes, BeaconStore, CustodyContext};
 use kzg::Kzg;
 use slot_clock::SlotClock;
 use std::fmt;
@@ -74,6 +74,7 @@ pub struct DataAvailabilityChecker<T: BeaconChainTypes> {
     availability_cache: Arc<DataAvailabilityCheckerInner<T>>,
     slot_clock: T::SlotClock,
     kzg: Arc<Kzg>,
+    custody_context: Arc<CustodyContext>,
     spec: Arc<ChainSpec>,
 }
 
@@ -111,17 +112,23 @@ impl<T: BeaconChainTypes> DataAvailabilityChecker<T> {
         slot_clock: T::SlotClock,
         kzg: Arc<Kzg>,
         store: BeaconStore<T>,
+        custody_context: Arc<CustodyContext>,
         spec: Arc<ChainSpec>,
     ) -> Result<Self, AvailabilityCheckError> {
         let inner = DataAvailabilityCheckerInner::new(OVERFLOW_LRU_CAPACITY, store, spec.clone())?;
         Ok(Self {
             availability_cache: Arc::new(inner),
             slot_clock,
             kzg,
+            custody_context,
             spec,
         })
     }
 
+    pub fn custody_context(&self) -> Arc<CustodyContext> {
+        self.custody_context.clone()
+    }
+
     /// Checks if the block root is currenlty in the availability cache awaiting import because
     /// of missing components.
     pub fn get_execution_valid_block(

beacon_node/beacon_chain/src/lib.rs

@@ -48,6 +48,7 @@ pub mod observed_data_sidecars;
 pub mod observed_operations;
 mod observed_slashable;
 mod persisted_beacon_chain;
+pub mod persisted_custody;
 mod persisted_fork_choice;
 mod pre_finalization_cache;
 pub mod proposer_prep_service;
@@ -59,6 +60,7 @@ pub mod summaries_dag;
 pub mod sync_committee_rewards;
 pub mod sync_committee_verification;
 pub mod test_utils;
+pub mod validator_custody;
 pub mod validator_monitor;
 pub mod validator_pubkey_cache;
 
@@ -100,3 +102,4 @@ pub use state_processing::per_block_processing::errors::{
 };
 pub use store;
 pub use types;
+pub use validator_custody::CustodyContext;

beacon_node/beacon_chain/src/persisted_custody.rs

@@ -0,0 +1,53 @@
+use crate::validator_custody::CustodyContextSsz;
+use ssz::{Decode, Encode};
+use std::sync::Arc;
+use store::{DBColumn, Error as StoreError, HotColdDB, ItemStore, StoreItem};
+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 fn load_custody_context<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>(
+    store: Arc<HotColdDB<E, Hot, Cold>>,
+) -> Option<CustodyContextSsz> {
+    let res: Result<Option<PersistedCustody>, _> =
+        store.get_item::<PersistedCustody>(&CUSTODY_DB_KEY);
+    // Load context from the store
+    match res {
+        Ok(Some(c)) => Some(c.0),
+        _ => None,
+    }
+}
+
+/// Attempt to persist the custody context object to `self.store`.
+pub fn persist_custody_context<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>(
+    store: Arc<HotColdDB<E, Hot, Cold>>,
+    custody_context: CustodyContextSsz,
+) -> Result<(), store::Error> {
+    store.put_item(&CUSTODY_DB_KEY, &PersistedCustody(custody_context))
+}
+
+/// Attempts to clear any custody context entries.
+pub fn clear_custody_context<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>(
+    store: Arc<HotColdDB<E, Hot, Cold>>,
+) -> Result<(), store::Error> {
+    store.hot_db.delete::<PersistedCustody>(&CUSTODY_DB_KEY)
+}
+
+impl StoreItem for PersistedCustody {
+    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, StoreError> {
+        let custody_context = CustodyContextSsz::from_ssz_bytes(bytes)?;
+
+        Ok(PersistedCustody(custody_context))
+    }
+}