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.
     ///
@@ -2988,7 +3018,6 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     pub async fn verify_block_for_gossip(
         self: &Arc<Self>,
         block: Arc<SignedBeaconBlock<T::EthSpec>>,
-        custody_columns_count: usize,
     ) -> Result<GossipVerifiedBlock<T>, BlockError> {
         let chain = self.clone();
         self.task_executor
@@ -2998,7 +3027,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
                     let slot = block.slot();
                     let graffiti_string = block.message().body().graffiti().as_utf8_lossy();
 
-                    match GossipVerifiedBlock::new(block, &chain, custody_columns_count) {
+                    match GossipVerifiedBlock::new(block, &chain) {
                         Ok(verified) => {
                             let commitments_formatted = verified.block.commitments_formatted();
                             debug!(
@@ -7187,6 +7216,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/block_verification.rs

@@ -685,7 +685,6 @@ pub struct GossipVerifiedBlock<T: BeaconChainTypes> {
     pub block_root: Hash256,
     parent: Option<PreProcessingSnapshot<T::EthSpec>>,
     consensus_context: ConsensusContext<T::EthSpec>,
-    custody_columns_count: usize,
 }
 
 /// A wrapper around a `SignedBeaconBlock` that indicates that all signatures (except the deposit
@@ -721,7 +720,6 @@ pub trait IntoGossipVerifiedBlock<T: BeaconChainTypes>: Sized {
     fn into_gossip_verified_block(
         self,
         chain: &BeaconChain<T>,
-        custody_columns_count: usize,
     ) -> Result<GossipVerifiedBlock<T>, BlockError>;
     fn inner_block(&self) -> Arc<SignedBeaconBlock<T::EthSpec>>;
 }
@@ -730,7 +728,6 @@ impl<T: BeaconChainTypes> IntoGossipVerifiedBlock<T> for GossipVerifiedBlock<T>
     fn into_gossip_verified_block(
         self,
         _chain: &BeaconChain<T>,
-        _custody_columns_count: usize,
     ) -> Result<GossipVerifiedBlock<T>, BlockError> {
         Ok(self)
     }
@@ -743,9 +740,8 @@ impl<T: BeaconChainTypes> IntoGossipVerifiedBlock<T> for Arc<SignedBeaconBlock<T
     fn into_gossip_verified_block(
         self,
         chain: &BeaconChain<T>,
-        custody_columns_count: usize,
     ) -> Result<GossipVerifiedBlock<T>, BlockError> {
-        GossipVerifiedBlock::new(self, chain, custody_columns_count)
+        GossipVerifiedBlock::new(self, chain)
     }
 
     fn inner_block(&self) -> Arc<SignedBeaconBlock<T::EthSpec>> {
@@ -821,7 +817,6 @@ impl<T: BeaconChainTypes> GossipVerifiedBlock<T> {
     pub fn new(
         block: Arc<SignedBeaconBlock<T::EthSpec>>,
         chain: &BeaconChain<T>,
-        custody_columns_count: usize,
     ) -> Result<Self, BlockError> {
         // If the block is valid for gossip we don't supply it to the slasher here because
         // we assume it will be transformed into a fully verified block. We *do* need to supply
@@ -831,22 +826,19 @@ impl<T: BeaconChainTypes> GossipVerifiedBlock<T> {
         // The `SignedBeaconBlock` and `SignedBeaconBlockHeader` have the same canonical root,
         // but it's way quicker to calculate root of the header since the hash of the tree rooted
         // at `BeaconBlockBody` is already computed in the header.
-        Self::new_without_slasher_checks(block, &header, chain, custody_columns_count).map_err(
-            |e| {
-                process_block_slash_info::<_, BlockError>(
-                    chain,
-                    BlockSlashInfo::from_early_error_block(header, e),
-                )
-            },
-        )
+        Self::new_without_slasher_checks(block, &header, chain).map_err(|e| {
+            process_block_slash_info::<_, BlockError>(
+                chain,
+                BlockSlashInfo::from_early_error_block(header, e),
+            )
+        })
     }
 
     /// As for new, but doesn't pass the block to the slasher.
     fn new_without_slasher_checks(
         block: Arc<SignedBeaconBlock<T::EthSpec>>,
         block_header: &SignedBeaconBlockHeader,
         chain: &BeaconChain<T>,
-        custody_columns_count: usize,
     ) -> Result<Self, BlockError> {
         // Ensure the block is the correct structure for the fork at `block.slot()`.
         block
@@ -1053,7 +1045,6 @@ impl<T: BeaconChainTypes> GossipVerifiedBlock<T> {
             block_root,
             parent,
             consensus_context,
-            custody_columns_count,
         })
     }
 
@@ -1201,7 +1192,6 @@ impl<T: BeaconChainTypes> SignatureVerifiedBlock<T> {
                 block: MaybeAvailableBlock::AvailabilityPending {
                     block_root: from.block_root,
                     block,
-                    custody_columns_count: from.custody_columns_count,
                 },
                 block_root: from.block_root,
                 parent: Some(parent),

beacon_node/beacon_chain/src/block_verification_types.rs

@@ -31,7 +31,6 @@ use types::{
 pub struct RpcBlock<E: EthSpec> {
     block_root: Hash256,
     block: RpcBlockInner<E>,
-    custody_columns_count: usize,
 }
 
 impl<E: EthSpec> Debug for RpcBlock<E> {
@@ -45,10 +44,6 @@ impl<E: EthSpec> RpcBlock<E> {
         self.block_root
     }
 
-    pub fn custody_columns_count(&self) -> usize {
-        self.custody_columns_count
-    }
-
     pub fn as_block(&self) -> &SignedBeaconBlock<E> {
         match &self.block {
             RpcBlockInner::Block(block) => block,
@@ -103,14 +98,12 @@ impl<E: EthSpec> RpcBlock<E> {
     pub fn new_without_blobs(
         block_root: Option<Hash256>,
         block: Arc<SignedBeaconBlock<E>>,
-        custody_columns_count: usize,
     ) -> Self {
         let block_root = block_root.unwrap_or_else(|| get_block_root(&block));
 
         Self {
             block_root,
             block: RpcBlockInner::Block(block),
-            custody_columns_count,
         }
     }
 
@@ -152,16 +145,13 @@ impl<E: EthSpec> RpcBlock<E> {
         Ok(Self {
             block_root,
             block: inner,
-            // Block is before PeerDAS
-            custody_columns_count: 0,
         })
     }
 
     pub fn new_with_custody_columns(
         block_root: Option<Hash256>,
         block: Arc<SignedBeaconBlock<E>>,
         custody_columns: Vec<CustodyDataColumn<E>>,
-        custody_columns_count: usize,
         spec: &ChainSpec,
     ) -> Result<Self, AvailabilityCheckError> {
         let block_root = block_root.unwrap_or_else(|| get_block_root(&block));
@@ -182,7 +172,6 @@ impl<E: EthSpec> RpcBlock<E> {
         Ok(Self {
             block_root,
             block: inner,
-            custody_columns_count,
         })
     }
 
@@ -250,12 +239,10 @@ impl<E: EthSpec> ExecutedBlock<E> {
             MaybeAvailableBlock::AvailabilityPending {
                 block_root: _,
                 block: pending_block,
-                custody_columns_count,
             } => Self::AvailabilityPending(AvailabilityPendingExecutedBlock::new(
                 pending_block,
                 import_data,
                 payload_verification_outcome,
-                custody_columns_count,
             )),
         }
     }
@@ -321,21 +308,18 @@ pub struct AvailabilityPendingExecutedBlock<E: EthSpec> {
     pub block: Arc<SignedBeaconBlock<E>>,
     pub import_data: BlockImportData<E>,
     pub payload_verification_outcome: PayloadVerificationOutcome,
-    pub custody_columns_count: usize,
 }
 
 impl<E: EthSpec> AvailabilityPendingExecutedBlock<E> {
     pub fn new(
         block: Arc<SignedBeaconBlock<E>>,
         import_data: BlockImportData<E>,
         payload_verification_outcome: PayloadVerificationOutcome,
-        custody_columns_count: usize,
     ) -> Self {
         Self {
             block,
             import_data,
             payload_verification_outcome,
-            custody_columns_count,
         }
     }
 

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)),