Skip to main content

fips_core/transport/
packet_channel.rs

1//! Priority-aware packet channel for transport receive paths.
2
3use super::{TransportAddr, TransportId};
4use std::mem;
5use std::ops::{Deref, DerefMut, Index};
6use std::sync::{
7    Arc, Mutex,
8    atomic::{AtomicUsize, Ordering::Relaxed},
9};
10use std::time::{SystemTime, UNIX_EPOCH};
11use tokio::sync::mpsc::{
12    Sender, UnboundedReceiver, UnboundedSender,
13    error::{TryRecvError, TrySendError},
14};
15
16pub(crate) trait PacketFastIngressSink: std::fmt::Debug + Send + Sync {
17    fn try_ingest_batch(&self, packets: &mut Vec<ReceivedPacket>) -> usize;
18}
19
20/// A packet received from a transport.
21#[derive(Clone, Debug)]
22pub struct ReceivedPacket {
23    /// Which transport received this packet.
24    pub transport_id: TransportId,
25    /// Remote peer address.
26    pub remote_addr: TransportAddr,
27    /// Packet data.
28    pub data: PacketBuffer,
29    /// Receipt timestamp (Unix milliseconds).
30    pub timestamp_ms: u64,
31    /// Monotonic timestamp for optional pipeline queue-wait profiling.
32    #[doc(hidden)]
33    pub trace_enqueued_at: Option<crate::perf_profile::TraceStamp>,
34    /// Monotonic timestamp captured when `PacketRx` takes ownership of a
35    /// channel item. Distinguishes mpsc/channel residence from rx-loop-owned
36    /// batch-tail residence in perf traces.
37    #[doc(hidden)]
38    pub trace_rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
39}
40
41impl ReceivedPacket {
42    /// Create a new received packet with current timestamp.
43    pub fn new(
44        transport_id: TransportId,
45        remote_addr: TransportAddr,
46        data: impl Into<PacketBuffer>,
47    ) -> Self {
48        Self::with_trace_timestamp(
49            transport_id,
50            remote_addr,
51            data,
52            received_timestamp_ms(),
53            crate::perf_profile::stamp(),
54        )
55    }
56
57    /// Create a received packet with explicit timestamp.
58    pub fn with_timestamp(
59        transport_id: TransportId,
60        remote_addr: TransportAddr,
61        data: impl Into<PacketBuffer>,
62        timestamp_ms: u64,
63    ) -> Self {
64        Self::with_trace_timestamp(
65            transport_id,
66            remote_addr,
67            data,
68            timestamp_ms,
69            crate::perf_profile::stamp(),
70        )
71    }
72
73    /// Create a received packet with explicit wall-clock and queue timestamps.
74    ///
75    /// UDP receive paths can drain several datagrams per kernel batch. Those
76    /// datagrams arrived close together, so sharing one wall-clock sample and
77    /// one queue trace stamp across the batch avoids per-packet clock reads
78    /// while preserving arrival order and queue residence visibility.
79    pub(crate) fn with_trace_timestamp(
80        transport_id: TransportId,
81        remote_addr: TransportAddr,
82        data: impl Into<PacketBuffer>,
83        timestamp_ms: u64,
84        trace_enqueued_at: Option<crate::perf_profile::TraceStamp>,
85    ) -> Self {
86        Self {
87            transport_id,
88            remote_addr,
89            data: data.into(),
90            timestamp_ms,
91            trace_enqueued_at,
92            trace_rx_loop_owned_at: None,
93        }
94    }
95
96    pub(crate) fn is_transport_priority(&self) -> bool {
97        is_transport_priority_packet(&self.data)
98    }
99}
100
101/// Byte storage for a received transport packet.
102///
103/// Receive/decrypt/drop paths carry this owner so pressure drops and endpoint
104/// delivery can recycle kernel receive buffers without an extra packet copy.
105#[derive(Debug, Default)]
106pub struct PacketBuffer {
107    data: Vec<u8>,
108    pool: Option<PacketBufferPool>,
109}
110
111impl PacketBuffer {
112    fn pooled(data: Vec<u8>, pool: PacketBufferPool) -> Self {
113        Self {
114            data,
115            pool: Some(pool),
116        }
117    }
118
119    pub fn new(data: Vec<u8>) -> Self {
120        Self { data, pool: None }
121    }
122
123    pub fn as_slice(&self) -> &[u8] {
124        &self.data
125    }
126
127    pub fn as_mut_slice(&mut self) -> &mut [u8] {
128        &mut self.data
129    }
130
131    pub fn len(&self) -> usize {
132        self.data.len()
133    }
134
135    pub fn is_empty(&self) -> bool {
136        self.data.is_empty()
137    }
138
139    pub fn into_vec(mut self) -> Vec<u8> {
140        self.pool = None;
141        mem::take(&mut self.data)
142    }
143
144    pub(crate) fn try_prepend_slices(&mut self, parts: &[&[u8]], reserve_tail: usize) -> bool {
145        let prefix_len = parts
146            .iter()
147            .fold(0usize, |total, part| total.saturating_add(part.len()));
148        if prefix_len == 0 {
149            return self.data.capacity().saturating_sub(self.data.len()) >= reserve_tail;
150        }
151
152        let len = self.data.len();
153        if self.data.capacity().saturating_sub(len) < prefix_len.saturating_add(reserve_tail) {
154            return false;
155        }
156
157        // Move the packet body right inside the existing allocation, then fill
158        // the newly opened header space. This is the Vec equivalent of the
159        // fixed headroom WireGuard-go keeps in its message buffers.
160        unsafe {
161            let ptr = self.data.as_mut_ptr();
162            std::ptr::copy(ptr, ptr.add(prefix_len), len);
163            let mut offset = 0usize;
164            for part in parts {
165                std::ptr::copy_nonoverlapping(part.as_ptr(), ptr.add(offset), part.len());
166                offset += part.len();
167            }
168            self.data.set_len(len + prefix_len);
169        }
170        true
171    }
172}
173
174impl Clone for PacketBuffer {
175    fn clone(&self) -> Self {
176        Self {
177            data: self.data.clone(),
178            pool: None,
179        }
180    }
181}
182
183impl Drop for PacketBuffer {
184    fn drop(&mut self) {
185        if let Some(pool) = self.pool.take() {
186            pool.put(mem::take(&mut self.data));
187        }
188    }
189}
190
191impl From<Vec<u8>> for PacketBuffer {
192    fn from(data: Vec<u8>) -> Self {
193        Self::new(data)
194    }
195}
196
197impl From<PacketBuffer> for Vec<u8> {
198    fn from(buffer: PacketBuffer) -> Self {
199        buffer.into_vec()
200    }
201}
202
203impl Deref for PacketBuffer {
204    type Target = Vec<u8>;
205
206    fn deref(&self) -> &Self::Target {
207        &self.data
208    }
209}
210
211impl DerefMut for PacketBuffer {
212    fn deref_mut(&mut self) -> &mut Self::Target {
213        &mut self.data
214    }
215}
216
217impl AsRef<[u8]> for PacketBuffer {
218    fn as_ref(&self) -> &[u8] {
219        &self.data
220    }
221}
222
223impl AsMut<[u8]> for PacketBuffer {
224    fn as_mut(&mut self) -> &mut [u8] {
225        &mut self.data
226    }
227}
228
229impl PartialEq for PacketBuffer {
230    fn eq(&self, other: &Self) -> bool {
231        self.data == other.data
232    }
233}
234
235impl Eq for PacketBuffer {}
236
237impl PartialEq<Vec<u8>> for PacketBuffer {
238    fn eq(&self, other: &Vec<u8>) -> bool {
239        self.data == *other
240    }
241}
242
243impl PartialEq<PacketBuffer> for Vec<u8> {
244    fn eq(&self, other: &PacketBuffer) -> bool {
245        *self == other.data
246    }
247}
248
249impl PartialEq<&[u8]> for PacketBuffer {
250    fn eq(&self, other: &&[u8]) -> bool {
251        self.data.as_slice() == *other
252    }
253}
254
255impl<const N: usize> PartialEq<[u8; N]> for PacketBuffer {
256    fn eq(&self, other: &[u8; N]) -> bool {
257        self.data.as_slice() == other
258    }
259}
260
261impl<const N: usize> PartialEq<&[u8; N]> for PacketBuffer {
262    fn eq(&self, other: &&[u8; N]) -> bool {
263        self.data.as_slice() == *other
264    }
265}
266
267pub(crate) fn received_timestamp_ms() -> u64 {
268    SystemTime::now()
269        .duration_since(UNIX_EPOCH)
270        .map(|d| d.as_millis() as u64)
271        .unwrap_or(0)
272}
273
274/// FMP packet shape that is visible before PM2 authenticates established data.
275///
276/// App payloads, TCP ACKs, pings, and established FSP/MMP frames are all opaque
277/// phase-0 data at this boundary. Only first-contact/link rekey handshakes get
278/// the unbounded reserve lane before PM2 can classify authenticated contents.
279const FMP_VERSION: u8 = 0;
280const FMP_PHASE_MSG1: u8 = 0x1;
281const FMP_PHASE_MSG2: u8 = 0x2;
282const FMP_COMMON_PREFIX_SIZE: usize = 4;
283const FMP_MSG1_WIRE_SIZE: usize = 114;
284const FMP_MSG2_WIRE_SIZE: usize = 69;
285
286fn is_transport_priority_packet(data: &[u8]) -> bool {
287    if data.len() < FMP_COMMON_PREFIX_SIZE {
288        return false;
289    }
290
291    let version = data[0] >> 4;
292    let phase = data[0] & 0x0F;
293    if version != FMP_VERSION {
294        return false;
295    }
296
297    matches!(
298        (phase, data.len()),
299        (FMP_PHASE_MSG1, FMP_MSG1_WIRE_SIZE) | (FMP_PHASE_MSG2, FMP_MSG2_WIRE_SIZE)
300    )
301}
302
303/// Number of receive-batch Vec containers retained for reuse.
304const PACKET_BATCH_POOL_LIMIT: usize = 256;
305/// Avoid pinning unusually large test/control batches in the hot-path pool.
306const PACKET_BATCH_MAX_RETAINED_CAPACITY: usize = 256;
307/// Number of packet byte buffers retained after pressure drops.
308const PACKET_BUFFER_POOL_LIMIT: usize = 4096;
309/// Avoid pinning oversized receive buffers in the hot-path pool.
310const PACKET_BUFFER_MAX_RETAINED_CAPACITY: usize = 16 * 1024;
311
312/// Packet count at which the transport receive channel is visibly backlogged.
313///
314/// This tracks packets still owned by the priority/bulk mpsc channels. Once a
315/// batched item is dequeued into `PacketRx`'s pending iterator, it no longer
316/// contributes to this counter; those packets are already inside the rx-loop
317/// owner's drain budget rather than waiting behind the transport channel.
318const TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER: usize = 4096;
319
320/// Channel sender for received packets.
321///
322/// The priority lane stays unbounded because control-shaped datagrams must keep
323/// making progress even when bulk is saturated. The bulk lane is bounded by the
324/// configured packet-channel capacity in packets, not receive-batch items, and
325/// uses nonblocking `try_send`: overload sheds bulk explicitly instead of
326/// hiding unbounded latency behind the rx loop.
327#[derive(Clone, Debug)]
328pub struct PacketTx {
329    priority: UnboundedSender<PacketQueueItem>,
330    bulk: Sender<PacketQueueItem>,
331    fast_ingress: Option<Arc<dyn PacketFastIngressSink>>,
332    batch_pool: PacketBatchPool,
333    buffer_pool: PacketBufferPool,
334    /// Packet-count ready hint for priority lane probes. Bulk batch tails check
335    /// this instead of touching an empty priority mpsc once per data packet.
336    priority_queued_packets: Arc<AtomicUsize>,
337    queued_packets: Arc<AtomicUsize>,
338    bulk_queued_packets: Arc<AtomicUsize>,
339    bulk_packet_capacity: usize,
340    track_backlog: bool,
341}
342
343/// Channel receiver for received packets.
344pub struct PacketRx {
345    priority: UnboundedReceiver<PacketQueueItem>,
346    bulk: tokio::sync::mpsc::Receiver<PacketQueueItem>,
347    priority_queued_packets: Arc<AtomicUsize>,
348    queued_packets: Arc<AtomicUsize>,
349    bulk_queued_packets: Arc<AtomicUsize>,
350    track_backlog: bool,
351    pending_priority: Option<PendingPackets>,
352    pending_bulk: Option<PendingPackets>,
353    priority_closed: bool,
354    bulk_closed: bool,
355}
356
357#[derive(Clone, Debug)]
358struct PacketBatchPool {
359    inner: Arc<Mutex<Vec<Vec<ReceivedPacket>>>>,
360}
361
362#[derive(Clone, Debug)]
363struct PacketBufferPool {
364    inner: Arc<Mutex<Vec<Vec<u8>>>>,
365    available: Arc<AtomicUsize>,
366}
367
368#[derive(Debug)]
369pub(crate) struct PacketBatch {
370    packets: Vec<ReceivedPacket>,
371    pool: Option<PacketBatchPool>,
372}
373
374#[derive(Debug)]
375enum PacketQueueItem {
376    One(ReceivedPacket),
377    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
378    Batch(PacketBatch),
379}
380
381#[derive(Clone, Copy)]
382enum PacketLane {
383    Priority,
384    Bulk,
385}
386
387#[derive(Clone, Copy)]
388enum PacketQueueTx {
389    Priority,
390    Bulk,
391}
392
393enum PacketSendFailure {
394    Closed(PacketQueueItem),
395    DroppedBulk(usize),
396}
397
398struct PendingPackets {
399    batch: PacketBatch,
400    rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
401}
402
403#[derive(Debug, PartialEq, Eq)]
404struct PacketQueueDequeueCounts {
405    total: usize,
406    priority: usize,
407    bulk: usize,
408}
409
410impl PacketQueueTx {
411    fn try_send(self, owner: &PacketTx, item: PacketQueueItem) -> Result<(), PacketSendFailure> {
412        match self {
413            PacketQueueTx::Priority => owner
414                .priority
415                .send(item)
416                .map_err(|error| PacketSendFailure::Closed(error.0)),
417            PacketQueueTx::Bulk => {
418                let packet_count = item.packet_count();
419                match owner.bulk.try_send(item) {
420                    Ok(()) => Ok(()),
421                    Err(TrySendError::Full(_item)) => {
422                        Err(PacketSendFailure::DroppedBulk(packet_count))
423                    }
424                    Err(TrySendError::Closed(item)) => Err(PacketSendFailure::Closed(item)),
425                }
426            }
427        }
428    }
429}
430
431impl PacketQueueItem {
432    fn packet_count(&self) -> usize {
433        match self {
434            PacketQueueItem::One(_) => 1,
435            PacketQueueItem::Batch(packets) => packets.len(),
436        }
437    }
438
439    fn dequeue_counts(&self, lane: PacketLane) -> PacketQueueDequeueCounts {
440        let total = self.packet_count();
441        match lane {
442            PacketLane::Priority => PacketQueueDequeueCounts {
443                total,
444                priority: total,
445                bulk: 0,
446            },
447            PacketLane::Bulk => PacketQueueDequeueCounts {
448                total,
449                priority: 0,
450                bulk: total,
451            },
452        }
453    }
454
455    fn queued_at(&self) -> Option<crate::perf_profile::TraceStamp> {
456        match self {
457            PacketQueueItem::One(packet) => packet.trace_enqueued_at,
458            PacketQueueItem::Batch(packets) => {
459                packets.first().and_then(|packet| packet.trace_enqueued_at)
460            }
461        }
462    }
463
464    fn record_dequeue_wait(&self, lane: PacketLane) {
465        let queued_at = self.queued_at();
466        if queued_at.is_none() {
467            return;
468        }
469        let counts = self.dequeue_counts(lane);
470        crate::perf_profile::record_since_split_count(
471            crate::perf_profile::Stage::TransportChannelWait,
472            crate::perf_profile::Stage::TransportPriorityChannelWait,
473            crate::perf_profile::Stage::TransportBulkChannelWait,
474            queued_at,
475            counts.total as u64,
476            counts.priority as u64,
477            counts.bulk as u64,
478        );
479    }
480}
481
482impl PacketBatchPool {
483    fn new() -> Self {
484        Self {
485            inner: Arc::new(Mutex::new(Vec::new())),
486        }
487    }
488
489    fn take(&self, capacity: usize) -> PacketBatch {
490        let packets = {
491            let mut guard = self.inner.lock().unwrap_or_else(|error| error.into_inner());
492            guard.pop()
493        };
494        if let Some(mut packets) = packets {
495            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBatchPoolReuse);
496            packets.clear();
497            if packets.capacity() >= capacity {
498                return PacketBatch::pooled(packets, self.clone());
499            }
500            packets.reserve(capacity.saturating_sub(packets.capacity()));
501            return PacketBatch::pooled(packets, self.clone());
502        }
503        crate::perf_profile::record_event(crate::perf_profile::Event::PacketBatchPoolFresh);
504        PacketBatch::pooled(Vec::with_capacity(capacity), self.clone())
505    }
506
507    fn put(&self, mut packets: Vec<ReceivedPacket>) {
508        packets.clear();
509        if packets.capacity() > PACKET_BATCH_MAX_RETAINED_CAPACITY {
510            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBatchPoolDiscard);
511            return;
512        }
513        let mut guard = self.inner.lock().unwrap_or_else(|error| error.into_inner());
514        if guard.len() < PACKET_BATCH_POOL_LIMIT {
515            guard.push(packets);
516            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBatchPoolReturn);
517        } else {
518            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBatchPoolDiscard);
519        }
520    }
521
522    #[cfg(test)]
523    fn cached_len(&self) -> usize {
524        self.inner
525            .lock()
526            .unwrap_or_else(|error| error.into_inner())
527            .len()
528    }
529}
530
531impl PacketBufferPool {
532    fn new() -> Self {
533        Self {
534            inner: Arc::new(Mutex::new(Vec::new())),
535            available: Arc::new(AtomicUsize::new(0)),
536        }
537    }
538
539    fn take(&self, capacity: usize) -> Vec<u8> {
540        if self.available.load(Relaxed) > 0 {
541            let buffer = {
542                let mut guard = self.inner.lock().unwrap_or_else(|error| error.into_inner());
543                guard.pop()
544            };
545            if let Some(mut buffer) = buffer {
546                self.available.fetch_sub(1, Relaxed);
547                crate::perf_profile::record_event(
548                    crate::perf_profile::Event::PacketBufferPoolReuse,
549                );
550                prepare_recv_buffer(&mut buffer, capacity);
551                return buffer;
552            }
553        }
554
555        crate::perf_profile::record_event(crate::perf_profile::Event::PacketBufferPoolFresh);
556        fresh_recv_buffer(capacity)
557    }
558
559    fn put(&self, mut buffer: Vec<u8>) {
560        buffer.clear();
561        if buffer.capacity() > PACKET_BUFFER_MAX_RETAINED_CAPACITY {
562            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBufferPoolDiscard);
563            return;
564        }
565
566        let mut guard = self.inner.lock().unwrap_or_else(|error| error.into_inner());
567        if guard.len() < PACKET_BUFFER_POOL_LIMIT {
568            guard.push(buffer);
569            self.available.fetch_add(1, Relaxed);
570            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBufferPoolReturn);
571        } else {
572            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBufferPoolDiscard);
573        }
574    }
575
576    #[cfg(test)]
577    fn cached_len(&self) -> usize {
578        self.available.load(Relaxed)
579    }
580}
581
582#[cfg(target_os = "macos")]
583fn fresh_recv_buffer(size: usize) -> Vec<u8> {
584    vec![0u8; size]
585}
586
587#[cfg(not(target_os = "macos"))]
588fn fresh_recv_buffer(size: usize) -> Vec<u8> {
589    Vec::with_capacity(size)
590}
591
592#[cfg(target_os = "macos")]
593fn prepare_recv_buffer(buffer: &mut Vec<u8>, size: usize) {
594    buffer.resize(size, 0);
595}
596
597#[cfg(not(target_os = "macos"))]
598fn prepare_recv_buffer(buffer: &mut Vec<u8>, size: usize) {
599    buffer.clear();
600    if buffer.capacity() < size {
601        buffer.reserve(size.saturating_sub(buffer.capacity()));
602    }
603}
604
605impl PacketBatch {
606    fn from_vec(packets: Vec<ReceivedPacket>) -> Self {
607        Self {
608            packets,
609            pool: None,
610        }
611    }
612
613    fn pooled(packets: Vec<ReceivedPacket>, pool: PacketBatchPool) -> Self {
614        Self {
615            packets,
616            pool: Some(pool),
617        }
618    }
619
620    pub(crate) fn push(&mut self, packet: ReceivedPacket) {
621        self.packets.push(packet);
622    }
623
624    pub(crate) fn is_empty(&self) -> bool {
625        self.packets.is_empty()
626    }
627
628    fn len(&self) -> usize {
629        self.packets.len()
630    }
631
632    fn first(&self) -> Option<&ReceivedPacket> {
633        self.packets.first()
634    }
635
636    fn iter(&self) -> impl Iterator<Item = &ReceivedPacket> {
637        self.packets.iter()
638    }
639
640    fn drain(&mut self) -> impl Iterator<Item = ReceivedPacket> + '_ {
641        self.packets.drain(..)
642    }
643
644    fn pop(&mut self) -> Option<ReceivedPacket> {
645        self.packets.pop()
646    }
647
648    fn reverse(&mut self) {
649        self.packets.reverse();
650    }
651
652    fn is_pooled(&self) -> bool {
653        self.pool.is_some()
654    }
655}
656
657impl From<Vec<ReceivedPacket>> for PacketBatch {
658    fn from(packets: Vec<ReceivedPacket>) -> Self {
659        Self::from_vec(packets)
660    }
661}
662
663impl Index<usize> for PacketBatch {
664    type Output = ReceivedPacket;
665
666    fn index(&self, index: usize) -> &Self::Output {
667        &self.packets[index]
668    }
669}
670
671impl Drop for PacketBatch {
672    fn drop(&mut self) {
673        let Some(pool) = self.pool.take() else {
674            return;
675        };
676        pool.put(mem::take(&mut self.packets));
677    }
678}
679
680impl PendingPackets {
681    fn new(
682        mut batch: PacketBatch,
683        rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
684    ) -> Self {
685        batch.reverse();
686        Self {
687            batch,
688            rx_loop_owned_at,
689        }
690    }
691
692    fn next(&mut self) -> Option<ReceivedPacket> {
693        let mut packet = self.batch.pop()?;
694        if let Some(rx_loop_owned_at) = self.rx_loop_owned_at {
695            packet.trace_rx_loop_owned_at = Some(rx_loop_owned_at);
696        }
697        Some(packet)
698    }
699
700    fn len(&self) -> usize {
701        self.batch.len()
702    }
703}
704
705impl PacketTx {
706    pub(crate) fn set_fast_ingress_sink(&mut self, sink: Arc<dyn PacketFastIngressSink>) {
707        self.fast_ingress = Some(sink);
708    }
709
710    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
711    pub(crate) fn try_fast_ingress_packet_batch(&self, batch: &mut PacketBatch) -> usize {
712        let Some(sink) = &self.fast_ingress else {
713            return 0;
714        };
715        sink.try_ingest_batch(&mut batch.packets)
716    }
717
718    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
719    pub(crate) fn packet_batch(&self, capacity: usize) -> PacketBatch {
720        self.batch_pool.take(capacity)
721    }
722
723    #[cfg_attr(not(any(target_os = "linux", target_os = "macos")), allow(dead_code))]
724    pub(crate) fn recv_buffer(&self, capacity: usize) -> Vec<u8> {
725        self.buffer_pool.take(capacity)
726    }
727
728    #[cfg_attr(not(any(target_os = "linux", target_os = "macos")), allow(dead_code))]
729    pub(crate) fn packet_buffer(&self, data: Vec<u8>) -> PacketBuffer {
730        PacketBuffer::pooled(data, self.buffer_pool.clone())
731    }
732
733    pub fn send(
734        &self,
735        packet: ReceivedPacket,
736    ) -> Result<(), tokio::sync::mpsc::error::SendError<ReceivedPacket>> {
737        let tx = if packet.is_transport_priority() {
738            PacketQueueTx::Priority
739        } else {
740            PacketQueueTx::Bulk
741        };
742        self.send_item(tx, PacketQueueItem::One(packet))
743            .map_err(|item| match item {
744                PacketQueueItem::One(packet) => tokio::sync::mpsc::error::SendError(packet),
745                PacketQueueItem::Batch(_) => {
746                    unreachable!("single packet send cannot fail with a batch item")
747                }
748            })
749    }
750
751    #[cfg(test)]
752    pub(crate) fn send_batch(&self, packets: Vec<ReceivedPacket>) -> Result<(), ()> {
753        self.send_packet_batch(PacketBatch::from_vec(packets))
754    }
755
756    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
757    pub(crate) fn send_packet_batch(&self, mut batch: PacketBatch) -> Result<(), ()> {
758        if batch.is_empty() {
759            return Ok(());
760        }
761
762        let packet_count = batch.len();
763        let priority_count = batch
764            .iter()
765            .filter(|packet| packet.is_transport_priority())
766            .count();
767        if priority_count == 0 || priority_count == packet_count {
768            let tx = if priority_count == 0 {
769                PacketQueueTx::Bulk
770            } else {
771                PacketQueueTx::Priority
772            };
773            return self.send_packet_items(tx, batch);
774        }
775
776        let mut priority_packets = self.packet_batch(priority_count);
777        let mut bulk_packets = self.packet_batch(packet_count - priority_count);
778        for packet in batch.drain() {
779            if packet.is_transport_priority() {
780                priority_packets.push(packet);
781            } else {
782                bulk_packets.push(packet);
783            }
784        }
785
786        self.send_packet_items(PacketQueueTx::Priority, priority_packets)?;
787        self.send_packet_items(PacketQueueTx::Bulk, bulk_packets)?;
788        Ok(())
789    }
790
791    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
792    fn send_packet_items(&self, tx: PacketQueueTx, mut packets: PacketBatch) -> Result<(), ()> {
793        if matches!(tx, PacketQueueTx::Bulk) {
794            return self.send_bulk_packet_items(packets);
795        }
796
797        let item = match packets.len() {
798            0 => return Ok(()),
799            1 if !packets.is_pooled() => {
800                PacketQueueItem::One(packets.pop().expect("one packet should be present"))
801            }
802            _ => PacketQueueItem::Batch(packets),
803        };
804        self.send_item(tx, item).map_err(|_| ())
805    }
806
807    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
808    fn send_bulk_packet_items(&self, mut packets: PacketBatch) -> Result<(), ()> {
809        let packet_count = packets.len();
810        if packet_count == 0 {
811            return Ok(());
812        }
813
814        let granted = self.try_reserve_bulk_packet_prefix(packet_count);
815        if granted == 0 {
816            crate::perf_profile::record_event_count(
817                crate::perf_profile::Event::TransportBulkDropped,
818                packet_count as u64,
819            );
820            return Ok(());
821        }
822
823        if granted < packet_count {
824            let dropped = packet_count - granted;
825            let _dropped_tail = packets.packets.split_off(granted);
826            crate::perf_profile::record_event_count(
827                crate::perf_profile::Event::TransportBulkDropped,
828                dropped as u64,
829            );
830        }
831
832        let item = match packets.len() {
833            0 => return Ok(()),
834            1 if !packets.is_pooled() => {
835                PacketQueueItem::One(packets.pop().expect("one packet should be present"))
836            }
837            _ => PacketQueueItem::Batch(packets),
838        };
839        self.send_reserved_item(PacketQueueTx::Bulk, item, Some(granted))
840            .map_err(|_| ())
841    }
842
843    fn send_item(&self, tx: PacketQueueTx, item: PacketQueueItem) -> Result<(), PacketQueueItem> {
844        let packet_count = item.packet_count();
845        let bulk_reserved = if matches!(tx, PacketQueueTx::Bulk) && packet_count > 0 {
846            if !self.try_reserve_bulk_packets(packet_count) {
847                crate::perf_profile::record_event_count(
848                    crate::perf_profile::Event::TransportBulkDropped,
849                    packet_count as u64,
850                );
851                return Ok(());
852            }
853            Some(packet_count)
854        } else {
855            None
856        };
857        self.send_reserved_item(tx, item, bulk_reserved)
858    }
859
860    fn send_reserved_item(
861        &self,
862        tx: PacketQueueTx,
863        item: PacketQueueItem,
864        bulk_reserved: Option<usize>,
865    ) -> Result<(), PacketQueueItem> {
866        let packet_count = item.packet_count();
867        debug_assert_eq!(
868            bulk_reserved,
869            matches!(tx, PacketQueueTx::Bulk)
870                .then_some(packet_count)
871                .filter(|count| *count > 0)
872        );
873        let priority_reserved = matches!(tx, PacketQueueTx::Priority)
874            .then_some(packet_count)
875            .filter(|count| *count > 0);
876        if let Some(count) = priority_reserved {
877            self.priority_queued_packets.fetch_add(count, Relaxed);
878        }
879
880        let tracked_count = if self.track_backlog {
881            Some(packet_count)
882        } else {
883            None
884        };
885        let previous = tracked_count.map(|count| self.queued_packets.fetch_add(count, Relaxed));
886        match tx.try_send(self, item) {
887            Ok(()) => {
888                if let (Some(count), Some(previous)) = (tracked_count, previous) {
889                    let queued = previous.saturating_add(count);
890                    if previous < TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER
891                        && queued >= TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER
892                    {
893                        crate::perf_profile::record_event(
894                            crate::perf_profile::Event::TransportChannelBacklogHigh,
895                        );
896                    }
897                }
898                Ok(())
899            }
900            Err(PacketSendFailure::Closed(item)) => {
901                if let Some(count) = tracked_count {
902                    self.queued_packets.fetch_sub(count, Relaxed);
903                }
904                if let Some(count) = priority_reserved {
905                    release_priority_packets(&self.priority_queued_packets, count);
906                }
907                if let Some(count) = bulk_reserved {
908                    self.release_bulk_packets(count);
909                }
910                Err(item)
911            }
912            Err(PacketSendFailure::DroppedBulk(dropped_count)) => {
913                if let Some(count) = tracked_count {
914                    self.queued_packets.fetch_sub(count, Relaxed);
915                }
916                if let Some(count) = priority_reserved {
917                    release_priority_packets(&self.priority_queued_packets, count);
918                }
919                if let Some(count) = bulk_reserved {
920                    self.release_bulk_packets(count);
921                }
922                crate::perf_profile::record_event_count(
923                    crate::perf_profile::Event::TransportBulkDropped,
924                    dropped_count as u64,
925                );
926                Ok(())
927            }
928        }
929    }
930
931    fn try_reserve_bulk_packets(&self, count: usize) -> bool {
932        self.bulk_queued_packets
933            .fetch_update(Relaxed, Relaxed, |current| {
934                current
935                    .checked_add(count)
936                    .filter(|next| *next <= self.bulk_packet_capacity)
937            })
938            .is_ok()
939    }
940
941    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
942    fn try_reserve_bulk_packet_prefix(&self, requested: usize) -> usize {
943        if requested == 0 {
944            return 0;
945        }
946
947        let mut current = self.bulk_queued_packets.load(Relaxed);
948        loop {
949            let available = self.bulk_packet_capacity.saturating_sub(current);
950            let granted = requested.min(available);
951            if granted == 0 {
952                return 0;
953            }
954            match self.bulk_queued_packets.compare_exchange_weak(
955                current,
956                current + granted,
957                Relaxed,
958                Relaxed,
959            ) {
960                Ok(_) => return granted,
961                Err(actual) => current = actual,
962            }
963        }
964    }
965
966    fn release_bulk_packets(&self, count: usize) {
967        release_reserved_bulk_packets(&self.bulk_queued_packets, count);
968    }
969
970    #[cfg(test)]
971    pub(crate) fn queued_packets(&self) -> usize {
972        self.queued_packets.load(Relaxed)
973    }
974
975    #[cfg(test)]
976    pub(crate) fn priority_queued_packets(&self) -> usize {
977        self.priority_queued_packets.load(Relaxed)
978    }
979
980    #[cfg(test)]
981    pub(crate) fn bulk_queued_packets(&self) -> usize {
982        self.bulk_queued_packets.load(Relaxed)
983    }
984}
985
986impl PacketRx {
987    pub(crate) fn priority_queued_packets(&self) -> usize {
988        self.priority_queued_packets.load(Relaxed)
989    }
990
991    pub(crate) fn priority_ready_packets(&self) -> usize {
992        self.pending_priority
993            .as_ref()
994            .map_or(0, PendingPackets::len)
995            .saturating_add(self.priority_queued_packets())
996    }
997
998    pub async fn recv(&mut self) -> Option<ReceivedPacket> {
999        loop {
1000            match self.try_recv() {
1001                Ok(packet) => return Some(packet),
1002                Err(TryRecvError::Disconnected) => return None,
1003                Err(TryRecvError::Empty) => {}
1004            }
1005
1006            tokio::select! {
1007                biased;
1008                item = self.priority.recv(), if !self.priority_closed => {
1009                    match item {
1010                        Some(item) => {
1011                            if let Some(packet) = self.packet_from_item(item, PacketLane::Priority) {
1012                                return Some(packet);
1013                            }
1014                        }
1015                        None => self.priority_closed = true,
1016                    }
1017                }
1018                item = self.bulk.recv(), if !self.bulk_closed => {
1019                    match item {
1020                        Some(item) => {
1021                            if let Some(packet) = self.packet_from_item(item, PacketLane::Bulk) {
1022                                return Some(packet);
1023                            }
1024                        }
1025                        None => self.bulk_closed = true,
1026                    }
1027                }
1028            }
1029        }
1030    }
1031
1032    pub fn try_recv(&mut self) -> Result<ReceivedPacket, TryRecvError> {
1033        if let Some(packet) = Self::take_pending(&mut self.pending_priority) {
1034            return Ok(packet);
1035        }
1036
1037        if self.should_probe_priority() {
1038            match self.priority.try_recv() {
1039                Ok(item) => {
1040                    if let Some(packet) = self.packet_from_item(item, PacketLane::Priority) {
1041                        return Ok(packet);
1042                    }
1043                }
1044                Err(TryRecvError::Empty) => {}
1045                Err(TryRecvError::Disconnected) => {
1046                    self.priority_closed = true;
1047                }
1048            }
1049        }
1050
1051        if let Some(packet) = Self::take_pending(&mut self.pending_bulk) {
1052            return Ok(packet);
1053        }
1054
1055        match self.bulk.try_recv() {
1056            Ok(item) => self
1057                .packet_from_item(item, PacketLane::Bulk)
1058                .ok_or(TryRecvError::Empty),
1059            Err(TryRecvError::Empty) => {
1060                if self.priority_closed && self.bulk_closed {
1061                    Err(TryRecvError::Disconnected)
1062                } else {
1063                    Err(TryRecvError::Empty)
1064                }
1065            }
1066            Err(TryRecvError::Disconnected) => {
1067                self.bulk_closed = true;
1068                if self.priority_closed {
1069                    Err(TryRecvError::Disconnected)
1070                } else {
1071                    Err(TryRecvError::Empty)
1072                }
1073            }
1074        }
1075    }
1076
1077    pub(crate) fn drain_ready<F>(&mut self, limit: usize, mut consume: F) -> usize
1078    where
1079        F: FnMut(ReceivedPacket) -> bool,
1080    {
1081        let mut drained = 0usize;
1082        while drained < limit {
1083            if !self.drain_pending_priority(limit, &mut drained, &mut consume) {
1084                break;
1085            }
1086            if drained >= limit {
1087                break;
1088            }
1089
1090            if self.should_probe_priority() {
1091                match self.priority.try_recv() {
1092                    Ok(item) => {
1093                        if !self.drain_item(
1094                            item,
1095                            PacketLane::Priority,
1096                            limit,
1097                            &mut drained,
1098                            &mut consume,
1099                        ) {
1100                            break;
1101                        }
1102                        continue;
1103                    }
1104                    Err(TryRecvError::Empty) => {}
1105                    Err(TryRecvError::Disconnected) => {
1106                        self.priority_closed = true;
1107                    }
1108                }
1109            }
1110            if drained >= limit {
1111                break;
1112            }
1113
1114            if !self.drain_pending_bulk(limit, &mut drained, &mut consume) {
1115                break;
1116            }
1117            if drained >= limit {
1118                break;
1119            }
1120
1121            match self.bulk.try_recv() {
1122                Ok(item) => {
1123                    if !self.drain_item(item, PacketLane::Bulk, limit, &mut drained, &mut consume) {
1124                        break;
1125                    }
1126                }
1127                Err(TryRecvError::Empty) => break,
1128                Err(TryRecvError::Disconnected) => {
1129                    self.bulk_closed = true;
1130                    break;
1131                }
1132            }
1133        }
1134        drained
1135    }
1136
1137    fn packet_from_item(
1138        &mut self,
1139        item: PacketQueueItem,
1140        lane: PacketLane,
1141    ) -> Option<ReceivedPacket> {
1142        item.record_dequeue_wait(lane);
1143        let packet_count = item.packet_count();
1144        if self.track_backlog {
1145            self.queued_packets.fetch_sub(packet_count, Relaxed);
1146        }
1147        if matches!(lane, PacketLane::Priority) {
1148            release_priority_packets(&self.priority_queued_packets, packet_count);
1149        }
1150        if matches!(lane, PacketLane::Bulk) {
1151            release_reserved_bulk_packets(&self.bulk_queued_packets, packet_count);
1152        }
1153        let rx_loop_owned_at = crate::perf_profile::stamp();
1154        match item {
1155            PacketQueueItem::One(mut packet) => {
1156                packet.trace_rx_loop_owned_at = rx_loop_owned_at;
1157                Some(packet)
1158            }
1159            PacketQueueItem::Batch(packets) => {
1160                let mut pending = PendingPackets::new(packets, rx_loop_owned_at);
1161                let packet = pending.next()?;
1162                if pending.len() > 0 {
1163                    match lane {
1164                        PacketLane::Priority => self.pending_priority = Some(pending),
1165                        PacketLane::Bulk => self.pending_bulk = Some(pending),
1166                    }
1167                }
1168                Some(packet)
1169            }
1170        }
1171    }
1172
1173    fn drain_item<F>(
1174        &mut self,
1175        item: PacketQueueItem,
1176        lane: PacketLane,
1177        limit: usize,
1178        drained: &mut usize,
1179        consume: &mut F,
1180    ) -> bool
1181    where
1182        F: FnMut(ReceivedPacket) -> bool,
1183    {
1184        if let Some(packet) = self.packet_from_item(item, lane) {
1185            *drained += 1;
1186            if !consume(packet) {
1187                return false;
1188            }
1189        }
1190
1191        match lane {
1192            PacketLane::Priority => self.drain_pending_priority(limit, drained, consume),
1193            PacketLane::Bulk => self.drain_pending_bulk(limit, drained, consume),
1194        }
1195    }
1196
1197    fn drain_pending_priority<F>(
1198        &mut self,
1199        limit: usize,
1200        drained: &mut usize,
1201        consume: &mut F,
1202    ) -> bool
1203    where
1204        F: FnMut(ReceivedPacket) -> bool,
1205    {
1206        while *drained < limit {
1207            let Some(packet) = Self::take_pending(&mut self.pending_priority) else {
1208                return true;
1209            };
1210            *drained += 1;
1211            if !consume(packet) {
1212                return false;
1213            }
1214        }
1215        true
1216    }
1217
1218    fn drain_pending_bulk<F>(&mut self, limit: usize, drained: &mut usize, consume: &mut F) -> bool
1219    where
1220        F: FnMut(ReceivedPacket) -> bool,
1221    {
1222        while *drained < limit {
1223            if self.should_probe_priority() {
1224                return true;
1225            }
1226            let Some(packet) = Self::take_pending(&mut self.pending_bulk) else {
1227                return true;
1228            };
1229            *drained += 1;
1230            if !consume(packet) {
1231                return false;
1232            }
1233        }
1234        true
1235    }
1236
1237    fn should_probe_priority(&self) -> bool {
1238        !self.priority_closed
1239            && (self.priority_queued_packets.load(Relaxed) > 0 || self.bulk_closed)
1240    }
1241
1242    fn take_pending(pending: &mut Option<PendingPackets>) -> Option<ReceivedPacket> {
1243        let packets = pending.as_mut()?;
1244        let packet = packets.next();
1245        if packets.len() == 0 {
1246            *pending = None;
1247        }
1248        packet
1249    }
1250}
1251
1252#[inline]
1253fn packet_channel_tracks_backlog() -> bool {
1254    cfg!(test) || crate::perf_profile::enabled()
1255}
1256
1257fn release_reserved_bulk_packets(counter: &AtomicUsize, count: usize) {
1258    if count == 0 {
1259        return;
1260    }
1261
1262    let previous = counter.fetch_sub(count, Relaxed);
1263    debug_assert!(
1264        previous >= count,
1265        "transport bulk queued packet accounting underflow"
1266    );
1267}
1268
1269fn release_priority_packets(counter: &AtomicUsize, count: usize) {
1270    if count == 0 {
1271        return;
1272    }
1273
1274    let previous = counter.fetch_sub(count, Relaxed);
1275    debug_assert!(
1276        previous >= count,
1277        "transport priority queued packet accounting underflow"
1278    );
1279}
1280
1281/// Create a packet channel.
1282///
1283/// The capacity applies to bulk packets. Priority traffic is intentionally
1284/// unbounded so small control-shaped packets can still wake the rx loop while a
1285/// bulk receiver is saturated.
1286pub fn packet_channel(buffer: usize) -> (PacketTx, PacketRx) {
1287    let (priority_tx, priority_rx) = tokio::sync::mpsc::unbounded_channel();
1288    let (bulk_tx, bulk_rx) = tokio::sync::mpsc::channel(buffer.max(1));
1289    let priority_queued_packets = Arc::new(AtomicUsize::new(0));
1290    let queued_packets = Arc::new(AtomicUsize::new(0));
1291    let bulk_queued_packets = Arc::new(AtomicUsize::new(0));
1292    let track_backlog = packet_channel_tracks_backlog();
1293    (
1294        PacketTx {
1295            priority: priority_tx,
1296            bulk: bulk_tx,
1297            fast_ingress: None,
1298            batch_pool: PacketBatchPool::new(),
1299            buffer_pool: PacketBufferPool::new(),
1300            priority_queued_packets: Arc::clone(&priority_queued_packets),
1301            queued_packets: Arc::clone(&queued_packets),
1302            bulk_queued_packets: Arc::clone(&bulk_queued_packets),
1303            bulk_packet_capacity: buffer.max(1),
1304            track_backlog,
1305        },
1306        PacketRx {
1307            priority: priority_rx,
1308            bulk: bulk_rx,
1309            priority_queued_packets,
1310            queued_packets,
1311            bulk_queued_packets,
1312            track_backlog,
1313            pending_priority: None,
1314            pending_bulk: None,
1315            priority_closed: false,
1316            bulk_closed: false,
1317        },
1318    )
1319}
1320
1321#[cfg(test)]
1322mod tests;