Leader qos part 2 - add stage to find packet's sender stake

core/benches/sigverify_stage.rs

@@ -159,7 +159,7 @@ fn bench_sigverify_stage(bencher: &mut Bencher) {
         for _ in 0..batches.len() {
             if let Some(batch) = batches.pop() {
                 sent_len += batch.packets.len();
-                packet_s.send(batch).unwrap();
+                packet_s.send(vec![batch]).unwrap();
             }
         }
         let mut received = 0;

core/src/find_packet_sender_stake_stage.rs

@@ -0,0 +1,185 @@
+use {
+    crossbeam_channel::{Receiver, RecvTimeoutError, Sender},
+    rayon::{prelude::*, ThreadPool},
+    solana_gossip::cluster_info::ClusterInfo,
+    solana_measure::measure::Measure,
+    solana_perf::packet::PacketBatch,
+    solana_rayon_threadlimit::get_thread_count,
+    solana_runtime::bank_forks::BankForks,
+    solana_sdk::timing::timestamp,
+    solana_streamer::streamer::{self, StreamerError},
+    std::{
+        cell::RefCell,
+        collections::HashMap,
+        net::IpAddr,
+        sync::{Arc, RwLock},
+        thread::{self, Builder, JoinHandle},
+        time::{Duration, Instant},
+    },
+};
+
+const IP_TO_STAKE_REFRESH_DURATION: Duration = Duration::from_secs(5);
+
+thread_local!(static PAR_THREAD_POOL: RefCell<ThreadPool> = RefCell::new(rayon::ThreadPoolBuilder::new()
+                    .num_threads(get_thread_count())
+                    .thread_name(|ix| format!("transaction_sender_stake_stage_{}", ix))
+                    .build()
+                    .unwrap()));
+
+pub type FindPacketSenderStakeSender = Sender<Vec<PacketBatch>>;
+pub type FindPacketSenderStakeReceiver = Receiver<Vec<PacketBatch>>;
+
+#[derive(Debug, Default)]
+struct FindPacketSenderStakeStats {
+    last_print: u64,
+    refresh_ip_to_stake_time: u64,
+    apply_sender_stakes_time: u64,
+    send_batches_time: u64,
+    receive_batches_time: u64,
+    total_batches: u64,
+    total_packets: u64,
+}
+
+impl FindPacketSenderStakeStats {
+    fn report(&mut self) {
+        let now = timestamp();
+        let elapsed_ms = now - self.last_print;
+        if elapsed_ms > 2000 {
+            datapoint_info!(
+                "find_packet_sender_stake-services_stats",
+                (
+                    "refresh_ip_to_stake_time",
+                    self.refresh_ip_to_stake_time as i64,
+                    i64
+                ),
+                (
+                    "apply_sender_stakes_time",
+                    self.apply_sender_stakes_time as i64,
+                    i64
+                ),
+                ("send_batches_time", self.send_batches_time as i64, i64),
+                (
+                    "receive_batches_time",
+                    self.receive_batches_time as i64,
+                    i64
+                ),
+                ("total_batches", self.total_batches as i64, i64),
+                ("total_packets", self.total_packets as i64, i64),
+            );
+            *self = FindPacketSenderStakeStats::default();
+            self.last_print = now;
+        }
+    }
+}
+
+pub struct FindPacketSenderStakeStage {
+    thread_hdl: JoinHandle<()>,
+}
+
+impl FindPacketSenderStakeStage {
+    pub fn new(
+        packet_receiver: streamer::PacketBatchReceiver,
+        sender: FindPacketSenderStakeSender,
+        bank_forks: Arc<RwLock<BankForks>>,
+        cluster_info: Arc<ClusterInfo>,
+    ) -> Self {
+        let mut stats = FindPacketSenderStakeStats::default();
+        let thread_hdl = Builder::new()
+            .name("find-packet-sender-stake".to_string())
+            .spawn(move || {
+                let mut last_stakes = Instant::now();
+                let mut ip_to_stake: HashMap<IpAddr, u64> = HashMap::new();
+                loop {
+                    let mut refresh_ip_to_stake_time = Measure::start("refresh_ip_to_stake_time");
+                    Self::try_refresh_ip_to_stake(
+                        &mut last_stakes,
+                        &mut ip_to_stake,
+                        bank_forks.clone(),
+                        cluster_info.clone(),
+                    );
+                    refresh_ip_to_stake_time.stop();
+                    stats.refresh_ip_to_stake_time = stats
+                        .refresh_ip_to_stake_time
+                        .saturating_add(refresh_ip_to_stake_time.as_us());
+
+                    match streamer::recv_packet_batches(&packet_receiver) {
+                        Ok((mut batches, num_packets, recv_duration)) => {
+                            let num_batches = batches.len();
+                            let mut apply_sender_stakes_time =
+                                Measure::start("apply_sender_stakes_time");
+                            Self::apply_sender_stakes(&mut batches, &ip_to_stake);
+                            apply_sender_stakes_time.stop();
+
+                            let mut send_batches_time = Measure::start("send_batches_time");
+                            if let Err(e) = sender.send(batches) {
+                                info!("Sender error: {:?}", e);
+                            }
+                            send_batches_time.stop();
+
+                            stats.apply_sender_stakes_time = stats
+                                .apply_sender_stakes_time
+                                .saturating_add(apply_sender_stakes_time.as_us());
+                            stats.send_batches_time = stats
+                                .send_batches_time
+                                .saturating_add(send_batches_time.as_us());
+                            stats.receive_batches_time = stats
+                                .receive_batches_time
+                                .saturating_add(recv_duration.as_nanos() as u64);
+                            stats.total_batches =
+                                stats.total_batches.saturating_add(num_batches as u64);
+                            stats.total_packets =
+                                stats.total_packets.saturating_add(num_packets as u64);
+                        }
+                        Err(e) => match e {
+                            StreamerError::RecvTimeout(RecvTimeoutError::Disconnected) => break,
+                            StreamerError::RecvTimeout(RecvTimeoutError::Timeout) => (),
+                            _ => error!("error: {:?}", e),
+                        },
+                    }
+
+                    stats.report();
+                }
+            })
+            .unwrap();
+        Self { thread_hdl }
+    }
+
+    fn try_refresh_ip_to_stake(
+        last_stakes: &mut Instant,
+        ip_to_stake: &mut HashMap<IpAddr, u64>,
+        bank_forks: Arc<RwLock<BankForks>>,
+        cluster_info: Arc<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();
+        }
+    }
+
+    fn apply_sender_stakes(batches: &mut [PacketBatch], ip_to_stake: &HashMap<IpAddr, u64>) {
+        PAR_THREAD_POOL.with(|thread_pool| {
+            thread_pool.borrow().install(|| {
+                batches
+                    .into_par_iter()
+                    .flat_map(|batch| batch.packets.par_iter_mut())
+                    .for_each(|packet| {
+                        packet.meta.sender_stake =
+                            *ip_to_stake.get(&packet.meta.addr().ip()).unwrap_or(&0);
+                    });
+            })
+        });
+    }
+
+    pub fn join(self) -> thread::Result<()> {
+        self.thread_hdl.join()
+    }
+}

core/src/lib.rs

@@ -24,6 +24,7 @@ pub mod cost_update_service;
 pub mod drop_bank_service;
 pub mod duplicate_repair_status;
 pub mod fetch_stage;
+pub mod find_packet_sender_stake_stage;
 pub mod fork_choice;
 pub mod gen_keys;
 pub mod heaviest_subtree_fork_choice;

core/src/shred_fetch_stage.rs

@@ -2,17 +2,17 @@
 
 use {
     crate::packet_hasher::PacketHasher,
-    crossbeam_channel::unbounded,
+    crossbeam_channel::{unbounded, Sender},
     lru::LruCache,
     solana_ledger::shred::{get_shred_slot_index_type, ShredFetchStats},
     solana_perf::{
         cuda_runtime::PinnedVec,
-        packet::{Packet, PacketBatchRecycler, PacketFlags},
+        packet::{Packet, PacketBatch, PacketBatchRecycler, PacketFlags},
         recycler::Recycler,
     },
     solana_runtime::bank_forks::BankForks,
     solana_sdk::clock::{Slot, DEFAULT_MS_PER_SLOT},
-    solana_streamer::streamer::{self, PacketBatchReceiver, PacketBatchSender},
+    solana_streamer::streamer::{self, PacketBatchReceiver},
     std::{
         net::UdpSocket,
         sync::{atomic::AtomicBool, Arc, RwLock},
@@ -65,7 +65,7 @@ impl ShredFetchStage {
     // updates packets received on a channel and sends them on another channel
     fn modify_packets<F>(
         recvr: PacketBatchReceiver,
-        sendr: PacketBatchSender,
+        sendr: Sender<Vec<PacketBatch>>,
         bank_forks: Option<Arc<RwLock<BankForks>>>,
         name: &'static str,
         modify: F,
@@ -125,7 +125,7 @@ impl ShredFetchStage {
                 stats = ShredFetchStats::default();
                 last_stats = Instant::now();
             }
-            if sendr.send(packet_batch).is_err() {
+            if sendr.send(vec![packet_batch]).is_err() {
                 break;
             }
         }
@@ -134,7 +134,7 @@ impl ShredFetchStage {
     fn packet_modifier<F>(
         sockets: Vec<Arc<UdpSocket>>,
         exit: &Arc<AtomicBool>,
-        sender: PacketBatchSender,
+        sender: Sender<Vec<PacketBatch>>,
         recycler: Recycler<PinnedVec<Packet>>,
         bank_forks: Option<Arc<RwLock<BankForks>>>,
         name: &'static str,
@@ -170,7 +170,7 @@ impl ShredFetchStage {
         sockets: Vec<Arc<UdpSocket>>,
         forward_sockets: Vec<Arc<UdpSocket>>,
         repair_socket: Arc<UdpSocket>,
-        sender: &PacketBatchSender,
+        sender: &Sender<Vec<PacketBatch>>,
         bank_forks: Option<Arc<RwLock<BankForks>>>,
         exit: &Arc<AtomicBool>,
     ) -> Self {

core/src/sigverify_stage.rs

@@ -6,17 +6,17 @@
 //! if perf-libs are available
 
 use {
-    crate::sigverify,
+    crate::{find_packet_sender_stake_stage, sigverify},
     core::time::Duration,
-    crossbeam_channel::{Receiver, RecvTimeoutError, SendError, Sender},
+    crossbeam_channel::{RecvTimeoutError, SendError, Sender},
     itertools::Itertools,
     solana_measure::measure::Measure,
     solana_perf::{
         packet::PacketBatch,
         sigverify::{count_valid_packets, shrink_batches, Deduper},
     },
     solana_sdk::timing,
-    solana_streamer::streamer::{self, PacketBatchReceiver, StreamerError},
+    solana_streamer::streamer::{self, StreamerError},
     std::{
         thread::{self, Builder, JoinHandle},
         time::Instant,
@@ -192,7 +192,7 @@ impl SigVerifier for DisabledSigVerifier {
 impl SigVerifyStage {
     #[allow(clippy::new_ret_no_self)]
     pub fn new<T: SigVerifier + 'static + Send + Clone>(
-        packet_receiver: Receiver<PacketBatch>,
+        packet_receiver: find_packet_sender_stake_stage::FindPacketSenderStakeReceiver,
         verified_sender: Sender<Vec<PacketBatch>>,
         verifier: T,
     ) -> Self {
@@ -227,12 +227,12 @@ impl SigVerifyStage {
 
     fn verifier<T: SigVerifier>(
         deduper: &Deduper,
-        recvr: &PacketBatchReceiver,
+        recvr: &find_packet_sender_stake_stage::FindPacketSenderStakeReceiver,
         sendr: &Sender<Vec<PacketBatch>>,
         verifier: &T,
         stats: &mut SigVerifierStats,
     ) -> Result<()> {
-        let (mut batches, num_packets, recv_duration) = streamer::recv_packet_batches(recvr)?;
+        let (mut batches, num_packets, recv_duration) = streamer::recv_vec_packet_batches(recvr)?;
 
         let batches_len = batches.len();
         debug!(
@@ -312,7 +312,7 @@ impl SigVerifyStage {
     }
 
     fn verifier_service<T: SigVerifier + 'static + Send + Clone>(
-        packet_receiver: PacketBatchReceiver,
+        packet_receiver: find_packet_sender_stake_stage::FindPacketSenderStakeReceiver,
         verified_sender: Sender<Vec<PacketBatch>>,
         verifier: &T,
     ) -> JoinHandle<()> {
@@ -358,7 +358,7 @@ impl SigVerifyStage {
     }
 
     fn verifier_services<T: SigVerifier + 'static + Send + Clone>(
-        packet_receiver: PacketBatchReceiver,
+        packet_receiver: find_packet_sender_stake_stage::FindPacketSenderStakeReceiver,
         verified_sender: Sender<Vec<PacketBatch>>,
         verifier: T,
     ) -> JoinHandle<()> {
@@ -445,7 +445,7 @@ mod tests {
         for _ in 0..batches.len() {
             if let Some(batch) = batches.pop() {
                 sent_len += batch.packets.len();
-                packet_s.send(batch).unwrap();
+                packet_s.send(vec![batch]).unwrap();
             }
         }
         let mut received = 0;