moves socket-addr-space check to get_retransmit_addrs

anza-xyz · Dec 16, 2024 · e6f1b71 · e6f1b71
1 parent 6954807
commit e6f1b71
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 9 deletions.
diff --git a/streamer/src/socket.rs b/streamer/src/socket.rs
@@ -16,6 +16,7 @@ impl SocketAddrSpace {
     }
 
     /// Returns true if the IP address is valid.
+    #[inline]
     #[must_use]
     pub fn check(&self, addr: &SocketAddr) -> bool {
         if matches!(self, SocketAddrSpace::Unspecified) {

diff --git a/turbine/benches/cluster_nodes.rs b/turbine/benches/cluster_nodes.rs
@@ -7,6 +7,7 @@ use {
     solana_gossip::contact_info::ContactInfo,
     solana_ledger::shred::{Shred, ShredFlags},
     solana_sdk::{clock::Slot, genesis_config::ClusterType, pubkey::Pubkey},
+    solana_streamer::socket::SocketAddrSpace,
     solana_turbine::{
         cluster_nodes::{make_test_cluster, new_cluster_nodes, ClusterNodes},
         retransmit_stage::RetransmitStage,
@@ -45,8 +46,12 @@ fn get_retransmit_peers_deterministic(
             0,
             0,
         );
-        let _retransmit_peers =
-            cluster_nodes.get_retransmit_addrs(slot_leader, &shred.id(), /*fanout:*/ 200);
+        let _retransmit_peers = cluster_nodes.get_retransmit_addrs(
+            slot_leader,
+            &shred.id(),
+            200, // fanout
+            &SocketAddrSpace::Unspecified,
+        );
     }
 }
 

diff --git a/turbine/src/cluster_nodes.rs b/turbine/src/cluster_nodes.rs
@@ -165,6 +165,7 @@ impl ClusterNodes<RetransmitStage> {
         slot_leader: &Pubkey,
         shred: &ShredId,
         fanout: usize,
+        socket_addr_space: &SocketAddrSpace,
     ) -> Result<(/*root_distance:*/ usize, Vec<SocketAddr>), Error> {
         let mut weighted_shuffle = self.weighted_shuffle.clone();
         // Exclude slot leader from list of nodes.
@@ -195,7 +196,10 @@ impl ClusterNodes<RetransmitStage> {
         };
         let (index, peers) =
             get_retransmit_peers(fanout, |(node, _)| node.pubkey() == self.pubkey, nodes);
-        let peers = peers.filter_map(|(_, addr)| addr).collect();
+        let peers = peers
+            .filter_map(|(_, addr)| addr)
+            .filter(|addr| socket_addr_space.check(addr))
+            .collect();
         let root_distance = get_root_distance(index, fanout);
         Ok((root_distance, peers))
     }

diff --git a/turbine/src/retransmit_stage.rs b/turbine/src/retransmit_stage.rs
@@ -328,12 +328,12 @@ fn retransmit_shred(
 ) -> Result<(/*root_distance:*/ usize, /*num_nodes:*/ usize), Error> {
     let mut compute_turbine_peers = Measure::start("turbine_start");
     let data_plane_fanout = cluster_nodes::get_data_plane_fanout(key.slot(), root_bank);
-    let (root_distance, addrs) =
-        cluster_nodes.get_retransmit_addrs(slot_leader, key, data_plane_fanout)?;
-    let addrs: Vec<_> = addrs
-        .into_iter()
-        .filter(|addr| socket_addr_space.check(addr))
-        .collect();
+    let (root_distance, addrs) = cluster_nodes.get_retransmit_addrs(
+        slot_leader,
+        key,
+        data_plane_fanout,
+        socket_addr_space,
+    )?;
     compute_turbine_peers.stop();
     stats
         .compute_turbine_peers_total