1use 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#[derive(Clone, Debug)]
22pub struct ReceivedPacket {
23 pub transport_id: TransportId,
25 pub remote_addr: TransportAddr,
27 pub data: PacketBuffer,
29 pub timestamp_ms: u64,
31 #[doc(hidden)]
33 pub trace_enqueued_at: Option<crate::perf_profile::TraceStamp>,
34 #[doc(hidden)]
38 pub trace_rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
39}
40
41impl ReceivedPacket {
42 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 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 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#[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 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
274const 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
303const PACKET_BATCH_POOL_LIMIT: usize = 256;
305const PACKET_BATCH_MAX_RETAINED_CAPACITY: usize = 256;
307const PACKET_BUFFER_POOL_LIMIT: usize = 4096;
309const PACKET_BUFFER_MAX_RETAINED_CAPACITY: usize = 16 * 1024;
311
312const TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER: usize = 4096;
319
320#[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 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
343pub 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
1281pub 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;