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 priority_ready_packets(&self) -> usize {
454        self.pending_priority
455            .as_ref()
456            .map_or(0, PendingPackets::len)
457            .saturating_add(self.priority_queued_packets())
458    }
459
460    pub async fn recv(&mut self) -> Option<ReceivedPacket> {
461        loop {
462            match self.try_recv() {
463                Ok(packet) => return Some(packet),
464                Err(TryRecvError::Disconnected) => return None,
465                Err(TryRecvError::Empty) => {}
466            }
467
468            tokio::select! {
469                biased;
470                item = self.priority.recv(), if !self.priority_closed => {
471                    match item {
472                        Some(item) => {
473                            if let Some(packet) = self.packet_from_item(item, PacketLane::Priority) {
474                                return Some(packet);
475                            }
476                        }
477                        None => self.priority_closed = true,
478                    }
479                }
480                item = self.bulk.recv(), if !self.bulk_closed => {
481                    match item {
482                        Some(item) => {
483                            if let Some(packet) = self.packet_from_item(item, PacketLane::Bulk) {
484                                return Some(packet);
485                            }
486                        }
487                        None => self.bulk_closed = true,
488                    }
489                }
490            }
491        }
492    }
493
494    pub fn try_recv(&mut self) -> Result<ReceivedPacket, TryRecvError> {
495        if let Some(packet) = Self::take_pending(&mut self.pending_priority) {
496            return Ok(packet);
497        }
498
499        if self.should_probe_priority() {
500            match self.priority.try_recv() {
501                Ok(item) => {
502                    if let Some(packet) = self.packet_from_item(item, PacketLane::Priority) {
503                        return Ok(packet);
504                    }
505                }
506                Err(TryRecvError::Empty) => {}
507                Err(TryRecvError::Disconnected) => {
508                    self.priority_closed = true;
509                }
510            }
511        }
512
513        if let Some(packet) = Self::take_pending(&mut self.pending_bulk) {
514            return Ok(packet);
515        }
516
517        match self.bulk.try_recv() {
518            Ok(item) => self
519                .packet_from_item(item, PacketLane::Bulk)
520                .ok_or(TryRecvError::Empty),
521            Err(TryRecvError::Empty) => {
522                if self.priority_closed && self.bulk_closed {
523                    Err(TryRecvError::Disconnected)
524                } else {
525                    Err(TryRecvError::Empty)
526                }
527            }
528            Err(TryRecvError::Disconnected) => {
529                self.bulk_closed = true;
530                if self.priority_closed {
531                    Err(TryRecvError::Disconnected)
532                } else {
533                    Err(TryRecvError::Empty)
534                }
535            }
536        }
537    }
538
539    fn packet_from_item(
540        &mut self,
541        item: PacketQueueItem,
542        lane: PacketLane,
543    ) -> Option<ReceivedPacket> {
544        item.record_dequeue_wait(lane);
545        let packet_count = item.packet_count();
546        if self.track_backlog {
547            self.queued_packets.fetch_sub(packet_count, Relaxed);
548        }
549        if matches!(lane, PacketLane::Priority) {
550            release_priority_packets(&self.priority_queued_packets, packet_count);
551        }
552        if matches!(lane, PacketLane::Bulk) {
553            release_reserved_bulk_packets(&self.bulk_queued_packets, packet_count);
554        }
555        let rx_loop_owned_at = crate::perf_profile::stamp();
556        match item {
557            PacketQueueItem::One(mut packet) => {
558                packet.trace_rx_loop_owned_at = rx_loop_owned_at;
559                Some(packet)
560            }
561            PacketQueueItem::Batch(packets) => {
562                let mut packets = packets.into_iter();
563                let mut packet = packets.next()?;
564                if let Some(rx_loop_owned_at) = rx_loop_owned_at {
565                    packet.trace_rx_loop_owned_at = Some(rx_loop_owned_at);
566                }
567                if packets.len() > 0 {
568                    let pending = PendingPackets::new(packets, rx_loop_owned_at);
569                    match lane {
570                        PacketLane::Priority => self.pending_priority = Some(pending),
571                        PacketLane::Bulk => self.pending_bulk = Some(pending),
572                    }
573                }
574                Some(packet)
575            }
576        }
577    }
578
579    fn should_probe_priority(&self) -> bool {
580        !self.priority_closed
581            && (self.priority_queued_packets.load(Relaxed) > 0 || self.bulk_closed)
582    }
583
584    fn take_pending(pending: &mut Option<PendingPackets>) -> Option<ReceivedPacket> {
585        let packets = pending.as_mut()?;
586        let packet = packets.next();
587        if packets.len() == 0 {
588            *pending = None;
589        }
590        packet
591    }
592}
593
594#[inline]
595fn packet_channel_tracks_backlog() -> bool {
596    cfg!(test) || crate::perf_profile::enabled()
597}
598
599fn release_reserved_bulk_packets(counter: &AtomicUsize, count: usize) {
600    if count == 0 {
601        return;
602    }
603
604    let previous = counter.fetch_sub(count, Relaxed);
605    debug_assert!(
606        previous >= count,
607        "transport bulk queued packet accounting underflow"
608    );
609}
610
611fn release_priority_packets(counter: &AtomicUsize, count: usize) {
612    if count == 0 {
613        return;
614    }
615
616    let previous = counter.fetch_sub(count, Relaxed);
617    debug_assert!(
618        previous >= count,
619        "transport priority queued packet accounting underflow"
620    );
621}
622
623/// Create a packet channel.
624///
625/// The capacity applies to bulk packets. Priority traffic is intentionally
626/// unbounded so small control-shaped packets can still wake the rx loop while a
627/// bulk receiver is saturated.
628pub fn packet_channel(buffer: usize) -> (PacketTx, PacketRx) {
629    let (priority_tx, priority_rx) = tokio::sync::mpsc::unbounded_channel();
630    let (bulk_tx, bulk_rx) = tokio::sync::mpsc::channel(buffer.max(1));
631    let priority_queued_packets = Arc::new(AtomicUsize::new(0));
632    let queued_packets = Arc::new(AtomicUsize::new(0));
633    let bulk_queued_packets = Arc::new(AtomicUsize::new(0));
634    let track_backlog = packet_channel_tracks_backlog();
635    (
636        PacketTx {
637            priority: priority_tx,
638            bulk: bulk_tx,
639            priority_queued_packets: Arc::clone(&priority_queued_packets),
640            queued_packets: Arc::clone(&queued_packets),
641            bulk_queued_packets: Arc::clone(&bulk_queued_packets),
642            bulk_packet_capacity: buffer.max(1),
643            track_backlog,
644        },
645        PacketRx {
646            priority: priority_rx,
647            bulk: bulk_rx,
648            priority_queued_packets,
649            queued_packets,
650            bulk_queued_packets,
651            track_backlog,
652            pending_priority: None,
653            pending_bulk: None,
654            priority_closed: false,
655            bulk_closed: false,
656        },
657    )
658}
659
660#[cfg(test)]
661mod tests;
fips_core/transport/packet_channel.rs

fips_core/transport/
packet_channel.rs
