feat: add select_random_peers to census (#1850)

Author: KolbyML

bin/portal-bridge/src/census/mod.rs

@@ -3,7 +3,7 @@ use std::{collections::HashSet, time::Duration};
 use discv5::enr::NodeId;
 use ethportal_api::{
     types::{network::Subnetwork, portal_wire::OfferTrace},
-    OverlayContentKey,
+    BeaconContentKey, HistoryContentKey, OverlayContentKey, StateContentKey,
 };
 use network::{Network, NetworkAction, NetworkInitializationConfig, NetworkManager};
 use peer::PeerInfo;
@@ -75,9 +75,22 @@ impl Census {
         content_key: &impl OverlayContentKey,
     ) -> Result<Vec<PeerInfo>, CensusError> {
         match subnetwork {
-            Subnetwork::History => self.history.select_peers(content_key),
-            Subnetwork::State => self.state.select_peers(content_key),
-            Subnetwork::Beacon => self.beacon.select_peers(content_key),
+            Subnetwork::History => self.history.select_peers(Some(content_key)),
+            Subnetwork::State => self.state.select_peers(Some(content_key)),
+            Subnetwork::Beacon => self.beacon.select_peers(Some(content_key)),
+            _ => Err(CensusError::UnsupportedSubnetwork(subnetwork)),
+        }
+    }
+
+    /// Selects random peers to receive content.
+    pub fn select_random_peers(
+        &self,
+        subnetwork: Subnetwork,
+    ) -> Result<Vec<PeerInfo>, CensusError> {
+        match subnetwork {
+            Subnetwork::History => self.history.select_peers(None::<&HistoryContentKey>),
+            Subnetwork::State => self.state.select_peers(None::<&StateContentKey>),
+            Subnetwork::Beacon => self.beacon.select_peers(None::<&BeaconContentKey>),
             _ => Err(CensusError::UnsupportedSubnetwork(subnetwork)),
         }
     }

bin/portal-bridge/src/census/network.rs

@@ -106,9 +106,12 @@ impl Network {
     }
 
     /// Selects peers to receive content.
+    ///
+    /// If content key is present, only peers interested in content will be returned.
+    /// Otherwise, all peers are eligible.
     pub fn select_peers(
         &self,
-        content_key: &impl OverlayContentKey,
+        content_key: Option<&impl OverlayContentKey>,
     ) -> Result<Vec<PeerInfo>, CensusError> {
         if self.peers.is_empty() {
             error!(

bin/portal-bridge/src/census/peers.rs

@@ -123,7 +123,10 @@ impl<W: Weight> Peers<W> {
     }
 
     /// Selects peers to receive content.
-    pub fn select_peers(&self, content_key: &impl OverlayContentKey) -> Vec<PeerInfo> {
+    ///
+    /// If content key is present, only peers interested in content will be returned.
+    /// Otherwise, all peers are eligible.
+    pub fn select_peers(&self, content_key: Option<&impl OverlayContentKey>) -> Vec<PeerInfo> {
         self.selector
             .select_peers(content_key, self.read().peers.values())
     }

bin/portal-bridge/src/census/scoring.rs

@@ -11,22 +11,25 @@ use super::peer::{Peer, PeerInfo};
 
 /// A trait for calculating peer's weight.
 pub trait Weight: Send + Sync {
+    /// Calculates peer's weight based on
+    /// - content key and id
+    /// - OfferResult's
+    /// - LivenessCheck's
     fn weight(
         &self,
-        content_key: &impl OverlayContentKey,
-        content_id: &[u8; 32],
+        content_key_and_id: &Option<(&impl OverlayContentKey, [u8; 32])>,
         peer: &Peer,
     ) -> u32;
 
     fn weight_all<'a>(
         &self,
-        content_key: &impl OverlayContentKey,
+        content_key: Option<&impl OverlayContentKey>,
         peers: impl IntoIterator<Item = &'a Peer>,
     ) -> impl Iterator<Item = (&'a Peer, u32)> {
-        let content_id = content_key.content_id();
+        let content_key_and_id = content_key.map(|key| (key, key.content_id()));
         peers
             .into_iter()
-            .map(move |peer| (peer, self.weight(content_key, &content_id, peer)))
+            .map(move |peer| (peer, self.weight(&content_key_and_id, peer)))
     }
 }
 
@@ -83,12 +86,13 @@ impl Default for AdditiveWeight {
 impl Weight for AdditiveWeight {
     fn weight(
         &self,
-        content_key: &impl OverlayContentKey,
-        content_id: &[u8; 32],
+        content_key_and_id: &Option<(&impl OverlayContentKey, [u8; 32])>,
         peer: &Peer,
     ) -> u32 {
-        if !peer.is_interested_in_content(content_key, content_id) {
-            return 0;
+        if let Some((content_key, content_id)) = content_key_and_id {
+            if !peer.is_interested_in_content(*content_key, content_id) {
+                return 0;
+            }
         }
 
         let liveness_weight = peer
@@ -145,7 +149,7 @@ impl<W: Weight> PeerSelector<W> {
     /// Selects up to `self.limit` peers based on their weights.
     pub fn select_peers<'a>(
         &self,
-        content_key: &impl OverlayContentKey,
+        content_key: Option<&impl OverlayContentKey>,
         peers: impl IntoIterator<Item = &'a Peer>,
     ) -> Vec<PeerInfo> {
         let weighted_peers = self