Separate staked and un-staked on quic tpu port (#24339)

core/src/lib.rs

@@ -59,6 +59,7 @@ pub mod sigverify;
 pub mod sigverify_shreds;
 pub mod sigverify_stage;
 pub mod snapshot_packager_service;
+pub mod staked_nodes_updater_service;
 pub mod stats_reporter_service;
 pub mod system_monitor_service;
 mod tower1_7_14;

core/src/staked_nodes_updater_service.rs

@@ -0,0 +1,76 @@
+use {
+    solana_gossip::cluster_info::ClusterInfo,
+    solana_runtime::bank_forks::BankForks,
+    std::{
+        collections::HashMap,
+        net::IpAddr,
+        sync::{
+            atomic::{AtomicBool, Ordering},
+            Arc, RwLock,
+        },
+        thread::{self, sleep, Builder, JoinHandle},
+        time::{Duration, Instant},
+    },
+};
+
+const IP_TO_STAKE_REFRESH_DURATION: Duration = Duration::from_secs(5);
+
+pub struct StakedNodesUpdaterService {
+    thread_hdl: JoinHandle<()>,
+}
+
+impl StakedNodesUpdaterService {
+    pub fn new(
+        exit: Arc<AtomicBool>,
+        cluster_info: Arc<ClusterInfo>,
+        bank_forks: Arc<RwLock<BankForks>>,
+        shared_staked_nodes: Arc<RwLock<HashMap<IpAddr, u64>>>,
+    ) -> Self {
+        let thread_hdl = Builder::new()
+            .name("sol-sn-updater".to_string())
+            .spawn(move || {
+                let mut last_stakes = Instant::now();
+                while !exit.load(Ordering::Relaxed) {
+                    let mut new_ip_to_stake = HashMap::new();
+                    Self::try_refresh_ip_to_stake(
+                        &mut last_stakes,
+                        &mut new_ip_to_stake,
+                        &bank_forks,
+                        &cluster_info,
+                    );
+                    let mut shared = shared_staked_nodes.write().unwrap();
+                    *shared = new_ip_to_stake;
+                }
+            })
+            .unwrap();
+
+        Self { thread_hdl }
+    }
+
+    fn try_refresh_ip_to_stake(
+        last_stakes: &mut Instant,
+        ip_to_stake: &mut HashMap<IpAddr, u64>,
+        bank_forks: &RwLock<BankForks>,
+        cluster_info: &ClusterInfo,
+    ) {
+        if last_stakes.elapsed() > IP_TO_STAKE_REFRESH_DURATION {
+            let root_bank = bank_forks.read().unwrap().root_bank();
+            let staked_nodes = root_bank.staked_nodes();
+            *ip_to_stake = cluster_info
+                .tvu_peers()
+                .into_iter()
+                .filter_map(|node| {
+                    let stake = staked_nodes.get(&node.id)?;
+                    Some((node.tvu.ip(), *stake))
+                })
+                .collect();
+            *last_stakes = Instant::now();
+        } else {
+            sleep(Duration::from_millis(1));
+        }
+    }
+
+    pub fn join(self) -> thread::Result<()> {
+        self.thread_hdl.join()
+    }
+}

core/src/tpu.rs

@@ -13,6 +13,7 @@ use {
         find_packet_sender_stake_stage::FindPacketSenderStakeStage,
         sigverify::TransactionSigVerifier,
         sigverify_stage::SigVerifyStage,
+        staked_nodes_updater_service::StakedNodesUpdaterService,
     },
     crossbeam_channel::{bounded, unbounded, Receiver, RecvTimeoutError},
     solana_gossip::cluster_info::ClusterInfo,
@@ -28,7 +29,9 @@ use {
         vote_sender_types::{ReplayVoteReceiver, ReplayVoteSender},
     },
     solana_sdk::signature::Keypair,
+    solana_streamer::quic::{spawn_server, MAX_STAKED_CONNECTIONS, MAX_UNSTAKED_CONNECTIONS},
     std::{
+        collections::HashMap,
         net::UdpSocket,
         sync::{atomic::AtomicBool, Arc, Mutex, RwLock},
         thread,
@@ -62,6 +65,7 @@ pub struct Tpu {
     tpu_quic_t: thread::JoinHandle<()>,
     find_packet_sender_stake_stage: FindPacketSenderStakeStage,
     vote_find_packet_sender_stake_stage: FindPacketSenderStakeStage,
+    staked_nodes_updater_service: StakedNodesUpdaterService,
 }
 
 impl Tpu {
@@ -132,13 +136,23 @@ impl Tpu {
 
         let (verified_sender, verified_receiver) = unbounded();
 
-        let tpu_quic_t = solana_streamer::quic::spawn_server(
+        let staked_nodes = Arc::new(RwLock::new(HashMap::new()));
+        let staked_nodes_updater_service = StakedNodesUpdaterService::new(
+            exit.clone(),
+            cluster_info.clone(),
+            bank_forks.clone(),
+            staked_nodes.clone(),
+        );
+        let tpu_quic_t = spawn_server(
             transactions_quic_sockets,
             keypair,
             cluster_info.my_contact_info().tpu.ip(),
             packet_sender,
             exit.clone(),
             MAX_QUIC_CONNECTIONS_PER_IP,
+            staked_nodes,
+            MAX_STAKED_CONNECTIONS,
+            MAX_UNSTAKED_CONNECTIONS,
         )
         .unwrap();
 
@@ -208,6 +222,7 @@ impl Tpu {
             tpu_quic_t,
             find_packet_sender_stake_stage,
             vote_find_packet_sender_stake_stage,
+            staked_nodes_updater_service,
         }
     }
 
@@ -236,6 +251,7 @@ impl Tpu {
             self.banking_stage.join(),
             self.find_packet_sender_stake_stage.join(),
             self.vote_find_packet_sender_stake_stage.join(),
+            self.staked_nodes_updater_service.join(),
         ];
         self.tpu_quic_t.join()?;
         let broadcast_result = self.broadcast_stage.join();