Fix custdoy context initialisation race condition.

Author: jimmygchen

beacon_node/beacon_chain/src/builder.rs

@@ -40,6 +40,7 @@ use std::time::Duration;
 use store::{Error as StoreError, HotColdDB, ItemStore, KeyValueStoreOp};
 use task_executor::{ShutdownReason, TaskExecutor};
 use tracing::{debug, error, info};
+use types::data_column_custody_group::{CustodyIndex, get_custody_groups_ordered};
 use types::{
     BeaconBlock, BeaconState, BlobSidecarList, ChainSpec, DataColumnSidecarList, Epoch, EthSpec,
     FixedBytesExtended, Hash256, Signature, SignedBeaconBlock, Slot,
@@ -102,6 +103,7 @@ pub struct BeaconChainBuilder<T: BeaconChainTypes> {
     task_executor: Option<TaskExecutor>,
     validator_monitor_config: Option<ValidatorMonitorConfig>,
     node_custody_type: NodeCustodyType,
+    node_id: Option<[u8; 32]>,
     rng: Option<Box<dyn RngCore + Send>>,
 }
 
@@ -141,6 +143,7 @@ where
             task_executor: None,
             validator_monitor_config: None,
             node_custody_type: NodeCustodyType::Fullnode,
+            node_id: None,
             rng: None,
         }
     }
@@ -647,6 +650,11 @@ where
         self
     }
 
+    pub fn node_id(mut self, node_id: [u8; 32]) -> Self {
+        self.node_id = Some(node_id);
+        self
+    }
+
     /// Sets the `BeaconChain` event handler backend.
     ///
     /// For example, provide `ServerSentEventHandler` as a `handler`.
@@ -740,6 +748,7 @@ where
             .genesis_state_root
             .ok_or("Cannot build without a genesis state root")?;
         let validator_monitor_config = self.validator_monitor_config.unwrap_or_default();
+        let node_id = self.node_id.ok_or("Cannot build without a node id")?;
         let rng = self.rng.ok_or("Cannot build without an RNG")?;
         let beacon_proposer_cache: Arc<Mutex<BeaconProposerCache>> = <_>::default();
 
@@ -929,6 +938,10 @@ where
             }
         };
 
+        let all_custody_groups_ordered =
+            get_custody_groups_ordered(node_id, self.spec.number_of_custody_groups, &self.spec)
+                .map_err(|e| format!("Failed to compute custody groups: {:?}", e))?;
+
         // Load the persisted custody context from the db and initialize
         // the context for this run
         let (custody_context, cgc_changed_opt) = if let Some(custody) =
@@ -942,11 +955,16 @@ where
                 custody,
                 self.node_custody_type,
                 head_epoch,
+                all_custody_groups_ordered,
                 &self.spec,
             )
         } else {
             (
-                CustodyContext::new(self.node_custody_type, &self.spec),
+                CustodyContext::new(
+                    self.node_custody_type,
+                    all_custody_groups_ordered,
+                    &self.spec,
+                ),
                 None,
             )
         };

beacon_node/beacon_chain/src/custody_context.rs

@@ -250,7 +250,7 @@ pub struct CustodyContext<E: EthSpec> {
     validator_registrations: RwLock<ValidatorRegistrations>,
     /// Stores an immutable, ordered list of all custody columns as determined by the node's NodeID
     /// on startup.
-    all_custody_columns_ordered: OnceLock<Box<[ColumnIndex]>>,
+    all_custody_columns_ordered: Vec<ColumnIndex>,
     _phantom_data: PhantomData<E>,
 }
 
@@ -290,6 +290,7 @@ impl<E: EthSpec> CustodyContext<E> {
         ssz_context: CustodyContextSsz,
         node_custody_type: NodeCustodyType,
         head_epoch: Epoch,
+        all_custody_groups_ordered: Vec<CustodyIndex>,
         spec: &ChainSpec,
     ) -> (Self, Option<CustodyCountChanged>) {
         let CustodyContextSsz {
@@ -347,6 +348,10 @@ impl<E: EthSpec> CustodyContext<E> {
             }
         }
 
+        // FIXME: remove unwrap
+        let all_custody_columns_ordered =
+            Self::compute_ordered_data_columns(all_custody_groups_ordered, spec).unwrap();
+
         let custody_context = CustodyContext {
             validator_custody_count: AtomicU64::new(validator_custody_at_head),
             validator_registrations: RwLock::new(ValidatorRegistrations {
@@ -355,7 +360,7 @@ impl<E: EthSpec> CustodyContext<E> {
                     .into_iter()
                     .collect(),
             }),
-            all_custody_columns_ordered: OnceLock::new(),
+            all_custody_columns_ordered,
             _phantom_data: PhantomData,
         };
 
@@ -370,22 +375,17 @@ impl<E: EthSpec> CustodyContext<E> {
     ///
     /// # Returns
     /// Ok(()) if initialization succeeds, Err with description string if it fails
-    pub fn init_ordered_data_columns_from_custody_groups(
-        &self,
+    fn compute_ordered_data_columns(
         all_custody_groups_ordered: Vec<CustodyIndex>,
         spec: &ChainSpec,
-    ) -> Result<(), String> {
+    ) -> Result<Vec<ColumnIndex>, String> {
         let mut ordered_custody_columns = vec![];
         for custody_index in all_custody_groups_ordered {
             let columns = compute_columns_for_custody_group::<E>(custody_index, spec)
                 .map_err(|e| format!("Failed to compute columns for custody group {e:?}"))?;
             ordered_custody_columns.extend(columns);
         }
-        self.all_custody_columns_ordered
-            .set(ordered_custody_columns.into_boxed_slice())
-            .map_err(|_| {
-                "Failed to initialise CustodyContext with computed custody columns".to_string()
-            })
+        Ok(ordered_custody_columns)
     }
 
     /// Register a new validator index and updates the list of validators if required.
@@ -497,11 +497,7 @@ impl<E: EthSpec> CustodyContext<E> {
     /// A slice of ordered column indices that should be sampled for this epoch based on the node's custody configuration
     pub fn sampling_columns_for_epoch(&self, epoch: Epoch, spec: &ChainSpec) -> &[ColumnIndex] {
         let num_of_columns_to_sample = self.num_of_data_columns_to_sample(epoch, spec);
-        let all_columns_ordered = self
-            .all_custody_columns_ordered
-            .get()
-            .expect("all_custody_columns_ordered should be initialized");
-        &all_columns_ordered[..num_of_columns_to_sample]
+        &self.all_custody_columns_ordered[..num_of_columns_to_sample]
     }
 
     /// Returns the ordered list of column indices that the node is assigned to custody
@@ -528,12 +524,7 @@ impl<E: EthSpec> CustodyContext<E> {
             self.custody_group_count_at_head(spec) as usize
         };
 
-        let all_columns_ordered = self
-            .all_custody_columns_ordered
-            .get()
-            .expect("all_custody_columns_ordered should be initialized");
-
-        &all_columns_ordered[..custody_group_count]
+        &self.all_custody_columns_ordered[..custody_group_count]
     }
 
     /// The node has completed backfill for this epoch. Update the internal records so the function

beacon_node/client/src/builder.rs

@@ -2,7 +2,7 @@ use crate::Client;
 use crate::compute_light_client_updates::{
     LIGHT_CLIENT_SERVER_CHANNEL_CAPACITY, compute_light_client_updates,
 };
-use crate::config::{ClientGenesis, Config as ClientConfig};
+use crate::config::{ClientGenesis, Config as ClientConfig, Config};
 use crate::notifier::spawn_notifier;
 use beacon_chain::attestation_simulator::start_attestation_simulator_service;
 use beacon_chain::data_availability_checker::start_availability_cache_maintenance_service;
@@ -28,6 +28,8 @@ use execution_layer::ExecutionLayer;
 use execution_layer::test_utils::generate_genesis_header;
 use futures::channel::mpsc::Receiver;
 use genesis::{DEFAULT_ETH1_BLOCK_HASH, interop_genesis_state};
+use lighthouse_network::discv5::enr::NodeId;
+use lighthouse_network::identity::Keypair;
 use lighthouse_network::{NetworkGlobals, prometheus_client::registry::Registry};
 use monitoring_api::{MonitoringHttpClient, ProcessType};
 use network::{NetworkConfig, NetworkSenders, NetworkService};
@@ -42,7 +44,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use store::database::interface::BeaconNodeBackend;
 use timer::spawn_timer;
 use tracing::{debug, info, warn};
-use types::data_column_custody_group::get_custody_groups_ordered;
+use types::data_column_custody_group::{CustodyIndex, get_custody_groups_ordered};
 use types::{
     BeaconState, BlobSidecarList, ChainSpec, EthSpec, ExecutionBlockHash, Hash256,
     SignedBeaconBlock, test_utils::generate_deterministic_keypairs,
@@ -154,6 +156,7 @@ where
         mut self,
         client_genesis: ClientGenesis,
         config: ClientConfig,
+        local_keypair: Keypair,
     ) -> Result<Self, String> {
         let store = self.store.clone();
         let chain_spec = self.chain_spec.clone();
@@ -203,6 +206,7 @@ where
             .event_handler(event_handler)
             .execution_layer(execution_layer)
             .node_custody_type(config.chain.node_custody_type)
+            .node_id(NodeId::from(local_keypair.public()).raw())
             .validator_monitor_config(config.validator_monitor.clone())
             .rng(Box::new(
                 StdRng::try_from_rng(&mut OsRng)
@@ -453,7 +457,11 @@ where
     }
 
     /// Starts the networking stack.
-    pub async fn network(mut self, config: Arc<NetworkConfig>) -> Result<Self, String> {
+    pub async fn network(
+        mut self,
+        config: Arc<NetworkConfig>,
+        local_keypair: Keypair,
+    ) -> Result<Self, String> {
         let beacon_chain = self
             .beacon_chain
             .clone()
@@ -481,12 +489,11 @@ where
             context.executor,
             libp2p_registry.as_mut(),
             beacon_processor_channels.beacon_processor_tx.clone(),
+            local_keypair,
         )
         .await
         .map_err(|e| format!("Failed to start network: {:?}", e))?;
 
-        init_custody_context(beacon_chain, &network_globals)?;
-
         self.network_globals = Some(network_globals);
         self.network_senders = Some(network_senders);
         self.libp2p_registry = libp2p_registry;
@@ -788,21 +795,6 @@ where
     }
 }
 
-fn init_custody_context<T: BeaconChainTypes>(
-    chain: Arc<BeaconChain<T>>,
-    network_globals: &NetworkGlobals<T::EthSpec>,
-) -> Result<(), String> {
-    let node_id = network_globals.local_enr().node_id().raw();
-    let spec = &chain.spec;
-    let custody_groups_ordered =
-        get_custody_groups_ordered(node_id, spec.number_of_custody_groups, spec)
-            .map_err(|e| format!("Failed to compute custody groups: {:?}", e))?;
-    chain
-        .data_availability_checker
-        .custody_context()
-        .init_ordered_data_columns_from_custody_groups(custody_groups_ordered, spec)
-}
-
 impl<TSlotClock, E, THotStore, TColdStore>
     ClientBuilder<Witness<TSlotClock, E, THotStore, TColdStore>>
 where

beacon_node/lighthouse_network/src/lib.rs

@@ -109,7 +109,7 @@ pub use discovery::Eth2Enr;
 pub use discv5;
 pub use gossipsub::{IdentTopic, MessageAcceptance, MessageId, Topic, TopicHash};
 pub use libp2p;
-pub use libp2p::{Multiaddr, multiaddr};
+pub use libp2p::{Multiaddr, identity, multiaddr};
 pub use libp2p::{PeerId, Swarm, core::ConnectedPoint};
 pub use peer_manager::{
     ConnectionDirection, PeerConnectionStatus, PeerInfo, PeerManager, SyncInfo, SyncStatus,

beacon_node/lighthouse_network/src/service/mod.rs

@@ -26,6 +26,7 @@ use gossipsub::{
     TopicScoreParams,
 };
 use gossipsub_scoring_parameters::{PeerScoreSettings, lighthouse_gossip_thresholds};
+use libp2p::identity::Keypair;
 use libp2p::multiaddr::{self, Multiaddr, Protocol as MProtocol};
 use libp2p::swarm::behaviour::toggle::Toggle;
 use libp2p::swarm::{NetworkBehaviour, Swarm, SwarmEvent};
@@ -171,11 +172,10 @@ impl<E: EthSpec> Network<E> {
         executor: task_executor::TaskExecutor,
         mut ctx: ServiceContext<'_>,
         custody_group_count: u64,
+        local_keypair: Keypair,
     ) -> Result<(Self, Arc<NetworkGlobals<E>>), String> {
         let config = ctx.config.clone();
         trace!("Libp2p Service starting");
-        // initialise the node's ID
-        let local_keypair = utils::load_private_key(&config);
 
         // Trusted peers will also be marked as explicit in GossipSub.
         // Cfr. https://github.com/libp2p/specs/blob/master/pubsub/gossipsub/gossipsub-v1.1.md#explicit-peering-agreements

beacon_node/network/src/service.rs

@@ -12,6 +12,7 @@ use futures::future::OptionFuture;
 use futures::prelude::*;
 
 use lighthouse_network::Enr;
+use lighthouse_network::identity::Keypair;
 use lighthouse_network::rpc::InboundRequestId;
 use lighthouse_network::rpc::RequestType;
 use lighthouse_network::rpc::methods::RpcResponse;
@@ -212,6 +213,7 @@ impl<T: BeaconChainTypes> NetworkService<T> {
         executor: task_executor::TaskExecutor,
         libp2p_registry: Option<&'_ mut Registry>,
         beacon_processor_send: BeaconProcessorSend<T::EthSpec>,
+        local_keypair: Keypair,
     ) -> Result<
         (
             NetworkService<T>,
@@ -284,6 +286,7 @@ impl<T: BeaconChainTypes> NetworkService<T> {
                 .data_availability_checker
                 .custody_context()
                 .custody_group_count_at_head(&beacon_chain.spec),
+            local_keypair,
         )
         .await?;
 
@@ -366,13 +369,15 @@ impl<T: BeaconChainTypes> NetworkService<T> {
         executor: task_executor::TaskExecutor,
         libp2p_registry: Option<&'_ mut Registry>,
         beacon_processor_send: BeaconProcessorSend<T::EthSpec>,
+        local_keypair: Keypair,
     ) -> Result<(Arc<NetworkGlobals<T::EthSpec>>, NetworkSenders<T::EthSpec>), String> {
         let (network_service, network_globals, network_senders) = Self::build(
             beacon_chain,
             config,
             executor.clone(),
             libp2p_registry,
             beacon_processor_send,
+            local_keypair,
         )
         .await?;
 

beacon_node/src/lib.rs

@@ -9,6 +9,7 @@ pub use client::{Client, ClientBuilder, ClientConfig, ClientGenesis};
 pub use config::{get_config, get_data_dir, set_network_config};
 use environment::RuntimeContext;
 pub use eth2_config::Eth2Config;
+use lighthouse_network::load_private_key;
 use slasher::{DatabaseBackendOverride, Slasher};
 use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
@@ -120,8 +121,9 @@ impl<E: EthSpec> ProductionBeaconNode<E> {
             builder
         };
 
+        let local_keypair = load_private_key(&client_config.network);
         let builder = builder
-            .beacon_chain_builder(client_genesis, client_config.clone())
+            .beacon_chain_builder(client_genesis, client_config.clone(), local_keypair)
             .await?;
         info!("Block production enabled");
 
@@ -133,7 +135,7 @@ impl<E: EthSpec> ProductionBeaconNode<E> {
 
         builder
             .build_beacon_chain()?
-            .network(Arc::new(client_config.network))
+            .network(Arc::new(client_config.network), local_keypair)
             .await?
             .notifier()?
             .http_metrics_config(client_config.http_metrics.clone())