1use {
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#[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 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), 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}