Skip to main content

miraland_streamer/
streamer.rs

1//! The `streamer` module defines a set of services for efficiently pulling data from UDP sockets.
2//!
3
4use {
5    crate::{
6        packet::{self, PacketBatch, PacketBatchRecycler, PACKETS_PER_BATCH},
7        sendmmsg::{batch_send, SendPktsError},
8        socket::SocketAddrSpace,
9    },
10    crossbeam_channel::{Receiver, RecvTimeoutError, SendError, Sender},
11    histogram::Histogram,
12    itertools::Itertools,
13    miraland_sdk::{packet::Packet, pubkey::Pubkey, timing::timestamp},
14    std::{
15        cmp::Reverse,
16        collections::HashMap,
17        net::{IpAddr, UdpSocket},
18        sync::{
19            atomic::{AtomicBool, AtomicUsize, Ordering},
20            Arc,
21        },
22        thread::{sleep, Builder, JoinHandle},
23        time::{Duration, Instant},
24    },
25    thiserror::Error,
26};
27
28// Total stake and nodes => stake map
29#[derive(Default)]
30pub struct StakedNodes {
31    stakes: Arc<HashMap<Pubkey, u64>>,
32    overrides: HashMap<Pubkey, u64>,
33    total_stake: u64,
34    max_stake: u64,
35    min_stake: u64,
36}
37
38pub type PacketBatchReceiver = Receiver<PacketBatch>;
39pub type PacketBatchSender = Sender<PacketBatch>;
40
41#[derive(Error, Debug)]
42pub enum StreamerError {
43    #[error("I/O error")]
44    Io(#[from] std::io::Error),
45
46    #[error("receive timeout error")]
47    RecvTimeout(#[from] RecvTimeoutError),
48
49    #[error("send packets error")]
50    Send(#[from] SendError<PacketBatch>),
51
52    #[error(transparent)]
53    SendPktsError(#[from] SendPktsError),
54}
55
56pub struct StreamerReceiveStats {
57    pub name: &'static str,
58    pub packets_count: AtomicUsize,
59    pub packet_batches_count: AtomicUsize,
60    pub full_packet_batches_count: AtomicUsize,
61    pub max_channel_len: AtomicUsize,
62}
63
64impl StreamerReceiveStats {
65    pub fn new(name: &'static str) -> Self {
66        Self {
67            name,
68            packets_count: AtomicUsize::default(),
69            packet_batches_count: AtomicUsize::default(),
70            full_packet_batches_count: AtomicUsize::default(),
71            max_channel_len: AtomicUsize::default(),
72        }
73    }
74
75    pub fn report(&self) {
76        datapoint_info!(
77            self.name,
78            (
79                "packets_count",
80                self.packets_count.swap(0, Ordering::Relaxed) as i64,
81                i64
82            ),
83            (
84                "packet_batches_count",
85                self.packet_batches_count.swap(0, Ordering::Relaxed) as i64,
86                i64
87            ),
88            (
89                "full_packet_batches_count",
90                self.full_packet_batches_count.swap(0, Ordering::Relaxed) as i64,
91                i64
92            ),
93            (
94                "channel_len",
95                self.max_channel_len.swap(0, Ordering::Relaxed) as i64,
96                i64
97            ),
98        );
99    }
100}
101
102pub type Result<T> = std::result::Result<T, StreamerError>;
103
104fn recv_loop(
105    socket: &UdpSocket,
106    exit: &AtomicBool,
107    packet_batch_sender: &PacketBatchSender,
108    recycler: &PacketBatchRecycler,
109    stats: &StreamerReceiveStats,
110    coalesce: Duration,
111    use_pinned_memory: bool,
112    in_vote_only_mode: Option<Arc<AtomicBool>>,
113) -> Result<()> {
114    loop {
115        let mut packet_batch = if use_pinned_memory {
116            PacketBatch::new_with_recycler(recycler, PACKETS_PER_BATCH, stats.name)
117        } else {
118            PacketBatch::with_capacity(PACKETS_PER_BATCH)
119        };
120        loop {
121            // Check for exit signal, even if socket is busy
122            // (for instance the leader transaction socket)
123            if exit.load(Ordering::Relaxed) {
124                return Ok(());
125            }
126
127            if let Some(ref in_vote_only_mode) = in_vote_only_mode {
128                if in_vote_only_mode.load(Ordering::Relaxed) {
129                    sleep(Duration::from_millis(1));
130                    continue;
131                }
132            }
133
134            if let Ok(len) = packet::recv_from(&mut packet_batch, socket, coalesce) {
135                if len > 0 {
136                    let StreamerReceiveStats {
137                        packets_count,
138                        packet_batches_count,
139                        full_packet_batches_count,
140                        max_channel_len,
141                        ..
142                    } = stats;
143
144                    packets_count.fetch_add(len, Ordering::Relaxed);
145                    packet_batches_count.fetch_add(1, Ordering::Relaxed);
146                    max_channel_len.fetch_max(packet_batch_sender.len(), Ordering::Relaxed);
147                    if len == PACKETS_PER_BATCH {
148                        full_packet_batches_count.fetch_add(1, Ordering::Relaxed);
149                    }
150
151                    packet_batch_sender.send(packet_batch)?;
152                }
153                break;
154            }
155        }
156    }
157}
158
159pub fn receiver(
160    socket: Arc<UdpSocket>,
161    exit: Arc<AtomicBool>,
162    packet_batch_sender: PacketBatchSender,
163    recycler: PacketBatchRecycler,
164    stats: Arc<StreamerReceiveStats>,
165    coalesce: Duration,
166    use_pinned_memory: bool,
167    in_vote_only_mode: Option<Arc<AtomicBool>>,
168) -> JoinHandle<()> {
169    let res = socket.set_read_timeout(Some(Duration::new(1, 0)));
170    assert!(res.is_ok(), "streamer::receiver set_read_timeout error");
171    Builder::new()
172        .name("mlnReceiver".to_string())
173        .spawn(move || {
174            let _ = recv_loop(
175                &socket,
176                &exit,
177                &packet_batch_sender,
178                &recycler,
179                &stats,
180                coalesce,
181                use_pinned_memory,
182                in_vote_only_mode,
183            );
184        })
185        .unwrap()
186}
187
188#[derive(Debug, Default)]
189struct SendStats {
190    bytes: u64,
191    count: u64,
192}
193
194#[derive(Default)]
195struct StreamerSendStats {
196    host_map: HashMap<IpAddr, SendStats>,
197    since: Option<Instant>,
198}
199
200impl StreamerSendStats {
201    fn report_stats(
202        name: &'static str,
203        host_map: HashMap<IpAddr, SendStats>,
204        sample_duration: Option<Duration>,
205    ) {
206        const MAX_REPORT_ENTRIES: usize = 5;
207        let sample_ms = sample_duration.map(|d| d.as_millis()).unwrap_or_default();
208        let mut hist = Histogram::default();
209        let mut byte_sum = 0;
210        let mut pkt_count = 0;
211        host_map.iter().for_each(|(_addr, host_stats)| {
212            hist.increment(host_stats.bytes).unwrap();
213            byte_sum += host_stats.bytes;
214            pkt_count += host_stats.count;
215        });
216
217        datapoint_info!(
218            name,
219            ("streamer-send-sample_duration_ms", sample_ms, i64),
220            ("streamer-send-host_count", host_map.len(), i64),
221            ("streamer-send-bytes_total", byte_sum, i64),
222            ("streamer-send-pkt_count_total", pkt_count, i64),
223            (
224                "streamer-send-host_bytes_min",
225                hist.minimum().unwrap_or_default(),
226                i64
227            ),
228            (
229                "streamer-send-host_bytes_max",
230                hist.maximum().unwrap_or_default(),
231                i64
232            ),
233            (
234                "streamer-send-host_bytes_mean",
235                hist.mean().unwrap_or_default(),
236                i64
237            ),
238            (
239                "streamer-send-host_bytes_90pct",
240                hist.percentile(90.0).unwrap_or_default(),
241                i64
242            ),
243            (
244                "streamer-send-host_bytes_50pct",
245                hist.percentile(50.0).unwrap_or_default(),
246                i64
247            ),
248            (
249                "streamer-send-host_bytes_10pct",
250                hist.percentile(10.0).unwrap_or_default(),
251                i64
252            ),
253        );
254
255        let num_entries = host_map.len();
256        let mut entries: Vec<_> = host_map.into_iter().collect();
257        if entries.len() > MAX_REPORT_ENTRIES {
258            entries.select_nth_unstable_by_key(MAX_REPORT_ENTRIES, |(_addr, stats)| {
259                Reverse(stats.bytes)
260            });
261            entries.truncate(MAX_REPORT_ENTRIES);
262        }
263        info!(
264            "streamer send {} hosts: count:{} {:?}",
265            name, num_entries, entries,
266        );
267    }
268
269    fn maybe_submit(&mut self, name: &'static str, sender: &Sender<Box<dyn FnOnce() + Send>>) {
270        const SUBMIT_CADENCE: Duration = Duration::from_secs(10);
271        const MAP_SIZE_REPORTING_THRESHOLD: usize = 1_000;
272        let elapsed = self.since.as_ref().map(Instant::elapsed);
273        if elapsed.map(|e| e < SUBMIT_CADENCE).unwrap_or_default()
274            && self.host_map.len() < MAP_SIZE_REPORTING_THRESHOLD
275        {
276            return;
277        }
278
279        let host_map = std::mem::take(&mut self.host_map);
280        let _ = sender.send(Box::new(move || {
281            Self::report_stats(name, host_map, elapsed);
282        }));
283
284        *self = Self {
285            since: Some(Instant::now()),
286            ..Self::default()
287        };
288    }
289
290    fn record(&mut self, pkt: &Packet) {
291        let ent = self.host_map.entry(pkt.meta().addr).or_default();
292        ent.count += 1;
293        ent.bytes += pkt.data(..).map(<[u8]>::len).unwrap_or_default() as u64;
294    }
295}
296
297impl StakedNodes {
298    pub fn new(stakes: Arc<HashMap<Pubkey, u64>>, overrides: HashMap<Pubkey, u64>) -> Self {
299        let values = stakes
300            .iter()
301            .filter(|(pubkey, _)| !overrides.contains_key(pubkey))
302            .map(|(_, &stake)| stake)
303            .chain(overrides.values().copied())
304            .filter(|&stake| stake > 0);
305        let total_stake = values.clone().sum();
306        let (min_stake, max_stake) = values.minmax().into_option().unwrap_or_default();
307        Self {
308            stakes,
309            overrides,
310            total_stake,
311            max_stake,
312            min_stake,
313        }
314    }
315
316    pub fn get_node_stake(&self, pubkey: &Pubkey) -> Option<u64> {
317        self.overrides
318            .get(pubkey)
319            .or_else(|| self.stakes.get(pubkey))
320            .filter(|&&stake| stake > 0)
321            .copied()
322    }
323
324    #[inline]
325    pub fn total_stake(&self) -> u64 {
326        self.total_stake
327    }
328
329    #[inline]
330    pub(super) fn min_stake(&self) -> u64 {
331        self.min_stake
332    }
333
334    #[inline]
335    pub(super) fn max_stake(&self) -> u64 {
336        self.max_stake
337    }
338}
339
340fn recv_send(
341    sock: &UdpSocket,
342    r: &PacketBatchReceiver,
343    socket_addr_space: &SocketAddrSpace,
344    stats: &mut Option<StreamerSendStats>,
345) -> Result<()> {
346    let timer = Duration::new(1, 0);
347    let packet_batch = r.recv_timeout(timer)?;
348    if let Some(stats) = stats {
349        packet_batch.iter().for_each(|p| stats.record(p));
350    }
351    let packets = packet_batch.iter().filter_map(|pkt| {
352        let addr = pkt.meta().socket_addr();
353        let data = pkt.data(..)?;
354        socket_addr_space.check(&addr).then_some((data, addr))
355    });
356    batch_send(sock, &packets.collect::<Vec<_>>())?;
357    Ok(())
358}
359
360pub fn recv_packet_batches(
361    recvr: &PacketBatchReceiver,
362) -> Result<(Vec<PacketBatch>, usize, Duration)> {
363    let timer = Duration::new(1, 0);
364    let packet_batch = recvr.recv_timeout(timer)?;
365    let recv_start = Instant::now();
366    trace!("got packets");
367    let mut num_packets = packet_batch.len();
368    let mut packet_batches = vec![packet_batch];
369    while let Ok(packet_batch) = recvr.try_recv() {
370        trace!("got more packets");
371        num_packets += packet_batch.len();
372        packet_batches.push(packet_batch);
373    }
374    let recv_duration = recv_start.elapsed();
375    trace!(
376        "packet batches len: {}, num packets: {}",
377        packet_batches.len(),
378        num_packets
379    );
380    Ok((packet_batches, num_packets, recv_duration))
381}
382
383pub fn responder(
384    name: &'static str,
385    sock: Arc<UdpSocket>,
386    r: PacketBatchReceiver,
387    socket_addr_space: SocketAddrSpace,
388    stats_reporter_sender: Option<Sender<Box<dyn FnOnce() + Send>>>,
389) -> JoinHandle<()> {
390    Builder::new()
391        .name(format!("mlnRspndr{name}"))
392        .spawn(move || {
393            let mut errors = 0;
394            let mut last_error = None;
395            let mut last_print = 0;
396            let mut stats = None;
397
398            if stats_reporter_sender.is_some() {
399                stats = Some(StreamerSendStats::default());
400            }
401
402            loop {
403                if let Err(e) = recv_send(&sock, &r, &socket_addr_space, &mut stats) {
404                    match e {
405                        StreamerError::RecvTimeout(RecvTimeoutError::Disconnected) => break,
406                        StreamerError::RecvTimeout(RecvTimeoutError::Timeout) => (),
407                        _ => {
408                            errors += 1;
409                            last_error = Some(e);
410                        }
411                    }
412                }
413                let now = timestamp();
414                if now - last_print > 1000 && errors != 0 {
415                    datapoint_info!(name, ("errors", errors, i64),);
416                    info!("{} last-error: {:?} count: {}", name, last_error, errors);
417                    last_print = now;
418                    errors = 0;
419                }
420                if let Some(ref stats_reporter_sender) = stats_reporter_sender {
421                    if let Some(ref mut stats) = stats {
422                        stats.maybe_submit(name, stats_reporter_sender);
423                    }
424                }
425            }
426        })
427        .unwrap()
428}
429
430#[cfg(test)]
431mod test {
432    use {
433        super::*,
434        crate::{
435            packet::{Packet, PacketBatch, PACKET_DATA_SIZE},
436            streamer::{receiver, responder},
437        },
438        crossbeam_channel::unbounded,
439        miraland_perf::recycler::Recycler,
440        std::{
441            io,
442            io::Write,
443            net::UdpSocket,
444            sync::{
445                atomic::{AtomicBool, Ordering},
446                Arc,
447            },
448            time::Duration,
449        },
450    };
451
452    fn get_packet_batches(r: PacketBatchReceiver, num_packets: &mut usize) {
453        for _ in 0..10 {
454            let packet_batch_res = r.recv_timeout(Duration::new(1, 0));
455            if packet_batch_res.is_err() {
456                continue;
457            }
458
459            *num_packets -= packet_batch_res.unwrap().len();
460
461            if *num_packets == 0 {
462                break;
463            }
464        }
465    }
466
467    #[test]
468    fn streamer_debug() {
469        write!(io::sink(), "{:?}", Packet::default()).unwrap();
470        write!(io::sink(), "{:?}", PacketBatch::default()).unwrap();
471    }
472    #[test]
473    fn streamer_send_test() {
474        let read = UdpSocket::bind("127.0.0.1:0").expect("bind");
475        read.set_read_timeout(Some(Duration::new(1, 0))).unwrap();
476
477        let addr = read.local_addr().unwrap();
478        let send = UdpSocket::bind("127.0.0.1:0").expect("bind");
479        let exit = Arc::new(AtomicBool::new(false));
480        let (s_reader, r_reader) = unbounded();
481        let stats = Arc::new(StreamerReceiveStats::new("test"));
482        let t_receiver = receiver(
483            Arc::new(read),
484            exit.clone(),
485            s_reader,
486            Recycler::default(),
487            stats.clone(),
488            Duration::from_millis(1), // coalesce
489            true,
490            None,
491        );
492        const NUM_PACKETS: usize = 5;
493        let t_responder = {
494            let (s_responder, r_responder) = unbounded();
495            let t_responder = responder(
496                "SendTest",
497                Arc::new(send),
498                r_responder,
499                SocketAddrSpace::Unspecified,
500                None,
501            );
502            let mut packet_batch = PacketBatch::default();
503            for i in 0..NUM_PACKETS {
504                let mut p = Packet::default();
505                {
506                    p.buffer_mut()[0] = i as u8;
507                    p.meta_mut().size = PACKET_DATA_SIZE;
508                    p.meta_mut().set_socket_addr(&addr);
509                }
510                packet_batch.push(p);
511            }
512            s_responder.send(packet_batch).expect("send");
513            t_responder
514        };
515
516        let mut packets_remaining = NUM_PACKETS;
517        get_packet_batches(r_reader, &mut packets_remaining);
518        assert_eq!(packets_remaining, 0);
519        exit.store(true, Ordering::Relaxed);
520        assert!(stats.packet_batches_count.load(Ordering::Relaxed) >= 1);
521        assert_eq!(stats.packets_count.load(Ordering::Relaxed), NUM_PACKETS);
522        assert_eq!(stats.full_packet_batches_count.load(Ordering::Relaxed), 0);
523        t_receiver.join().expect("join");
524        t_responder.join().expect("join");
525    }
526}