fips_core/transport/
packet_channel.rs

1//! Priority-aware packet channel for transport receive paths.
2
3use super::{TransportAddr, TransportId};
4use std::sync::{
5    Arc,
6    atomic::{AtomicUsize, Ordering::Relaxed},
7};
8use std::time::{SystemTime, UNIX_EPOCH};
9use std::vec::IntoIter;
10use tokio::sync::mpsc::{
11    Sender, UnboundedReceiver, UnboundedSender,
12    error::{TryRecvError, TrySendError},
13};
14
15/// A packet received from a transport.
16#[derive(Clone, Debug)]
17pub struct ReceivedPacket {
18    /// Which transport received this packet.
19    pub transport_id: TransportId,
20    /// Remote peer address.
21    pub remote_addr: TransportAddr,
22    /// Packet data.
23    pub data: Vec<u8>,
24    /// Receipt timestamp (Unix milliseconds).
25    pub timestamp_ms: u64,
26    /// Monotonic timestamp for optional pipeline queue-wait profiling.
27    #[doc(hidden)]
28    pub trace_enqueued_at: Option<crate::perf_profile::TraceStamp>,
29    /// Monotonic timestamp captured when `PacketRx` takes ownership of a
30    /// channel item. Distinguishes mpsc/channel residence from rx-loop-owned
31    /// batch-tail residence in perf traces.
32    #[doc(hidden)]
33    pub trace_rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
34}
35
36impl ReceivedPacket {
37    /// Create a new received packet with current timestamp.
38    pub fn new(transport_id: TransportId, remote_addr: TransportAddr, data: Vec<u8>) -> Self {
39        Self::with_trace_timestamp(
40            transport_id,
41            remote_addr,
42            data,
43            received_timestamp_ms(),
44            crate::perf_profile::stamp(),
45        )
46    }
47
48    /// Create a received packet with explicit timestamp.
49    pub fn with_timestamp(
50        transport_id: TransportId,
51        remote_addr: TransportAddr,
52        data: Vec<u8>,
53        timestamp_ms: u64,
54    ) -> Self {
55        Self::with_trace_timestamp(
56            transport_id,
57            remote_addr,
58            data,
59            timestamp_ms,
60            crate::perf_profile::stamp(),
61        )
62    }
63
64    /// Create a received packet with explicit wall-clock and queue timestamps.
65    ///
66    /// UDP receive paths can drain several datagrams per kernel batch. Those
67    /// datagrams arrived close together, so sharing one wall-clock sample and
68    /// one queue trace stamp across the batch avoids per-packet clock reads
69    /// while preserving arrival order and queue residence visibility.
70    pub(crate) fn with_trace_timestamp(
71        transport_id: TransportId,
72        remote_addr: TransportAddr,
73        data: Vec<u8>,
74        timestamp_ms: u64,
75        trace_enqueued_at: Option<crate::perf_profile::TraceStamp>,
76    ) -> Self {
77        Self {
78            transport_id,
79            remote_addr,
80            data,
81            timestamp_ms,
82            trace_enqueued_at,
83            trace_rx_loop_owned_at: None,
84        }
85    }
86
87    pub(crate) fn is_priority_sized(&self) -> bool {
88        self.data.len() <= PRIORITY_PACKET_MAX_LEN
89    }
90}
91
92pub(crate) fn received_timestamp_ms() -> u64 {
93    SystemTime::now()
94        .duration_since(UNIX_EPOCH)
95        .map(|d| d.as_millis() as u64)
96        .unwrap_or(0)
97}
98
99/// Wire-size threshold for keeping transport receive work out of the bulk
100/// FIFO. Most heartbeat, MMP, rekey, ping, and handshake-shaped datagrams are
101/// comfortably below this; full-size endpoint payloads are not.
102const PRIORITY_PACKET_MAX_LEN: usize = 512;
103
104/// Packet count at which the transport receive channel is visibly backlogged.
105///
106/// This tracks packets still owned by the priority/bulk mpsc channels. Once a
107/// batched item is dequeued into `PacketRx`'s pending iterator, it no longer
108/// contributes to this counter; those packets are already inside the rx-loop
109/// owner's drain budget rather than waiting behind the transport channel.
110const TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER: usize = 4096;
111
112/// Channel sender for received packets.
113///
114/// The priority lane stays unbounded because control-shaped datagrams must keep
115/// making progress even when bulk is saturated. The bulk lane is bounded by the
116/// configured packet-channel capacity in packets, not receive-batch items, and
117/// uses nonblocking `try_send`: overload sheds bulk explicitly instead of
118/// hiding unbounded latency behind the rx loop.
119#[derive(Clone, Debug)]
120pub struct PacketTx {
121    priority: UnboundedSender<PacketQueueItem>,
122    bulk: Sender<PacketQueueItem>,
123    /// Packet-count ready hint for priority lane probes. Bulk batch tails check
124    /// this instead of touching an empty priority mpsc once per data packet.
125    priority_queued_packets: Arc<AtomicUsize>,
126    queued_packets: Arc<AtomicUsize>,
127    bulk_queued_packets: Arc<AtomicUsize>,
128    bulk_packet_capacity: usize,
129    track_backlog: bool,
130}
131
132/// Channel receiver for received packets.
133pub struct PacketRx {
134    priority: UnboundedReceiver<PacketQueueItem>,
135    bulk: tokio::sync::mpsc::Receiver<PacketQueueItem>,
136    priority_queued_packets: Arc<AtomicUsize>,
137    queued_packets: Arc<AtomicUsize>,
138    bulk_queued_packets: Arc<AtomicUsize>,
139    track_backlog: bool,
140    pending_priority: Option<PendingPackets>,
141    pending_bulk: Option<PendingPackets>,
142    priority_closed: bool,
143    bulk_closed: bool,
144}
145
146#[derive(Debug)]
147enum PacketQueueItem {
148    One(ReceivedPacket),
149    #[cfg_attr(not(any(target_os = "linux", target_os = "macos")), allow(dead_code))]
150    Batch(Vec<ReceivedPacket>),
151}
152
153#[derive(Clone, Copy)]
154enum PacketLane {
155    Priority,
156    Bulk,
157}
158
159#[derive(Clone, Copy)]
160enum PacketQueueTx {
161    Priority,
162    Bulk,
163}
164
165enum PacketSendFailure {
166    Closed(PacketQueueItem),
167    DroppedBulk(usize),
168}
169
170struct PendingPackets {
171    packets: IntoIter<ReceivedPacket>,
172    rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
173}
174
175#[derive(Debug, PartialEq, Eq)]
176struct PacketQueueDequeueCounts {
177    total: usize,
178    priority: usize,
179    bulk: usize,
180}
181
182impl PacketQueueTx {
183    fn try_send(self, owner: &PacketTx, item: PacketQueueItem) -> Result<(), PacketSendFailure> {
184        match self {
185            PacketQueueTx::Priority => owner
186                .priority
187                .send(item)
188                .map_err(|error| PacketSendFailure::Closed(error.0)),
189            PacketQueueTx::Bulk => {
190                let packet_count = item.packet_count();
191                match owner.bulk.try_send(item) {
192                    Ok(()) => Ok(()),
193                    Err(TrySendError::Full(_item)) => {
194                        Err(PacketSendFailure::DroppedBulk(packet_count))
195                    }
196                    Err(TrySendError::Closed(item)) => Err(PacketSendFailure::Closed(item)),
197                }
198            }
199        }
200    }
201}
202
203impl PacketQueueItem {
204    fn packet_count(&self) -> usize {
205        match self {
206            PacketQueueItem::One(_) => 1,
207            PacketQueueItem::Batch(packets) => packets.len(),
208        }
209    }
210
211    fn dequeue_counts(&self, lane: PacketLane) -> PacketQueueDequeueCounts {
212        let total = self.packet_count();
213        match lane {
214            PacketLane::Priority => PacketQueueDequeueCounts {
215                total,
216                priority: total,
217                bulk: 0,
218            },
219            PacketLane::Bulk => PacketQueueDequeueCounts {
220                total,
221                priority: 0,
222                bulk: total,
223            },
224        }
225    }
226
227    fn queued_at(&self) -> Option<crate::perf_profile::TraceStamp> {
228        match self {
229            PacketQueueItem::One(packet) => packet.trace_enqueued_at,
230            PacketQueueItem::Batch(packets) => {
231                packets.first().and_then(|packet| packet.trace_enqueued_at)
232            }
233        }
234    }
235
236    fn record_dequeue_wait(&self, lane: PacketLane) {
237        let queued_at = self.queued_at();
238        if queued_at.is_none() {
239            return;
240        }
241        let counts = self.dequeue_counts(lane);
242        crate::perf_profile::record_since_split_count(
243            crate::perf_profile::Stage::TransportChannelWait,
244            crate::perf_profile::Stage::TransportPriorityChannelWait,
245            crate::perf_profile::Stage::TransportBulkChannelWait,
246            queued_at,
247            counts.total as u64,
248            counts.priority as u64,
249            counts.bulk as u64,
250        );
251    }
252}
253
254impl PendingPackets {
255    fn new(
256        packets: IntoIter<ReceivedPacket>,
257        rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
258    ) -> Self {
259        Self {
260            packets,
261            rx_loop_owned_at,
262        }
263    }
264
265    fn next(&mut self) -> Option<ReceivedPacket> {
266        let mut packet = self.packets.next()?;
267        if let Some(rx_loop_owned_at) = self.rx_loop_owned_at {
268            packet.trace_rx_loop_owned_at = Some(rx_loop_owned_at);
269        }
270        Some(packet)
271    }
272
273    fn len(&self) -> usize {
274        self.packets.len()
275    }
276}
277
278impl PacketTx {
279    pub fn send(
280        &self,
281        packet: ReceivedPacket,
282    ) -> Result<(), tokio::sync::mpsc::error::SendError<ReceivedPacket>> {
283        let tx = if packet.data.len() <= PRIORITY_PACKET_MAX_LEN {
284            PacketQueueTx::Priority
285        } else {
286            PacketQueueTx::Bulk
287        };
288        self.send_item(tx, PacketQueueItem::One(packet))
289            .map_err(|item| match item {
290                PacketQueueItem::One(packet) => tokio::sync::mpsc::error::SendError(packet),
291                PacketQueueItem::Batch(_) => {
292                    unreachable!("single packet send cannot fail with a batch item")
293                }
294            })
295    }
296
297    #[cfg_attr(not(any(target_os = "linux", target_os = "macos")), allow(dead_code))]
298    pub(crate) fn send_batch(&self, packets: Vec<ReceivedPacket>) -> Result<(), ()> {
299        if packets.is_empty() {
300            return Ok(());
301        }
302
303        let packet_count = packets.len();
304        let priority_count = packets
305            .iter()
306            .filter(|packet| packet.is_priority_sized())
307            .count();
308        if priority_count == 0 || priority_count == packet_count {
309            let tx = if priority_count == 0 {
310                PacketQueueTx::Bulk
311            } else {
312                PacketQueueTx::Priority
313            };
314            return self.send_packet_items(tx, packets);
315        }
316
317        let mut priority_packets = Vec::with_capacity(priority_count);
318        let mut bulk_packets = Vec::with_capacity(packet_count - priority_count);
319        for packet in packets {
320            if packet.is_priority_sized() {
321                priority_packets.push(packet);
322            } else {
323                bulk_packets.push(packet);
324            }
325        }
326
327        self.send_packet_items(PacketQueueTx::Priority, priority_packets)?;
328        self.send_packet_items(PacketQueueTx::Bulk, bulk_packets)?;
329        Ok(())
330    }
331
332    #[cfg_attr(not(any(target_os = "linux", target_os = "macos")), allow(dead_code))]
333    fn send_packet_items(
334        &self,
335        tx: PacketQueueTx,
336        mut packets: Vec<ReceivedPacket>,
337    ) -> Result<(), ()> {
338        let item = match packets.len() {
339            0 => return Ok(()),
340            1 => PacketQueueItem::One(packets.pop().expect("one packet should be present")),
341            _ => PacketQueueItem::Batch(packets),
342        };
343        self.send_item(tx, item).map_err(|_| ())
344    }
345
346    fn send_item(&self, tx: PacketQueueTx, item: PacketQueueItem) -> Result<(), PacketQueueItem> {
347        let packet_count = item.packet_count();
348        let bulk_reserved = matches!(tx, PacketQueueTx::Bulk)
349            .then_some(packet_count)
350            .filter(|count| *count > 0);
351        let priority_reserved = matches!(tx, PacketQueueTx::Priority)
352            .then_some(packet_count)
353            .filter(|count| *count > 0);
354        if let Some(count) = priority_reserved {
355            self.priority_queued_packets.fetch_add(count, Relaxed);
356        }
357        if let Some(count) = bulk_reserved
358            && !self.try_reserve_bulk_packets(count)
359        {
360            crate::perf_profile::record_event_count(
361                crate::perf_profile::Event::TransportBulkDropped,
362                count as u64,
363            );
364            return Ok(());
365        }
366
367        let tracked_count = if self.track_backlog {
368            Some(packet_count)
369        } else {
370            None
371        };
372        let previous = tracked_count.map(|count| self.queued_packets.fetch_add(count, Relaxed));
373        match tx.try_send(self, item) {
374            Ok(()) => {
375                if let (Some(count), Some(previous)) = (tracked_count, previous) {
376                    let queued = previous.saturating_add(count);
377                    if previous < TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER
378                        && queued >= TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER
379                    {
380                        crate::perf_profile::record_event(
381                            crate::perf_profile::Event::TransportChannelBacklogHigh,
382                        );
383                    }
384                }
385                Ok(())
386            }
387            Err(PacketSendFailure::Closed(item)) => {
388                if let Some(count) = tracked_count {
389                    self.queued_packets.fetch_sub(count, Relaxed);
390                }
391                if let Some(count) = priority_reserved {
392                    release_priority_packets(&self.priority_queued_packets, count);
393                }
394                if let Some(count) = bulk_reserved {
395                    self.release_bulk_packets(count);
396                }
397                Err(item)
398            }
399            Err(PacketSendFailure::DroppedBulk(dropped_count)) => {
400                if let Some(count) = tracked_count {
401                    self.queued_packets.fetch_sub(count, Relaxed);
402                }
403                if let Some(count) = priority_reserved {
404                    release_priority_packets(&self.priority_queued_packets, count);
405                }
406                if let Some(count) = bulk_reserved {
407                    self.release_bulk_packets(count);
408                }
409                crate::perf_profile::record_event_count(
410                    crate::perf_profile::Event::TransportBulkDropped,
411                    dropped_count as u64,
412                );
413                Ok(())
414            }
415        }
416    }
417
418    fn try_reserve_bulk_packets(&self, count: usize) -> bool {
419        self.bulk_queued_packets
420            .fetch_update(Relaxed, Relaxed, |current| {
421                current
422                    .checked_add(count)
423                    .filter(|next| *next <= self.bulk_packet_capacity)
424            })
425            .is_ok()
426    }
427
428    fn release_bulk_packets(&self, count: usize) {
429        release_reserved_bulk_packets(&self.bulk_queued_packets, count);
430    }
431
432    #[cfg(test)]
433    pub(crate) fn queued_packets(&self) -> usize {
434        self.queued_packets.load(Relaxed)
435    }
436
437    #[cfg(test)]
438    pub(crate) fn priority_queued_packets(&self) -> usize {
439        self.priority_queued_packets.load(Relaxed)
440    }
441
442    #[cfg(test)]
443    pub(crate) fn bulk_queued_packets(&self) -> usize {
444        self.bulk_queued_packets.load(Relaxed)
445    }
446}
447
448impl PacketRx {
449    pub(crate) fn priority_queued_packets(&self) -> usize {
450        self.priority_queued_packets.load(Relaxed)
451    }
452
453    pub(crate) fn ready_packets(&self) -> usize {
454        self.pending_priority
455            .as_ref()
456            .map_or(0, PendingPackets::len)
457            .saturating_add(self.pending_bulk.as_ref().map_or(0, PendingPackets::len))
458            .saturating_add(self.priority_queued_packets())
459            .saturating_add(self.bulk_queued_packets.load(Relaxed))
460    }
461
462    pub(crate) fn priority_ready_packets(&self) -> usize {
463        self.pending_priority
464            .as_ref()
465            .map_or(0, PendingPackets::len)
466            .saturating_add(self.priority_queued_packets())
467    }
468
469    pub async fn recv(&mut self) -> Option<ReceivedPacket> {
470        loop {
471            match self.try_recv() {
472                Ok(packet) => return Some(packet),
473                Err(TryRecvError::Disconnected) => return None,
474                Err(TryRecvError::Empty) => {}
475            }
476
477            tokio::select! {
478                biased;
479                item = self.priority.recv(), if !self.priority_closed => {
480                    match item {
481                        Some(item) => {
482                            if let Some(packet) = self.packet_from_item(item, PacketLane::Priority) {
483                                return Some(packet);
484                            }
485                        }
486                        None => self.priority_closed = true,
487                    }
488                }
489                item = self.bulk.recv(), if !self.bulk_closed => {
490                    match item {
491                        Some(item) => {
492                            if let Some(packet) = self.packet_from_item(item, PacketLane::Bulk) {
493                                return Some(packet);
494                            }
495                        }
496                        None => self.bulk_closed = true,
497                    }
498                }
499            }
500        }
501    }
502
503    pub fn try_recv(&mut self) -> Result<ReceivedPacket, TryRecvError> {
504        if let Some(packet) = Self::take_pending(&mut self.pending_priority) {
505            return Ok(packet);
506        }
507
508        if self.should_probe_priority() {
509            match self.priority.try_recv() {
510                Ok(item) => {
511                    if let Some(packet) = self.packet_from_item(item, PacketLane::Priority) {
512                        return Ok(packet);
513                    }
514                }
515                Err(TryRecvError::Empty) => {}
516                Err(TryRecvError::Disconnected) => {
517                    self.priority_closed = true;
518                }
519            }
520        }
521
522        if let Some(packet) = Self::take_pending(&mut self.pending_bulk) {
523            return Ok(packet);
524        }
525
526        match self.bulk.try_recv() {
527            Ok(item) => self
528                .packet_from_item(item, PacketLane::Bulk)
529                .ok_or(TryRecvError::Empty),
530            Err(TryRecvError::Empty) => {
531                if self.priority_closed && self.bulk_closed {
532                    Err(TryRecvError::Disconnected)
533                } else {
534                    Err(TryRecvError::Empty)
535                }
536            }
537            Err(TryRecvError::Disconnected) => {
538                self.bulk_closed = true;
539                if self.priority_closed {
540                    Err(TryRecvError::Disconnected)
541                } else {
542                    Err(TryRecvError::Empty)
543                }
544            }
545        }
546    }
547
548    fn packet_from_item(
549        &mut self,
550        item: PacketQueueItem,
551        lane: PacketLane,
552    ) -> Option<ReceivedPacket> {
553        item.record_dequeue_wait(lane);
554        let packet_count = item.packet_count();
555        if self.track_backlog {
556            self.queued_packets.fetch_sub(packet_count, Relaxed);
557        }
558        if matches!(lane, PacketLane::Priority) {
559            release_priority_packets(&self.priority_queued_packets, packet_count);
560        }
561        if matches!(lane, PacketLane::Bulk) {
562            release_reserved_bulk_packets(&self.bulk_queued_packets, packet_count);
563        }
564        let rx_loop_owned_at = crate::perf_profile::stamp();
565        match item {
566            PacketQueueItem::One(mut packet) => {
567                packet.trace_rx_loop_owned_at = rx_loop_owned_at;
568                Some(packet)
569            }
570            PacketQueueItem::Batch(packets) => {
571                let mut packets = packets.into_iter();
572                let mut packet = packets.next()?;
573                if let Some(rx_loop_owned_at) = rx_loop_owned_at {
574                    packet.trace_rx_loop_owned_at = Some(rx_loop_owned_at);
575                }
576                if packets.len() > 0 {
577                    let pending = PendingPackets::new(packets, rx_loop_owned_at);
578                    match lane {
579                        PacketLane::Priority => self.pending_priority = Some(pending),
580                        PacketLane::Bulk => self.pending_bulk = Some(pending),
581                    }
582                }
583                Some(packet)
584            }
585        }
586    }
587
588    fn should_probe_priority(&self) -> bool {
589        !self.priority_closed
590            && (self.priority_queued_packets.load(Relaxed) > 0 || self.bulk_closed)
591    }
592
593    fn take_pending(pending: &mut Option<PendingPackets>) -> Option<ReceivedPacket> {
594        let packets = pending.as_mut()?;
595        let packet = packets.next();
596        if packets.len() == 0 {
597            *pending = None;
598        }
599        packet
600    }
601}
602
603#[inline]
604fn packet_channel_tracks_backlog() -> bool {
605    cfg!(test) || crate::perf_profile::enabled()
606}
607
608fn release_reserved_bulk_packets(counter: &AtomicUsize, count: usize) {
609    if count == 0 {
610        return;
611    }
612
613    let previous = counter.fetch_sub(count, Relaxed);
614    debug_assert!(
615        previous >= count,
616        "transport bulk queued packet accounting underflow"
617    );
618}
619
620fn release_priority_packets(counter: &AtomicUsize, count: usize) {
621    if count == 0 {
622        return;
623    }
624
625    let previous = counter.fetch_sub(count, Relaxed);
626    debug_assert!(
627        previous >= count,
628        "transport priority queued packet accounting underflow"
629    );
630}
631
632/// Create a packet channel.
633///
634/// The capacity applies to bulk packets. Priority traffic is intentionally
635/// unbounded so small control-shaped packets can still wake the rx loop while a
636/// bulk receiver is saturated.
637pub fn packet_channel(buffer: usize) -> (PacketTx, PacketRx) {
638    let (priority_tx, priority_rx) = tokio::sync::mpsc::unbounded_channel();
639    let (bulk_tx, bulk_rx) = tokio::sync::mpsc::channel(buffer.max(1));
640    let priority_queued_packets = Arc::new(AtomicUsize::new(0));
641    let queued_packets = Arc::new(AtomicUsize::new(0));
642    let bulk_queued_packets = Arc::new(AtomicUsize::new(0));
643    let track_backlog = packet_channel_tracks_backlog();
644    (
645        PacketTx {
646            priority: priority_tx,
647            bulk: bulk_tx,
648            priority_queued_packets: Arc::clone(&priority_queued_packets),
649            queued_packets: Arc::clone(&queued_packets),
650            bulk_queued_packets: Arc::clone(&bulk_queued_packets),
651            bulk_packet_capacity: buffer.max(1),
652            track_backlog,
653        },
654        PacketRx {
655            priority: priority_rx,
656            bulk: bulk_rx,
657            priority_queued_packets,
658            queued_packets,
659            bulk_queued_packets,
660            track_backlog,
661            pending_priority: None,
662            pending_bulk: None,
663            priority_closed: false,
664            bulk_closed: false,
665        },
666    )
667}
668
669#[cfg(test)]
670mod tests;
fips_core/transport/packet_channel.rs

fips_core/transport/
packet_channel.rs
