More progress

beacon_node/lighthouse_network/src/discovery/mod.rs

@@ -1194,9 +1194,8 @@ impl<E: EthSpec> Discovery<E> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::types::CGCUpdates;
     use libp2p::identity::secp256k1;
-    use types::{BitVector, MinimalEthSpec, SubnetId};
+    use types::{BitVector, CGCUpdates, MinimalEthSpec, SubnetId};
 
     type E = MinimalEthSpec;
 

beacon_node/lighthouse_network/src/peer_manager/mod.rs

@@ -31,7 +31,7 @@ pub use peerdb::peer_info::{
 };
 use peerdb::score::{PeerAction, ReportSource};
 pub use peerdb::sync_status::{SyncInfo, SyncStatus};
-use std::collections::{hash_map::Entry, HashMap, HashSet};
+use std::collections::{hash_map::Entry, HashMap};
 use std::net::IpAddr;
 use strum::IntoEnumIterator;
 use types::data_column_custody_group::{
@@ -1451,7 +1451,7 @@ impl<E: EthSpec> PeerManager<E> {
         &self,
         peer_id: &PeerId,
         custody_group_count: u64,
-    ) -> Result<HashSet<CustodyIndex>, String> {
+    ) -> Result<Vec<CustodyIndex>, String> {
         // If we don't have a node id, we cannot compute the custody duties anyway
         let node_id = peer_id_to_node_id(peer_id)?;
         let spec = &self.network_globals.spec;

beacon_node/lighthouse_network/src/service/mod.rs

@@ -15,7 +15,7 @@ use crate::rpc::{
 };
 use crate::types::{
     all_topics_at_fork, core_topics_to_subscribe, is_fork_non_core_topic, subnet_from_topic_hash,
-    CGCUpdates, GossipEncoding, GossipKind, GossipTopic, SnappyTransform, Subnet, SubnetDiscovery,
+    GossipEncoding, GossipKind, GossipTopic, SnappyTransform, Subnet, SubnetDiscovery,
 };
 use crate::EnrExt;
 use crate::Eth2Enr;
@@ -42,7 +42,7 @@ use types::{
     consts::altair::SYNC_COMMITTEE_SUBNET_COUNT, EnrForkId, Epoch, EthSpec, ForkContext, Slot,
     SubnetId,
 };
-use types::{ChainSpec, ForkName};
+use types::{CGCUpdates, ChainSpec, ForkName};
 use utils::{build_transport, strip_peer_id, Context as ServiceContext};
 
 pub mod api_types;
@@ -173,6 +173,7 @@ impl<E: EthSpec> Network<E> {
     pub async fn new(
         executor: task_executor::TaskExecutor,
         mut ctx: ServiceContext<'_>,
+        cgc_updates: Option<CGCUpdates>,
     ) -> Result<(Self, Arc<NetworkGlobals<E>>), String> {
         let config = ctx.config.clone();
         trace!("Libp2p Service starting");
@@ -197,10 +198,12 @@ impl<E: EthSpec> Network<E> {
         )?;
 
         // TODO: Load from disk, and check consistency with DB somewhere
-        let initial_cgc = ctx
-            .chain_spec
-            .custody_group_count(config.subscribe_all_data_column_subnets);
-        let cgc_updates = CGCUpdates::new(initial_cgc);
+        let cgc_updates = cgc_updates.unwrap_or_else(|| {
+            CGCUpdates::new(
+                ctx.chain_spec
+                    .custody_group_count(config.subscribe_all_data_column_subnets),
+            )
+        });
 
         // Construct the metadata
         let globals = NetworkGlobals::new(

beacon_node/lighthouse_network/src/types/globals.rs

@@ -11,7 +11,7 @@ use std::sync::Arc;
 use types::data_column_custody_group::{
     compute_columns_for_custody_group, compute_subnets_from_custody_group, get_custody_groups,
 };
-use types::{ChainSpec, ColumnIndex, DataColumnSubnetId, EthSpec, Slot};
+use types::{CGCUpdates, ChainSpec, ColumnIndex, DataColumnSubnetId, EthSpec, Slot};
 
 pub struct NetworkGlobals<E: EthSpec> {
     /// The current local ENR.
@@ -39,12 +39,6 @@ pub struct NetworkGlobals<E: EthSpec> {
     pub spec: Arc<ChainSpec>,
 }
 
-pub struct CGCUpdates {
-    initial_value: u64,
-    updates: Vec<(Slot, u64)>,
-    // TODO(das): Track backfilled CGC
-}
-
 impl<E: EthSpec> NetworkGlobals<E> {
     pub fn new(
         enr: Enr,
@@ -145,9 +139,8 @@ impl<E: EthSpec> NetworkGlobals<E> {
     }
 
     pub fn sampling_subnets(&self, slot: Slot) -> &[DataColumnSubnetId] {
-        let cgc = self.custody_group_count(slot) as usize;
-        // Returns as many elements as possible, can't panic as it's upper bounded by len
-        &self.all_sampling_subnets[..self.all_sampling_subnets.len().min(cgc)]
+        let cgc = self.custody_group_count(slot);
+        self.sampling_subnets_for_cgc(cgc)
     }
 
     pub fn sampling_columns(&self, slot: Slot) -> &[ColumnIndex] {
@@ -156,6 +149,11 @@ impl<E: EthSpec> NetworkGlobals<E> {
         &self.all_sampling_columns[..self.all_sampling_columns.len().min(cgc)]
     }
 
+    pub fn sampling_subnets_for_cgc(&self, cgc: u64) -> &[DataColumnSubnetId] {
+        // Returns as many elements as possible, can't panic as it's upper bounded by len
+        &self.all_sampling_subnets[..self.all_sampling_subnets.len().min(cgc as usize)]
+    }
+
     fn public_custody_group_count(&self) -> u64 {
         // TODO(das): delay announcing the public custody count for the duration of the pruning
         // period
@@ -163,7 +161,7 @@ impl<E: EthSpec> NetworkGlobals<E> {
     }
 
     /// Returns the custody group count (CGC)
-    fn custody_group_count(&self, slot: Slot) -> u64 {
+    pub fn custody_group_count(&self, slot: Slot) -> u64 {
         self.cgc_updates.read().at_slot(slot)
     }
 
@@ -176,8 +174,14 @@ impl<E: EthSpec> NetworkGlobals<E> {
     }
 
     /// Adds a new CGC value update
-    pub fn add_cgc_update(&self, update: (Slot, u64)) {
-        self.cgc_updates.write().add_latest_update(update);
+    pub fn add_cgc_update(&self, update_start_slot: Slot, cgc: u64) -> Result<(), String> {
+        self.cgc_updates
+            .write()
+            .add_latest_update((update_start_slot, cgc))
+    }
+
+    pub fn dump_cgc_updates(&self) -> CGCUpdates {
+        self.cgc_updates.read().clone()
     }
 
     /// Returns the number of libp2p connected peers.
@@ -276,30 +280,6 @@ impl<E: EthSpec> NetworkGlobals<E> {
     }
 }
 
-impl CGCUpdates {
-    pub fn new(initial_value: u64) -> Self {
-        Self {
-            initial_value,
-            updates: vec![],
-        }
-    }
-
-    fn at_slot(&self, slot: Slot) -> u64 {
-        // TODO: Test and fix logic
-        for (update_slot, cgc) in &self.updates {
-            if slot > *update_slot {
-                return *cgc;
-            }
-        }
-
-        self.initial_value
-    }
-
-    fn add_latest_update(&mut self, update: (Slot, u64)) {
-        self.updates.push(update);
-    }
-}
-
 #[cfg(test)]
 mod test {
     use super::*;

beacon_node/lighthouse_network/src/types/mod.rs

@@ -11,7 +11,7 @@ pub type EnrSyncCommitteeBitfield<E> = BitVector<<E as EthSpec>::SyncCommitteeSu
 pub type Enr = discv5::enr::Enr<discv5::enr::CombinedKey>;
 
 pub use eth2::lighthouse::sync_state::{BackFillState, SyncState};
-pub use globals::{CGCUpdates, NetworkGlobals};
+pub use globals::NetworkGlobals;
 pub use pubsub::{PubsubMessage, SnappyTransform};
 pub use subnet::{Subnet, SubnetDiscovery};
 pub use topics::{

beacon_node/lighthouse_network/tests/common.rs

@@ -117,7 +117,7 @@ pub async fn build_libp2p_instance(
         libp2p_registry: None,
     };
     Libp2pInstance(
-        LibP2PService::new(executor, libp2p_context)
+        LibP2PService::new(executor, libp2p_context, None)
             .await
             .expect("should build libp2p instance")
             .0,