refactor: census peer discovery algorithm (#1553)

Cargo.toml

@@ -94,6 +94,7 @@ ethereum_ssz_derive = "0.7.1"
 ethportal-api = { path = "ethportal-api" }
 futures = "0.3.23"
 hex = "0.4.3"
+itertools = "0.13.0"
 jsonrpsee = "0.24.4"
 keccak-hash = "0.10.0"
 lazy_static = "1.4.0"

ethportal-api/src/types/node_id.rs

@@ -16,3 +16,36 @@ pub fn generate_random_node_id(target_bucket_idx: u8, local_node_id: NodeId) ->
 
     raw_bytes.into()
 }
+
+/// Generates `2 ^ unique_bits` random `NodeId`s.
+///
+/// The most significant bits of each `NodeId` will be unique.
+///
+/// The `unique_bits` has to be in `(0..=8)` range, panics otherwise. Can be easily upgraded to
+/// support wider range.
+pub fn generate_random_node_ids(unique_bits: u32) -> Vec<NodeId> {
+    assert!(
+        (0..=8).contains(&unique_bits),
+        "Invalid bits value: {unique_bits}"
+    );
+
+    let insignificant_bits = u8::BITS - unique_bits;
+    let insignificant_bits_mask = u8::MAX.checked_shr(unique_bits).unwrap_or_default();
+
+    (0usize..1 << unique_bits)
+        .map(|index| {
+            // shift bits to most significant positions
+            let unique_bits = (index as u8)
+                .checked_shl(insignificant_bits)
+                .unwrap_or_default();
+
+            let mut node_id = rand::random::<[u8; 32]>();
+
+            // set most significant bits of the first byte
+            node_id[0] &= insignificant_bits_mask;
+            node_id[0] |= unique_bits;
+
+            NodeId::from(node_id)
+        })
+        .collect()
+}

portal-bridge/Cargo.toml

@@ -24,6 +24,7 @@ eth_trie.workspace = true
 ethereum_ssz.workspace = true
 ethportal-api.workspace = true
 futures.workspace = true
+itertools.workspace = true
 jsonrpsee = { workspace = true, features = [
   "async-client",
   "client",

portal-bridge/src/census/mod.rs

@@ -6,7 +6,7 @@ use tokio::task::JoinHandle;
 use tracing::{error, info, Instrument};
 
 use crate::cli::BridgeConfig;
-use network::Network;
+use network::{Network, NetworkAction, NetworkInitializationConfig, NetworkManager};
 
 mod network;
 mod peers;
@@ -40,6 +40,9 @@ pub struct Census {
 }
 
 impl Census {
+    const SUPPORTED_SUBNETWORKS: [Subnetwork; 3] =
+        [Subnetwork::Beacon, Subnetwork::History, Subnetwork::State];
+
     pub fn new(client: HttpClient, bridge_config: &BridgeConfig) -> Self {
         Self {
             history: Network::new(client.clone(), Subnetwork::History, bridge_config),
@@ -71,50 +74,74 @@ impl Census {
         &mut self,
         subnetworks: impl IntoIterator<Item = Subnetwork>,
     ) -> Result<JoinHandle<()>, CensusError> {
+        info!("Initializing census");
+
         if self.initialized {
             return Err(CensusError::AlreadyInitialized);
         }
         self.initialized = true;
 
-        let subnetworks = HashSet::from_iter(subnetworks);
+        let subnetworks = HashSet::<Subnetwork>::from_iter(subnetworks);
         if subnetworks.is_empty() {
             return Err(CensusError::FailedInitialization("No subnetwork"));
         }
         for subnetwork in &subnetworks {
-            info!("Initializing {subnetwork} subnetwork");
-            match subnetwork {
-                Subnetwork::History => self.history.init().await?,
-                Subnetwork::State => self.state.init().await?,
-                Subnetwork::Beacon => self.beacon.init().await?,
-                _ => return Err(CensusError::UnsupportedSubnetwork(*subnetwork)),
+            if !Self::SUPPORTED_SUBNETWORKS.contains(subnetwork) {
+                return Err(CensusError::UnsupportedSubnetwork(*subnetwork));
             }
         }
 
-        Ok(self.start_background_service(subnetworks))
-    }
+        let initialization_config = NetworkInitializationConfig::default();
+
+        let mut beacon_manager = if subnetworks.contains(&Subnetwork::Beacon) {
+            self.beacon.init(&initialization_config).await?;
+            Some(self.beacon.create_manager())
+        } else {
+            None
+        };
+        let mut history_manager = if subnetworks.contains(&Subnetwork::History) {
+            self.history.init(&initialization_config).await?;
+            Some(self.history.create_manager())
+        } else {
+            None
+        };
+        let mut state_manager = if subnetworks.contains(&Subnetwork::State) {
+            self.state.init(&initialization_config).await?;
+            Some(self.state.create_manager())
+        } else {
+            None
+        };
 
-    /// Starts background service that is responsible for keeping view of the network up to date.
-    ///
-    /// Selects available tasks and runs them. Tasks are provided by enabled subnetworks.
-    fn start_background_service(&self, subnetworks: HashSet<Subnetwork>) -> JoinHandle<()> {
-        let mut history_network = self.history.clone();
-        let mut state_network = self.state.clone();
-        let mut beacon_network = self.beacon.clone();
         let service = async move {
             loop {
                 tokio::select! {
-                    peer = history_network.peer_to_process(), if subnetworks.contains(&Subnetwork::History) => {
-                        history_network.process_peer(peer).await;
+                    Some(action) = next_action(&mut beacon_manager) => {
+                        if let Some(manager) = &mut beacon_manager {
+                            manager.execute_action(action).await;
+                        }
                     }
-                    peer = state_network.peer_to_process(), if subnetworks.contains(&Subnetwork::State) => {
-                        state_network.process_peer(peer).await;
+                    Some(action) = next_action(&mut history_manager) => {
+                        if let Some(manager) = &mut history_manager {
+                            manager.execute_action(action).await;
+                        }
                     }
-                    peer = beacon_network.peer_to_process(), if subnetworks.contains(&Subnetwork::Beacon) => {
-                        beacon_network.process_peer(peer).await;
+                    Some(action) = next_action(&mut state_manager) => {
+                        if let Some(manager) = &mut state_manager {
+                            manager.execute_action(action).await;
+                        }
                     }
                 }
             }
         };
-        tokio::spawn(service.instrument(tracing::trace_span!("census").or_current()))
+        Ok(tokio::spawn(
+            service.instrument(tracing::trace_span!("census").or_current()),
+        ))
+    }
+}
+
+async fn next_action(manager: &mut Option<NetworkManager>) -> Option<NetworkAction> {
+    match manager {
+        Some(manager) => Some(manager.next_action().await),
+        None => None,
     }
 }