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