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::sync::{
6    Arc, Mutex,
7    atomic::{AtomicUsize, Ordering::Relaxed},
8};
9use tokio::sync::mpsc::{
10    Sender, UnboundedReceiver, UnboundedSender,
11    error::{TryRecvError, TrySendError},
12};
13
14pub(crate) trait PacketFastIngressSink: std::fmt::Debug + Send + Sync {
15    fn try_ingest_batch(&self, packets: &mut Vec<ReceivedPacket>) -> usize;
16}
17
18/// A packet received from a transport.
19#[derive(Clone, Debug)]
20pub struct ReceivedPacket {
21    /// Which transport received this packet.
22    pub transport_id: TransportId,
23    /// Remote peer address.
24    pub remote_addr: TransportAddr,
25    /// Packet data.
26    pub data: PacketBuffer,
27    /// Receipt timestamp (Unix milliseconds).
28    pub timestamp_ms: u64,
29    /// Monotonic timestamp for optional pipeline queue-wait profiling.
30    #[doc(hidden)]
31    pub trace_enqueued_at: Option<crate::perf_profile::TraceStamp>,
32    /// Monotonic timestamp captured when `PacketRx` takes ownership of a
33    /// channel item. Distinguishes mpsc/channel residence from rx-loop-owned
34    /// batch-tail residence in perf traces.
35    #[doc(hidden)]
36    pub trace_rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
37}
38
39impl ReceivedPacket {
40    /// Create a received packet with explicit timestamp.
41    pub fn with_timestamp(
42        transport_id: TransportId,
43        remote_addr: TransportAddr,
44        data: PacketBuffer,
45        timestamp_ms: u64,
46    ) -> Self {
47        Self::with_trace_timestamp(
48            transport_id,
49            remote_addr,
50            data,
51            timestamp_ms,
52            crate::perf_profile::stamp(),
53        )
54    }
55
56    /// Create a received packet with explicit wall-clock and queue timestamps.
57    ///
58    /// UDP receive paths can drain several datagrams per kernel batch. Those
59    /// datagrams arrived close together, so sharing one wall-clock sample and
60    /// one queue trace stamp across the batch avoids per-packet clock reads
61    /// while preserving arrival order and queue residence visibility.
62    pub(crate) fn with_trace_timestamp(
63        transport_id: TransportId,
64        remote_addr: TransportAddr,
65        data: PacketBuffer,
66        timestamp_ms: u64,
67        trace_enqueued_at: Option<crate::perf_profile::TraceStamp>,
68    ) -> Self {
69        Self {
70            transport_id,
71            remote_addr,
72            data,
73            timestamp_ms,
74            trace_enqueued_at,
75            trace_rx_loop_owned_at: None,
76        }
77    }
78
79    pub(crate) fn is_transport_priority(&self) -> bool {
80        is_transport_priority_packet(self.data.as_slice())
81    }
82}
83
84/// Byte storage for a received transport packet.
85///
86/// Receive/decrypt/drop paths carry this owner so pressure drops and endpoint
87/// delivery can recycle kernel receive buffers without an extra packet copy.
88#[derive(Debug, Default)]
89pub struct PacketBuffer {
90    data: Vec<u8>,
91    start: usize,
92    pool: Option<PacketBufferPool>,
93}
94
95impl PacketBuffer {
96    #[cfg(any(test, target_os = "linux", target_os = "macos"))]
97    fn pooled(data: Vec<u8>, pool: PacketBufferPool) -> Self {
98        Self {
99            data,
100            start: 0,
101            pool: Some(pool),
102        }
103    }
104
105    pub fn new(data: Vec<u8>) -> Self {
106        Self {
107            data,
108            start: 0,
109            pool: None,
110        }
111    }
112
113    pub fn as_slice(&self) -> &[u8] {
114        &self.data[self.start..]
115    }
116
117    pub fn as_mut_slice(&mut self) -> &mut [u8] {
118        &mut self.data[self.start..]
119    }
120
121    pub fn len(&self) -> usize {
122        self.data.len().saturating_sub(self.start)
123    }
124
125    pub fn is_empty(&self) -> bool {
126        self.len() == 0
127    }
128
129    pub fn into_vec(mut self) -> Vec<u8> {
130        self.pool = None;
131        if self.start > 0 {
132            self.data.drain(..self.start);
133            self.start = 0;
134        }
135        mem::take(&mut self.data)
136    }
137
138    pub(crate) fn trim_front(&mut self, len: usize) -> bool {
139        if len > self.len() {
140            return false;
141        }
142        self.start += len;
143        true
144    }
145
146    pub(crate) fn truncate(&mut self, len: usize) {
147        if len < self.len() {
148            self.data.truncate(self.start + len);
149        }
150    }
151
152    pub(crate) fn extend_from_slice(&mut self, bytes: &[u8]) {
153        self.data.extend_from_slice(bytes);
154    }
155
156    pub(crate) fn try_prepend_slices(&mut self, parts: &[&[u8]], reserve_tail: usize) -> bool {
157        let prefix_len = parts
158            .iter()
159            .fold(0usize, |total, part| total.saturating_add(part.len()));
160        if prefix_len == 0 {
161            return self.data.capacity().saturating_sub(self.data.len()) >= reserve_tail;
162        }
163
164        let len = self.data.len();
165        if self.start >= prefix_len && self.data.capacity().saturating_sub(len) >= reserve_tail {
166            let new_start = self.start - prefix_len;
167            let mut offset = new_start;
168            for part in parts {
169                self.data[offset..offset + part.len()].copy_from_slice(part);
170                offset += part.len();
171            }
172            self.start = new_start;
173            return true;
174        }
175
176        if self.data.capacity().saturating_sub(len) < prefix_len.saturating_add(reserve_tail) {
177            return false;
178        }
179
180        // Move the packet body right inside the existing allocation, then fill
181        // the newly opened header space. This is the Vec equivalent of the
182        // fixed headroom WireGuard-go keeps in its message buffers.
183        unsafe {
184            let ptr = self.data.as_mut_ptr();
185            std::ptr::copy(
186                ptr.add(self.start),
187                ptr.add(self.start + prefix_len),
188                self.len(),
189            );
190            let mut offset = self.start;
191            for part in parts {
192                std::ptr::copy_nonoverlapping(part.as_ptr(), ptr.add(offset), part.len());
193                offset += part.len();
194            }
195            self.data.set_len(len + prefix_len);
196        }
197        true
198    }
199
200    pub(crate) fn replace_visible_prefix(&mut self, remove_len: usize, prefix: &[u8]) -> bool {
201        if remove_len > self.len() {
202            return false;
203        }
204
205        let prefix_len = prefix.len();
206        let tail_len = self.len() - remove_len;
207        if prefix_len >= remove_len {
208            let grow = prefix_len - remove_len;
209            if grow > 0 && self.start >= grow {
210                let new_start = self.start - grow;
211                self.data[new_start..new_start + prefix_len].copy_from_slice(prefix);
212                self.start = new_start;
213                return true;
214            }
215
216            let len = self.data.len();
217            if grow > 0 {
218                self.data.reserve(grow);
219                unsafe {
220                    let ptr = self.data.as_mut_ptr();
221                    std::ptr::copy(
222                        ptr.add(self.start + remove_len),
223                        ptr.add(self.start + prefix_len),
224                        tail_len,
225                    );
226                    self.data.set_len(len + grow);
227                }
228            }
229            self.data[self.start..self.start + prefix_len].copy_from_slice(prefix);
230            return true;
231        }
232
233        let shrink = remove_len - prefix_len;
234        if tail_len > 0 {
235            self.data.copy_within(
236                self.start + remove_len..self.start + remove_len + tail_len,
237                self.start + prefix_len,
238            );
239        }
240        self.data.truncate(self.data.len() - shrink);
241        self.data[self.start..self.start + prefix_len].copy_from_slice(prefix);
242        true
243    }
244
245    pub(crate) fn recycle_batch(packets: &mut [Self]) {
246        let Some(pool) = packets.first().and_then(|packet| packet.pool.clone()) else {
247            return;
248        };
249        if packets.iter().all(|packet| {
250            packet
251                .pool
252                .as_ref()
253                .is_some_and(|packet_pool| pool.shares_storage(packet_pool))
254        }) {
255            for packet in packets.iter_mut() {
256                packet.pool = None;
257            }
258            pool.put_batch(packets);
259        }
260    }
261}
262
263impl Clone for PacketBuffer {
264    fn clone(&self) -> Self {
265        Self {
266            data: self.as_slice().to_vec(),
267            start: 0,
268            pool: None,
269        }
270    }
271}
272
273impl AsRef<[u8]> for PacketBuffer {
274    fn as_ref(&self) -> &[u8] {
275        self.as_slice()
276    }
277}
278
279impl Drop for PacketBuffer {
280    fn drop(&mut self) {
281        if let Some(pool) = self.pool.take() {
282            pool.put(mem::take(&mut self.data));
283        }
284    }
285}
286
287impl PartialEq for PacketBuffer {
288    fn eq(&self, other: &Self) -> bool {
289        self.as_slice() == other.as_slice()
290    }
291}
292
293impl Eq for PacketBuffer {}
294
295/// FMP packet shape that is visible before dataplane authenticates established data.
296///
297/// Bulk app data is opaque phase-0 data here, so the transport channel only
298/// promotes exact control-sized frames that can be identified from public wire
299/// length: handshakes, link heartbeats, and fixed-size link MMP reports.
300const FMP_VERSION: u8 = crate::node::wire::FMP_VERSION;
301const FMP_PHASE_ESTABLISHED: u8 = crate::node::wire::PHASE_ESTABLISHED;
302const FMP_PHASE_MSG1: u8 = crate::node::wire::PHASE_MSG1;
303const FMP_PHASE_MSG2: u8 = crate::node::wire::PHASE_MSG2;
304const FMP_COMMON_PREFIX_SIZE: usize = crate::node::wire::COMMON_PREFIX_SIZE;
305const FMP_ESTABLISHED_HEADER_SIZE: usize = crate::node::wire::ESTABLISHED_HEADER_SIZE;
306const FMP_MSG1_WIRE_SIZE: usize = crate::node::wire::MSG1_WIRE_SIZE;
307const FMP_MSG2_WIRE_SIZE: usize = crate::node::wire::MSG2_WIRE_SIZE;
308const AEAD_TAG_SIZE: usize = crate::noise::TAG_SIZE;
309const FMP_HEARTBEAT_PLAINTEXT_SIZE: usize = 4 + 1;
310const FMP_MMP_SENDER_REPORT_PLAINTEXT_SIZE: usize = crate::mmp::SENDER_REPORT_WIRE_SIZE;
311const FMP_MMP_RECEIVER_REPORT_PLAINTEXT_SIZE: usize = crate::mmp::RECEIVER_REPORT_WIRE_SIZE;
312
313fn is_transport_priority_packet(data: &[u8]) -> bool {
314    if data.len() < FMP_COMMON_PREFIX_SIZE {
315        return false;
316    }
317
318    let version = data[0] >> 4;
319    let phase = data[0] & 0x0F;
320    if version != FMP_VERSION {
321        return false;
322    }
323
324    match phase {
325        FMP_PHASE_MSG1 => data.len() == FMP_MSG1_WIRE_SIZE,
326        FMP_PHASE_MSG2 => data.len() == FMP_MSG2_WIRE_SIZE,
327        FMP_PHASE_ESTABLISHED => is_fmp_established_priority_packet(data),
328        _ => false,
329    }
330}
331
332fn is_fmp_established_priority_packet(data: &[u8]) -> bool {
333    if data.len() < FMP_ESTABLISHED_HEADER_SIZE.saturating_add(AEAD_TAG_SIZE) {
334        return false;
335    }
336
337    let payload_len = usize::from(u16::from_le_bytes([data[2], data[3]]));
338    let expected_len = FMP_ESTABLISHED_HEADER_SIZE
339        .saturating_add(payload_len)
340        .saturating_add(AEAD_TAG_SIZE);
341    if data.len() != expected_len {
342        return false;
343    }
344
345    matches!(
346        payload_len,
347        FMP_HEARTBEAT_PLAINTEXT_SIZE
348            | FMP_MMP_SENDER_REPORT_PLAINTEXT_SIZE
349            | FMP_MMP_RECEIVER_REPORT_PLAINTEXT_SIZE
350    )
351}
352
353/// Number of receive-batch Vec containers retained for reuse.
354const PACKET_BATCH_POOL_LIMIT: usize = 256;
355/// Avoid pinning unusually large test/control batches in the hot-path pool.
356const PACKET_BATCH_MAX_RETAINED_CAPACITY: usize = 256;
357/// Number of packet byte buffers retained after pressure drops.
358const PACKET_BUFFER_POOL_LIMIT: usize = 4096;
359/// Avoid pinning oversized receive buffers in the hot-path pool.
360const PACKET_BUFFER_MAX_RETAINED_CAPACITY: usize = 16 * 1024;
361
362/// Packet count at which the transport receive channel is visibly backlogged.
363///
364/// This tracks packets still owned by the priority/bulk mpsc channels. Once a
365/// batched item is dequeued into `PacketRx`'s pending iterator, it no longer
366/// contributes to this counter; those packets are already inside the rx-loop
367/// owner's drain budget rather than waiting behind the transport channel.
368const TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER: usize = 16_384;
369
370/// Channel sender for received packets.
371///
372/// The priority lane stays unbounded because control-shaped datagrams must keep
373/// making progress even when bulk is saturated. The bulk lane is bounded by the
374/// configured packet-channel capacity in packets, not receive-batch items, and
375/// uses nonblocking `try_send`: overload sheds bulk explicitly instead of
376/// hiding unbounded latency behind the rx loop.
377#[derive(Clone, Debug)]
378pub struct PacketTx {
379    priority: UnboundedSender<PacketQueueItem>,
380    bulk: Sender<PacketQueueItem>,
381    fast_ingress: Option<Arc<dyn PacketFastIngressSink>>,
382    batch_pool: PacketBatchPool,
383    #[cfg(any(test, target_os = "linux", target_os = "macos"))]
384    buffer_pool: PacketBufferPool,
385    /// Packet-count ready hint for priority lane probes. Bulk batch tails check
386    /// this instead of touching an empty priority mpsc once per data packet.
387    priority_queued_packets: Arc<AtomicUsize>,
388    queued_packets: Arc<AtomicUsize>,
389    bulk_queued_packets: Arc<AtomicUsize>,
390    bulk_packet_capacity: usize,
391    track_backlog: bool,
392}
393
394/// Channel receiver for received packets.
395pub struct PacketRx {
396    priority: UnboundedReceiver<PacketQueueItem>,
397    bulk: tokio::sync::mpsc::Receiver<PacketQueueItem>,
398    priority_queued_packets: Arc<AtomicUsize>,
399    queued_packets: Arc<AtomicUsize>,
400    bulk_queued_packets: Arc<AtomicUsize>,
401    track_backlog: bool,
402    pending_priority: Option<PendingPackets>,
403    pending_bulk: Option<PendingPackets>,
404    priority_closed: bool,
405    bulk_closed: bool,
406}
407
408#[derive(Clone, Debug)]
409struct PacketBatchPool {
410    inner: Arc<Mutex<Vec<Vec<ReceivedPacket>>>>,
411}
412
413#[derive(Clone, Debug)]
414struct PacketBufferPool {
415    inner: Arc<Mutex<Vec<Vec<u8>>>>,
416    available: Arc<AtomicUsize>,
417}
418
419#[derive(Debug)]
420pub(crate) struct PacketBatch {
421    packets: Vec<ReceivedPacket>,
422    pool: Option<PacketBatchPool>,
423}
424
425#[derive(Debug)]
426enum PacketQueueItem {
427    One(ReceivedPacket),
428    Batch(PacketBatch),
429}
430
431#[derive(Clone, Copy)]
432enum PacketLane {
433    Priority,
434    Bulk,
435}
436
437#[derive(Clone, Copy)]
438enum PacketQueueTx {
439    Priority,
440    Bulk,
441}
442
443enum PacketSendFailure {
444    Closed(PacketQueueItem),
445    DroppedBulk(usize),
446}
447
448struct PendingPackets {
449    batch: PacketBatch,
450    rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
451}
452
453#[derive(Debug, PartialEq, Eq)]
454struct PacketQueueDequeueCounts {
455    total: usize,
456    priority: usize,
457    bulk: usize,
458}
459
460impl PacketQueueTx {
461    fn try_send(self, owner: &PacketTx, item: PacketQueueItem) -> Result<(), PacketSendFailure> {
462        match self {
463            PacketQueueTx::Priority => owner
464                .priority
465                .send(item)
466                .map_err(|error| PacketSendFailure::Closed(error.0)),
467            PacketQueueTx::Bulk => {
468                let packet_count = item.packet_count();
469                match owner.bulk.try_send(item) {
470                    Ok(()) => Ok(()),
471                    Err(TrySendError::Full(_item)) => {
472                        Err(PacketSendFailure::DroppedBulk(packet_count))
473                    }
474                    Err(TrySendError::Closed(item)) => Err(PacketSendFailure::Closed(item)),
475                }
476            }
477        }
478    }
479}
480
481impl PacketQueueItem {
482    fn packet_count(&self) -> usize {
483        match self {
484            PacketQueueItem::One(_) => 1,
485            PacketQueueItem::Batch(packets) => packets.packets.len(),
486        }
487    }
488
489    fn dequeue_counts(&self, lane: PacketLane) -> PacketQueueDequeueCounts {
490        let total = self.packet_count();
491        match lane {
492            PacketLane::Priority => PacketQueueDequeueCounts {
493                total,
494                priority: total,
495                bulk: 0,
496            },
497            PacketLane::Bulk => PacketQueueDequeueCounts {
498                total,
499                priority: 0,
500                bulk: total,
501            },
502        }
503    }
504
505    fn queued_at(&self) -> Option<crate::perf_profile::TraceStamp> {
506        match self {
507            PacketQueueItem::One(packet) => packet.trace_enqueued_at,
508            PacketQueueItem::Batch(packets) => packets
509                .packets
510                .first()
511                .and_then(|packet| packet.trace_enqueued_at),
512        }
513    }
514
515    fn record_dequeue_wait(&self, lane: PacketLane) {
516        let queued_at = self.queued_at();
517        if queued_at.is_none() {
518            return;
519        }
520        let counts = self.dequeue_counts(lane);
521        crate::perf_profile::record_since_split_count(
522            crate::perf_profile::Stage::TransportChannelWait,
523            crate::perf_profile::Stage::TransportPriorityChannelWait,
524            crate::perf_profile::Stage::TransportBulkChannelWait,
525            queued_at,
526            counts.total as u64,
527            counts.priority as u64,
528            counts.bulk as u64,
529        );
530    }
531}
532
533impl PacketBatchPool {
534    fn new() -> Self {
535        Self {
536            inner: Arc::new(Mutex::new(Vec::new())),
537        }
538    }
539
540    fn take(&self, capacity: usize) -> PacketBatch {
541        let packets = {
542            let mut guard = self.inner.lock().unwrap_or_else(|error| error.into_inner());
543            guard.pop()
544        };
545        if let Some(mut packets) = packets {
546            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBatchPoolReuse);
547            packets.clear();
548            if packets.capacity() >= capacity {
549                return PacketBatch::pooled(packets, self.clone());
550            }
551            packets.reserve(capacity.saturating_sub(packets.capacity()));
552            return PacketBatch::pooled(packets, self.clone());
553        }
554        crate::perf_profile::record_event(crate::perf_profile::Event::PacketBatchPoolFresh);
555        PacketBatch::pooled(Vec::with_capacity(capacity), self.clone())
556    }
557
558    fn put(&self, mut packets: Vec<ReceivedPacket>) {
559        packets.clear();
560        if packets.capacity() > PACKET_BATCH_MAX_RETAINED_CAPACITY {
561            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBatchPoolDiscard);
562            return;
563        }
564        let mut guard = self.inner.lock().unwrap_or_else(|error| error.into_inner());
565        if guard.len() < PACKET_BATCH_POOL_LIMIT {
566            guard.push(packets);
567            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBatchPoolReturn);
568        } else {
569            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBatchPoolDiscard);
570        }
571    }
572}
573
574impl PacketBufferPool {
575    #[cfg(any(test, target_os = "linux", target_os = "macos"))]
576    fn new() -> Self {
577        Self {
578            inner: Arc::new(Mutex::new(Vec::new())),
579            available: Arc::new(AtomicUsize::new(0)),
580        }
581    }
582
583    #[cfg(any(test, target_os = "linux", target_os = "macos"))]
584    fn take(&self, capacity: usize) -> Vec<u8> {
585        if self.available.load(Relaxed) > 0 {
586            let buffer = {
587                let mut guard = self.inner.lock().unwrap_or_else(|error| error.into_inner());
588                guard.pop()
589            };
590            if let Some(mut buffer) = buffer {
591                self.available.fetch_sub(1, Relaxed);
592                crate::perf_profile::record_event(
593                    crate::perf_profile::Event::PacketBufferPoolReuse,
594                );
595                prepare_recv_buffer(&mut buffer, capacity);
596                return buffer;
597            }
598        }
599
600        crate::perf_profile::record_event(crate::perf_profile::Event::PacketBufferPoolFresh);
601        fresh_recv_buffer(capacity)
602    }
603
604    fn put(&self, mut buffer: Vec<u8>) {
605        buffer.clear();
606        if buffer.capacity() > PACKET_BUFFER_MAX_RETAINED_CAPACITY {
607            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBufferPoolDiscard);
608            return;
609        }
610
611        let mut guard = self.inner.lock().unwrap_or_else(|error| error.into_inner());
612        if guard.len() < PACKET_BUFFER_POOL_LIMIT {
613            guard.push(buffer);
614            self.available.fetch_add(1, Relaxed);
615            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBufferPoolReturn);
616        } else {
617            crate::perf_profile::record_event(crate::perf_profile::Event::PacketBufferPoolDiscard);
618        }
619    }
620
621    fn shares_storage(&self, other: &Self) -> bool {
622        Arc::ptr_eq(&self.inner, &other.inner)
623    }
624
625    fn put_batch(&self, packets: &mut [PacketBuffer]) {
626        let mut returned = 0usize;
627        let mut discarded = 0usize;
628        let mut guard = self.inner.lock().unwrap_or_else(|error| error.into_inner());
629        let available_slots = PACKET_BUFFER_POOL_LIMIT.saturating_sub(guard.len());
630        for packet in packets {
631            packet.start = 0;
632            let mut buffer = mem::take(&mut packet.data);
633            buffer.clear();
634            if buffer.capacity() <= PACKET_BUFFER_MAX_RETAINED_CAPACITY
635                && returned < available_slots
636            {
637                guard.push(buffer);
638                returned += 1;
639            } else {
640                discarded += 1;
641            }
642        }
643        if returned > 0 {
644            self.available.fetch_add(returned, Relaxed);
645        }
646        drop(guard);
647        if returned > 0 {
648            crate::perf_profile::record_event_count(
649                crate::perf_profile::Event::PacketBufferPoolReturn,
650                returned as u64,
651            );
652        }
653        if discarded > 0 {
654            crate::perf_profile::record_event_count(
655                crate::perf_profile::Event::PacketBufferPoolDiscard,
656                discarded as u64,
657            );
658        }
659    }
660}
661
662#[cfg(target_os = "macos")]
663fn fresh_recv_buffer(size: usize) -> Vec<u8> {
664    vec![0u8; size]
665}
666
667#[cfg(all(any(test, target_os = "linux"), not(target_os = "macos")))]
668fn fresh_recv_buffer(size: usize) -> Vec<u8> {
669    Vec::with_capacity(size)
670}
671
672#[cfg(target_os = "macos")]
673fn prepare_recv_buffer(buffer: &mut Vec<u8>, size: usize) {
674    buffer.resize(size, 0);
675}
676
677#[cfg(all(any(test, target_os = "linux"), not(target_os = "macos")))]
678fn prepare_recv_buffer(buffer: &mut Vec<u8>, size: usize) {
679    buffer.clear();
680    if buffer.capacity() < size {
681        buffer.reserve(size.saturating_sub(buffer.capacity()));
682    }
683}
684
685impl PacketBatch {
686    fn pooled(packets: Vec<ReceivedPacket>, pool: PacketBatchPool) -> Self {
687        Self {
688            packets,
689            pool: Some(pool),
690        }
691    }
692
693    pub(crate) fn push(&mut self, packet: ReceivedPacket) {
694        self.packets.push(packet);
695    }
696
697    pub(crate) fn is_empty(&self) -> bool {
698        self.packets.is_empty()
699    }
700
701    pub(crate) fn as_slice(&self) -> &[ReceivedPacket] {
702        &self.packets
703    }
704}
705
706impl Drop for PacketBatch {
707    fn drop(&mut self) {
708        let Some(pool) = self.pool.take() else {
709            return;
710        };
711        pool.put(mem::take(&mut self.packets));
712    }
713}
714
715impl PendingPackets {
716    fn new(
717        mut batch: PacketBatch,
718        rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
719    ) -> Self {
720        batch.packets.reverse();
721        Self {
722            batch,
723            rx_loop_owned_at,
724        }
725    }
726
727    fn next(&mut self) -> Option<ReceivedPacket> {
728        let mut packet = self.batch.packets.pop()?;
729        if let Some(rx_loop_owned_at) = self.rx_loop_owned_at {
730            packet.trace_rx_loop_owned_at = Some(rx_loop_owned_at);
731        }
732        Some(packet)
733    }
734}
735
736impl PacketTx {
737    pub(crate) fn set_fast_ingress_sink(&mut self, sink: Arc<dyn PacketFastIngressSink>) {
738        self.fast_ingress = Some(sink);
739    }
740
741    pub(crate) fn try_fast_ingress_packet_batch(&self, batch: &mut PacketBatch) -> usize {
742        let Some(sink) = &self.fast_ingress else {
743            return 0;
744        };
745        sink.try_ingest_batch(&mut batch.packets)
746    }
747
748    pub(crate) fn packet_batch(&self, capacity: usize) -> PacketBatch {
749        self.batch_pool.take(capacity)
750    }
751
752    #[cfg(any(test, target_os = "linux", target_os = "macos"))]
753    pub(crate) fn recv_buffer(&self, capacity: usize) -> Vec<u8> {
754        self.buffer_pool.take(capacity)
755    }
756
757    #[cfg(any(test, target_os = "linux", target_os = "macos"))]
758    pub(crate) fn packet_buffer(&self, data: Vec<u8>) -> PacketBuffer {
759        PacketBuffer::pooled(data, self.buffer_pool.clone())
760    }
761
762    pub fn send(
763        &self,
764        packet: ReceivedPacket,
765    ) -> Result<(), tokio::sync::mpsc::error::SendError<ReceivedPacket>> {
766        let tx = if packet.is_transport_priority() {
767            PacketQueueTx::Priority
768        } else {
769            PacketQueueTx::Bulk
770        };
771        self.send_item(tx, PacketQueueItem::One(packet))
772            .map_err(|item| match item {
773                PacketQueueItem::One(packet) => tokio::sync::mpsc::error::SendError(packet),
774                PacketQueueItem::Batch(_) => {
775                    unreachable!("single packet send cannot fail with a batch item")
776                }
777            })
778    }
779
780    pub(crate) fn send_packet_batch(&self, mut batch: PacketBatch) -> Result<(), ()> {
781        if batch.is_empty() {
782            return Ok(());
783        }
784
785        let packet_count = batch.packets.len();
786        let priority_count = batch
787            .packets
788            .iter()
789            .filter(|packet| packet.is_transport_priority())
790            .count();
791        if priority_count == 0 || priority_count == packet_count {
792            let tx = if priority_count == 0 {
793                PacketQueueTx::Bulk
794            } else {
795                PacketQueueTx::Priority
796            };
797            return self.send_packet_items(tx, batch);
798        }
799
800        let mut priority_packets = self.packet_batch(priority_count);
801        let mut bulk_packets = self.packet_batch(packet_count - priority_count);
802        for packet in batch.packets.drain(..) {
803            if packet.is_transport_priority() {
804                priority_packets.push(packet);
805            } else {
806                bulk_packets.push(packet);
807            }
808        }
809
810        self.send_packet_items(PacketQueueTx::Priority, priority_packets)?;
811        self.send_packet_items(PacketQueueTx::Bulk, bulk_packets)?;
812        Ok(())
813    }
814
815    fn send_packet_items(&self, tx: PacketQueueTx, packets: PacketBatch) -> Result<(), ()> {
816        if matches!(tx, PacketQueueTx::Bulk) {
817            return self.send_bulk_packet_items(packets);
818        }
819
820        let item = match packets.packets.len() {
821            0 => return Ok(()),
822            _ => PacketQueueItem::Batch(packets),
823        };
824        self.send_item(tx, item).map_err(|_| ())
825    }
826
827    fn send_bulk_packet_items(&self, mut packets: PacketBatch) -> Result<(), ()> {
828        let packet_count = packets.packets.len();
829        if packet_count == 0 {
830            return Ok(());
831        }
832
833        let granted = self.try_reserve_bulk_packet_prefix(packet_count);
834        if granted == 0 {
835            crate::perf_profile::record_event_count(
836                crate::perf_profile::Event::TransportBulkDropped,
837                packet_count as u64,
838            );
839            return Ok(());
840        }
841
842        if granted < packet_count {
843            let dropped = packet_count - granted;
844            let _dropped_tail = packets.packets.split_off(granted);
845            crate::perf_profile::record_event_count(
846                crate::perf_profile::Event::TransportBulkDropped,
847                dropped as u64,
848            );
849        }
850
851        let item = match packets.packets.len() {
852            0 => return Ok(()),
853            _ => PacketQueueItem::Batch(packets),
854        };
855        self.send_reserved_item(PacketQueueTx::Bulk, item, Some(granted))
856            .map_err(|_| ())
857    }
858
859    fn send_item(&self, tx: PacketQueueTx, item: PacketQueueItem) -> Result<(), PacketQueueItem> {
860        let packet_count = item.packet_count();
861        let bulk_reserved = if matches!(tx, PacketQueueTx::Bulk) && packet_count > 0 {
862            if !self.try_reserve_bulk_packets(packet_count) {
863                crate::perf_profile::record_event_count(
864                    crate::perf_profile::Event::TransportBulkDropped,
865                    packet_count as u64,
866                );
867                return Ok(());
868            }
869            Some(packet_count)
870        } else {
871            None
872        };
873        self.send_reserved_item(tx, item, bulk_reserved)
874    }
875
876    fn send_reserved_item(
877        &self,
878        tx: PacketQueueTx,
879        item: PacketQueueItem,
880        bulk_reserved: Option<usize>,
881    ) -> Result<(), PacketQueueItem> {
882        let packet_count = item.packet_count();
883        debug_assert_eq!(
884            bulk_reserved,
885            matches!(tx, PacketQueueTx::Bulk)
886                .then_some(packet_count)
887                .filter(|count| *count > 0)
888        );
889        let priority_reserved = matches!(tx, PacketQueueTx::Priority)
890            .then_some(packet_count)
891            .filter(|count| *count > 0);
892        if let Some(count) = priority_reserved {
893            self.priority_queued_packets.fetch_add(count, Relaxed);
894        }
895
896        let tracked_count = if self.track_backlog {
897            Some(packet_count)
898        } else {
899            None
900        };
901        let previous = tracked_count.map(|count| self.queued_packets.fetch_add(count, Relaxed));
902        match tx.try_send(self, item) {
903            Ok(()) => {
904                if let (Some(count), Some(previous)) = (tracked_count, previous) {
905                    let queued = previous.saturating_add(count);
906                    if previous < TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER
907                        && queued >= TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER
908                    {
909                        crate::perf_profile::record_event(
910                            crate::perf_profile::Event::TransportChannelBacklogHigh,
911                        );
912                    }
913                }
914                Ok(())
915            }
916            Err(PacketSendFailure::Closed(item)) => {
917                if let Some(count) = tracked_count {
918                    self.queued_packets.fetch_sub(count, Relaxed);
919                }
920                if let Some(count) = priority_reserved {
921                    release_priority_packets(&self.priority_queued_packets, count);
922                }
923                if let Some(count) = bulk_reserved {
924                    self.release_bulk_packets(count);
925                }
926                Err(item)
927            }
928            Err(PacketSendFailure::DroppedBulk(dropped_count)) => {
929                if let Some(count) = tracked_count {
930                    self.queued_packets.fetch_sub(count, Relaxed);
931                }
932                if let Some(count) = priority_reserved {
933                    release_priority_packets(&self.priority_queued_packets, count);
934                }
935                if let Some(count) = bulk_reserved {
936                    self.release_bulk_packets(count);
937                }
938                crate::perf_profile::record_event_count(
939                    crate::perf_profile::Event::TransportBulkDropped,
940                    dropped_count as u64,
941                );
942                Ok(())
943            }
944        }
945    }
946
947    fn try_reserve_bulk_packets(&self, count: usize) -> bool {
948        self.bulk_queued_packets
949            .fetch_update(Relaxed, Relaxed, |current| {
950                current
951                    .checked_add(count)
952                    .filter(|next| *next <= self.bulk_packet_capacity)
953            })
954            .is_ok()
955    }
956
957    fn try_reserve_bulk_packet_prefix(&self, requested: usize) -> usize {
958        if requested == 0 {
959            return 0;
960        }
961
962        let mut current = self.bulk_queued_packets.load(Relaxed);
963        loop {
964            let available = self.bulk_packet_capacity.saturating_sub(current);
965            let granted = requested.min(available);
966            if granted == 0 {
967                return 0;
968            }
969            match self.bulk_queued_packets.compare_exchange_weak(
970                current,
971                current + granted,
972                Relaxed,
973                Relaxed,
974            ) {
975                Ok(_) => return granted,
976                Err(actual) => current = actual,
977            }
978        }
979    }
980
981    fn release_bulk_packets(&self, count: usize) {
982        release_reserved_bulk_packets(&self.bulk_queued_packets, count);
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, |packets| packets.batch.packets.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.batch.packets.is_empty() {
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.batch.packets.is_empty() {
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            #[cfg(any(test, target_os = "linux", target_os = "macos"))]
1300            buffer_pool: PacketBufferPool::new(),
1301            priority_queued_packets: Arc::clone(&priority_queued_packets),
1302            queued_packets: Arc::clone(&queued_packets),
1303            bulk_queued_packets: Arc::clone(&bulk_queued_packets),
1304            bulk_packet_capacity: buffer.max(1),
1305            track_backlog,
1306        },
1307        PacketRx {
1308            priority: priority_rx,
1309            bulk: bulk_rx,
1310            priority_queued_packets,
1311            queued_packets,
1312            bulk_queued_packets,
1313            track_backlog,
1314            pending_priority: None,
1315            pending_bulk: None,
1316            priority_closed: false,
1317            bulk_closed: false,
1318        },
1319    )
1320}
1321
1322#[cfg(test)]
1323mod tests;
fips_core/transport/packet_channel.rs

fips_core/transport/
packet_channel.rs
