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
16#[derive(Clone, Debug)]
18pub struct ReceivedPacket {
19 pub transport_id: TransportId,
21 pub remote_addr: TransportAddr,
23 pub data: PacketBuffer,
25 pub timestamp_ms: u64,
27 #[doc(hidden)]
29 pub trace_enqueued_at: Option<crate::perf_profile::TraceStamp>,
30 #[doc(hidden)]
34 pub trace_rx_loop_owned_at: Option<crate::perf_profile::TraceStamp>,
35}
36
37impl ReceivedPacket {
38 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 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 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#[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
216const PRIORITY_PACKET_MAX_LEN: usize = 512;
220const PACKET_BATCH_POOL_LIMIT: usize = 256;
222const PACKET_BATCH_MAX_RETAINED_CAPACITY: usize = 256;
224const PACKET_BUFFER_POOL_LIMIT: usize = 4096;
226const PACKET_BUFFER_MAX_RETAINED_CAPACITY: usize = 16 * 1024;
228
229const TRANSPORT_CHANNEL_BACKLOG_HIGH_WATER: usize = 4096;
236
237#[derive(Clone, Debug)]
245pub struct PacketTx {
246 priority: UnboundedSender<PacketQueueItem>,
247 bulk: Sender<PacketQueueItem>,
248 batch_pool: PacketBatchPool,
249 buffer_pool: PacketBufferPool,
250 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
259pub 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
1060pub 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;