From 7b506e25a7f14b989e07fd54831bc19967dc5ab1 Mon Sep 17 00:00:00 2001 From: KKRainbow <443152178@qq.com> Date: Mon, 27 Jul 2026 21:32:57 +0800 Subject: [PATCH] perf(data-plane): reduce per-packet synchronization overhead (#2453) * perf(stats): avoid per-update clock reads Perf profiles show quanta::get_now consuming 2.9-4.3% of data-plane CPU because every counter update refreshes a high-resolution timestamp. Track metric activity with the existing 60-second cleanup cadence instead. Relaxed 32-bit epochs preserve the three-minute retention window, support 32-bit targets, and remove repeated clock reads from packet processing. * perf(data-plane): reduce per-packet synchronization Perf profiles showed per-packet config Arc cloning, bounded-channel permit futures, duplicate peer lookups, and default connection UUID lookups consuming CPU in both TCP and UDP data paths. Borrow stable config snapshots, use nonblocking channel fast paths with the existing backpressure fallback, reuse direct peer lookups, and cache the selected connection while preserving close and reselection behavior. Add focused tests for channel backpressure and cached connection invalidation. * fix(peer): serialize default connection cache updates The profile-guided default connection cache could republish a connection after the close task removed it, leaving a stale cache while another connection remained live. Serialize only cache-miss selection/publication and connection removal. The per-packet cache-hit path remains lock-free, while close and selection can no longer race to resurrect a removed connection. --- easytier-core/src/config/runtime.rs | 5 + easytier-core/src/foundation/stats.rs | 113 ++++++++++++---------- easytier-core/src/peers/conn/peer.rs | 75 ++++++++++---- easytier-core/src/peers/conn/peer_conn.rs | 3 +- easytier-core/src/peers/conn/peer_map.rs | 4 + easytier-core/src/peers/context.rs | 20 ++-- easytier-core/src/peers/mod.rs | 12 +++ easytier-core/src/peers/peer_manager.rs | 4 +- easytier-core/src/peers/tests.rs | 93 +++++++++++++++++- easytier-core/src/tunnel/mpsc.rs | 38 +++++++- 10 files changed, 284 insertions(+), 83 deletions(-) diff --git a/easytier-core/src/config/runtime.rs b/easytier-core/src/config/runtime.rs index d90c9026..c3a64a0f 100644 --- a/easytier-core/src/config/runtime.rs +++ b/easytier-core/src/config/runtime.rs @@ -81,6 +81,11 @@ impl CoreRuntimeConfigStore { self.inner.snapshot.load_full() } + pub(crate) fn with_snapshot(&self, read: impl FnOnce(&CoreInstanceRuntimeConfig) -> T) -> T { + let snapshot = self.inner.snapshot.load(); + read(&snapshot) + } + pub fn replace(&self, config: CoreInstanceRuntimeConfig) { let _update = self.inner.update.lock(); self.inner.snapshot.store(Arc::new(config)); diff --git a/easytier-core/src/foundation/stats.rs b/easytier-core/src/foundation/stats.rs index 5f7354a4..fcbd25cb 100644 --- a/easytier-core/src/foundation/stats.rs +++ b/easytier-core/src/foundation/stats.rs @@ -1,14 +1,19 @@ use dashmap::DashMap; -use quanta::Instant; use serde::{Deserialize, Serialize}; use std::cell::UnsafeCell; use std::fmt; -use std::sync::{Arc, Mutex}; +use std::sync::{ + Arc, Mutex, + atomic::{AtomicU32, Ordering}, +}; use std::time::Duration; use tokio_util::task::AbortOnDropHandle; use crate::foundation::time::interval; +const METRIC_CLEANUP_INTERVAL: Duration = Duration::from_secs(60); +const METRIC_RETENTION_EPOCHS: u32 = 3; + #[derive(Clone, Debug, PartialEq, Eq)] pub struct RpcMetricLabels { pub network_name: String, @@ -535,43 +540,44 @@ impl UnsafeCounter { unsafe impl Send for UnsafeCounter {} unsafe impl Sync for UnsafeCounter {} -/// MetricData contains both the counter and last update timestamp -/// Uses UnsafeCell for lock-free access +/// MetricData contains both the counter and its last active cleanup epoch. #[derive(Debug)] struct MetricData { counter: UnsafeCounter, - last_updated: UnsafeCell, + activity_epoch: Arc, + last_updated_epoch: AtomicU32, } impl MetricData { - fn new() -> Self { + fn new(activity_epoch: Arc) -> Self { + let last_updated_epoch = activity_epoch.load(Ordering::Relaxed); Self { counter: UnsafeCounter::new(), - last_updated: UnsafeCell::new(Instant::now()), + activity_epoch, + last_updated_epoch: AtomicU32::new(last_updated_epoch), } } - /// Update the last_updated timestamp - /// # Safety - /// This method is unsafe because it uses UnsafeCell. The caller must ensure - /// that no other thread is accessing this timestamp simultaneously. - unsafe fn touch(&self) { - let ptr = self.last_updated.get(); - unsafe { - *ptr = Instant::now(); - } + fn touch(&self) { + let current_epoch = self.activity_epoch.load(Ordering::Relaxed); + self.last_updated_epoch + .store(current_epoch, Ordering::Relaxed); } - /// Get the last updated timestamp - /// # Safety - /// This method is unsafe because it uses UnsafeCell. The caller must ensure - /// that no other thread is modifying this timestamp simultaneously. - unsafe fn get_last_updated(&self) -> Instant { - let ptr = self.last_updated.get(); - unsafe { *ptr } + fn last_updated_epoch(&self) -> u32 { + self.last_updated_epoch.load(Ordering::Relaxed) } } +fn cleanup_metrics(counters: &DashMap>, current_epoch: u32) { + counters.retain(|_, metric_data| { + Arc::strong_count(metric_data) > 1 + || current_epoch.saturating_sub(metric_data.last_updated_epoch()) + <= METRIC_RETENTION_EPOCHS + }); + counters.shrink_to_fit(); +} + // MetricData is Send + Sync because the safety is guaranteed by the caller unsafe impl Send for MetricData {} unsafe impl Sync for MetricData {} @@ -620,16 +626,16 @@ impl CounterHandle { pub fn add(&self, delta: u64) { unsafe { self.metric_data.counter.add(delta); - self.metric_data.touch(); } + self.metric_data.touch(); } /// Increment the counter by 1 pub fn inc(&self) { unsafe { self.metric_data.counter.inc(); - self.metric_data.touch(); } + self.metric_data.touch(); } /// Get the current value of the counter @@ -641,16 +647,16 @@ impl CounterHandle { pub fn reset(&self) { unsafe { self.metric_data.counter.reset(); - self.metric_data.touch(); } + self.metric_data.touch(); } /// Set the counter to a specific value pub fn set(&self, value: u64) { unsafe { self.metric_data.counter.set(value); - self.metric_data.touch(); } + self.metric_data.touch(); } } @@ -671,6 +677,7 @@ impl MetricSnapshot { /// StatsManager manages global statistics with high performance counters pub struct StatsManager { counters: Arc>>, + activity_epoch: Arc, cleanup_task: Mutex>>, } @@ -679,6 +686,7 @@ impl StatsManager { pub fn new() -> Self { let manager = Self { counters: Arc::new(DashMap::new()), + activity_epoch: Arc::new(AtomicU32::new(0)), cleanup_task: Mutex::new(None), }; manager.start_cleanup_task(); @@ -698,24 +706,19 @@ impl StatsManager { return; }; let counters = Arc::downgrade(&self.counters); + let activity_epoch = Arc::clone(&self.activity_epoch); *cleanup_task = Some(AbortOnDropHandle::new(runtime.spawn(async move { - let mut interval = interval(Duration::from_secs(60)); // Check every minute + let mut interval = interval(METRIC_CLEANUP_INTERVAL); loop { interval.tick().await; - let Some(cutoff_time) = Instant::now().checked_sub(Duration::from_secs(180)) else { - continue; - }; + let current_epoch = activity_epoch.fetch_add(1, Ordering::Relaxed) + 1; let Some(counters) = counters.upgrade() else { break; }; - counters.retain(|_, metric_data: &mut Arc| { - Arc::strong_count(metric_data) > 1 - || unsafe { metric_data.get_last_updated() > cutoff_time } - }); - counters.shrink_to_fit(); + cleanup_metrics(&counters, current_epoch); } }))); } @@ -735,7 +738,7 @@ impl StatsManager { let metric_data = self .counters .entry(key.clone()) - .or_insert_with(|| Arc::new(MetricData::new())) + .or_insert_with(|| Arc::new(MetricData::new(Arc::clone(&self.activity_epoch)))) .clone(); CounterHandle::new(metric_data, key) @@ -1245,27 +1248,37 @@ mod tests { #[tokio::test] async fn test_cleanup_keeps_metrics_with_live_handles() { let stats = StatsManager::new(); + stats.stop_cleanup_task().await; + stats.activity_epoch.store(0, Ordering::Relaxed); + let counter = stats.get_simple_counter(MetricName::TrafficBytesForwarded); counter.set(1); - let cutoff_time = Instant::now().checked_add(Duration::from_secs(1)).unwrap(); - stats - .counters - .retain(|_, metric_data: &mut Arc| { - Arc::strong_count(metric_data) > 1 - || unsafe { metric_data.get_last_updated() > cutoff_time } - }); + let expired_epoch = METRIC_RETENTION_EPOCHS + 1; + cleanup_metrics(&stats.counters, expired_epoch); assert_eq!(stats.metric_count(), 1); assert_eq!(stats.get_all_metrics().len(), 1); drop(counter); - stats - .counters - .retain(|_, metric_data: &mut Arc| { - Arc::strong_count(metric_data) > 1 - || unsafe { metric_data.get_last_updated() > cutoff_time } - }); + cleanup_metrics(&stats.counters, expired_epoch); + assert_eq!(stats.metric_count(), 0); + } + + #[tokio::test] + async fn test_cleanup_retains_recently_updated_metrics_for_three_epochs() { + let stats = StatsManager::new(); + stats.stop_cleanup_task().await; + stats.activity_epoch.store(0, Ordering::Relaxed); + + let counter = stats.get_simple_counter(MetricName::TrafficBytesForwarded); + counter.set(1); + drop(counter); + + cleanup_metrics(&stats.counters, METRIC_RETENTION_EPOCHS); + assert_eq!(stats.metric_count(), 1); + + cleanup_metrics(&stats.counters, METRIC_RETENTION_EPOCHS + 1); assert_eq!(stats.metric_count(), 0); } diff --git a/easytier-core/src/peers/conn/peer.rs b/easytier-core/src/peers/conn/peer.rs index 04daa012..8a2f1b34 100644 --- a/easytier-core/src/peers/conn/peer.rs +++ b/easytier-core/src/peers/conn/peer.rs @@ -1,8 +1,9 @@ use std::sync::Arc; +use arc_swap::ArcSwapOption; use crossbeam::atomic::AtomicCell; use dashmap::{DashMap, DashSet}; -use parking_lot::RwLock; +use parking_lot::{Mutex, RwLock}; use tokio::{select, sync::mpsc}; @@ -38,11 +39,12 @@ pub struct Peer { shutdown_notifier: Arc, - default_conn_id: Arc>, + default_conn: Arc>, + default_conn_update_lock: Arc>, peer_identity_type: Arc>>, peer_public_key: Arc>>>, #[allow(dead_code)] - default_conn_id_clear_task: AbortOnDropHandle<()>, + default_conn_clear_task: AbortOnDropHandle<()>, } impl Peer { @@ -58,10 +60,14 @@ impl Peer { let peer_identity_type_copy = peer_identity_type.clone(); let peer_public_key = Arc::new(RwLock::new(None)); let peer_public_key_copy = peer_public_key.clone(); + let default_conn = Arc::new(ArcSwapOption::empty()); + let default_conn_update_lock = Arc::new(Mutex::new(())); let conns_copy = conns.clone(); let shutdown_notifier_copy = shutdown_notifier.clone(); let context_copy = context.clone(); + let default_conn_copy = default_conn.clone(); + let default_conn_update_lock_copy = default_conn_update_lock.clone(); let close_event_listener = AbortOnDropHandle::new(tokio::spawn( async move { loop { @@ -77,7 +83,22 @@ impl Peer { "notified that peer conn is closed", ); - if let Some((_, conn)) = conns_copy.remove(&ret) { + let removed_conn = { + let _update_guard = default_conn_update_lock_copy.lock(); + let removed_conn = conns_copy.remove(&ret); + if let Some((_, conn)) = removed_conn.as_ref() { + let cached_conn = default_conn_copy.load(); + if cached_conn + .as_ref() + .is_some_and(|cached| Arc::ptr_eq(cached, conn)) + { + default_conn_copy.store(None); + } + } + removed_conn + }; + + if let Some((_, conn)) = removed_conn { context_copy.issue_event(PeerEvent::PeerConnRemoved( conn.get_conn_info(), )); @@ -103,15 +124,13 @@ impl Peer { )), )); - let default_conn_id = Arc::new(AtomicCell::new(PeerConnId::default())); - let conns_copy = conns.clone(); - let default_conn_id_copy = default_conn_id.clone(); - let default_conn_id_clear_task = AbortOnDropHandle::new(tokio::spawn(async move { + let default_conn_copy = default_conn.clone(); + let default_conn_clear_task = AbortOnDropHandle::new(tokio::spawn(async move { loop { crate::foundation::time::sleep(std::time::Duration::from_secs(5)).await; if conns_copy.len() > 1 { - default_conn_id_copy.store(PeerConnId::default()); + default_conn_copy.store(None); } } })); @@ -126,10 +145,11 @@ impl Peer { close_event_listener, shutdown_notifier, - default_conn_id, + default_conn, + default_conn_update_lock, peer_identity_type, peer_public_key, - default_conn_id_clear_task, + default_conn_clear_task, } } @@ -186,29 +206,38 @@ impl Peer { Ok(()) } - async fn select_conn(&self) -> Option { - let default_conn_id = self.default_conn_id.load(); - if let Some(conn) = self.conns.get(&default_conn_id) { - return Some(conn.clone()); + fn select_conn(&self) -> Option { + let _update_guard = self.default_conn_update_lock.lock(); + if let Some(conn) = self.default_conn.load_full() { + return Some(conn); } // find a conn with the smallest latency let mut min_latency = u64::MAX; + let mut selected = None; for conn in self.conns.iter() { let latency = conn.value().get_stats().latency_us; if latency < min_latency { min_latency = latency; - self.default_conn_id.store(conn.get_conn_id()); + selected = Some(conn.value().clone()); } } - self.conns - .get(&self.default_conn_id.load()) - .map(|conn| conn.clone()) + if let Some(conn) = selected.as_ref() { + self.default_conn.store(Some(conn.clone())); + } + selected } pub async fn send_msg(&self, msg: ZCPacket) -> Result<(), Error> { - let Some(conn) = self.select_conn().await else { + let default_conn = self.default_conn.load(); + if let Some(conn) = default_conn.as_ref() { + conn.send_msg(msg).await?; + return Ok(()); + } + drop(default_conn); + + let Some(conn) = self.select_conn() else { return Err(Error::PeerNoConnectionError(self.peer_node_id)); }; conn.send_msg(msg).await?; @@ -264,7 +293,11 @@ impl Peer { } pub fn get_default_conn_id(&self) -> PeerConnId { - self.default_conn_id.load() + self.default_conn + .load() + .as_ref() + .map(|conn| conn.get_conn_id()) + .unwrap_or_default() } pub fn get_peer_identity_type(&self) -> Option { diff --git a/easytier-core/src/peers/conn/peer_conn.rs b/easytier-core/src/peers/conn/peer_conn.rs index 867d3da5..377abe97 100644 --- a/easytier-core/src/peers/conn/peer_conn.rs +++ b/easytier-core/src/peers/conn/peer_conn.rs @@ -35,6 +35,7 @@ use super::{ use crate::peers::{ PacketRecvChan, context::{ArcPeerContext, NetworkIdentity, NetworkSecretDigest}, + send_packet_to_chan, }; use crate::{ config::PeerId, @@ -1310,7 +1311,7 @@ impl PeerConn { if let Err(e) = ctrl_sender.send(zc_packet) { tracing::error!(?e, "peer conn send ctrl resp error"); } - } else if sender.send(zc_packet).await.is_err() { + } else if send_packet_to_chan(&sender, zc_packet).await.is_err() { break; } diff --git a/easytier-core/src/peers/conn/peer_map.rs b/easytier-core/src/peers/conn/peer_map.rs index 4a89496c..769dc0da 100644 --- a/easytier-core/src/peers/conn/peer_map.rs +++ b/easytier-core/src/peers/conn/peer_map.rs @@ -143,6 +143,10 @@ impl PeerMap { peer_id == self.my_peer_id || self.peer_map.contains_key(&peer_id) } + pub(crate) fn is_self(&self, peer_id: PeerId) -> bool { + peer_id == self.my_peer_id + } + pub async fn send_msg_directly(&self, msg: ZCPacket, dst_peer_id: PeerId) -> Result<(), Error> { if dst_peer_id == self.my_peer_id { let packet_send = self.packet_send.clone(); diff --git a/easytier-core/src/peers/context.rs b/easytier-core/src/peers/context.rs index cd63a0c6..149a7128 100644 --- a/easytier-core/src/peers/context.rs +++ b/easytier-core/src/peers/context.rs @@ -763,7 +763,8 @@ impl PeerContext for CorePeerContext { } fn packet_policy(&self) -> PeerPacketPolicy { - PeerPacketPolicy::from_flags(&self.snapshot().flags) + self.config + .with_snapshot(|snapshot| PeerPacketPolicy::from_flags(&snapshot.peer.flags)) } fn host_routing_policy(&self) -> HostRoutingPolicy { @@ -792,13 +793,16 @@ impl PeerContext for CorePeerContext { } fn ipv4(&self) -> Option { - self.snapshot() - .runtime - .core - .routes - .ipv4 - .as_ref() - .and_then(config_ipv4) + self.config.with_snapshot(|snapshot| { + snapshot + .peer + .runtime + .core + .routes + .ipv4 + .as_ref() + .and_then(config_ipv4) + }) } fn ipv6(&self) -> Option { diff --git a/easytier-core/src/peers/mod.rs b/easytier-core/src/peers/mod.rs index 244c3a8b..5332056e 100644 --- a/easytier-core/src/peers/mod.rs +++ b/easytier-core/src/peers/mod.rs @@ -21,6 +21,7 @@ pub(crate) mod test_support; mod tests; use crate::packet::ZCPacket; +use tokio::sync::mpsc::error::{SendError, TrySendError}; pub type PacketRecvChan = tokio::sync::mpsc::Sender; pub type PacketRecvChanReceiver = tokio::sync::mpsc::Receiver; @@ -29,6 +30,17 @@ pub fn create_packet_recv_chan() -> (PacketRecvChan, PacketRecvChanReceiver) { tokio::sync::mpsc::channel(128) } +pub(crate) async fn send_packet_to_chan( + sender: &PacketRecvChan, + packet: ZCPacket, +) -> Result<(), SendError> { + match sender.try_send(packet) { + Ok(()) => Ok(()), + Err(TrySendError::Full(packet)) => sender.send(packet).await, + Err(TrySendError::Closed(packet)) => Err(SendError(packet)), + } +} + pub async fn recv_packet_from_chan( packet_recv_chan_receiver: &mut PacketRecvChanReceiver, ) -> Result { diff --git a/easytier-core/src/peers/peer_manager.rs b/easytier-core/src/peers/peer_manager.rs index 73befbd6..396f34ac 100644 --- a/easytier-core/src/peers/peer_manager.rs +++ b/easytier-core/src/peers/peer_manager.rs @@ -3153,7 +3153,9 @@ pub(crate) async fn send_msg_internal( && (peers.has_peer(gateway) || foreign_network_client.has_next_hop(gateway)) { relay_peer_map.send_msg(msg, dst_peer_id, policy).await - } else if peers.has_peer(dst_peer_id) { + } else if let Some(peer) = peers.get_peer_by_id(dst_peer_id) { + peer.send_msg(msg).await + } else if peers.is_self(dst_peer_id) { peers.send_msg_directly(msg, dst_peer_id).await } else if foreign_network_client.has_next_hop(dst_peer_id) { foreign_network_client.send_msg(msg, dst_peer_id).await diff --git a/easytier-core/src/peers/tests.rs b/easytier-core/src/peers/tests.rs index 9a398eb3..a2a93c3b 100644 --- a/easytier-core/src/peers/tests.rs +++ b/easytier-core/src/peers/tests.rs @@ -5,7 +5,11 @@ use crate::foundation::time::{Duration, timeout}; use crate::{ packet::{PacketType, ZCPacket}, peers::{ - conn::{peer_conn::PeerConn, peer_map::PeerMap, peer_session::PeerSessionStore}, + conn::{ + peer_conn::{PeerConn, PeerConnId}, + peer_map::PeerMap, + peer_session::PeerSessionStore, + }, context::NetworkIdentity, create_packet_recv_chan, error::Error, @@ -111,3 +115,90 @@ async fn peer_map_forwards_packet_over_memory_tunnel() { .unwrap(); assert_eq!(received.payload(), b"hello"); } + +#[tokio::test] +async fn peer_map_reselects_cached_connection_after_close() { + let peer_session_store = Arc::new(PeerSessionStore::new()); + let (client_tunnel_a, server_tunnel_a) = create_ring_tunnel_pair(); + let (client_tunnel_b, server_tunnel_b) = create_ring_tunnel_pair(); + let client_ctx = Arc::new(NoopPeerContext::default()); + let server_ctx = Arc::new(NoopPeerContext::default()); + + let mut client_conn_a = PeerConn::new( + 1, + client_ctx.clone(), + client_tunnel_a, + peer_session_store.clone(), + ); + let mut server_conn_a = PeerConn::new( + 2, + server_ctx.clone(), + server_tunnel_a, + peer_session_store.clone(), + ); + let mut client_conn_b = PeerConn::new( + 1, + client_ctx.clone(), + client_tunnel_b, + peer_session_store.clone(), + ); + let mut server_conn_b = + PeerConn::new(2, server_ctx.clone(), server_tunnel_b, peer_session_store); + + let (client_a_ret, server_a_ret, client_b_ret, server_b_ret) = tokio::join!( + client_conn_a.do_handshake_as_client(), + server_conn_a.do_handshake_as_server(), + client_conn_b.do_handshake_as_client(), + server_conn_b.do_handshake_as_server(), + ); + client_a_ret.unwrap(); + server_a_ret.unwrap(); + client_b_ret.unwrap(); + server_b_ret.unwrap(); + + let (client_tx, _client_rx) = create_packet_recv_chan(); + let (server_tx, mut server_rx) = create_packet_recv_chan(); + let client_map = PeerMap::new(client_tx, client_ctx, 1); + let server_map = PeerMap::new(server_tx, server_ctx, 2); + + client_map.add_new_peer_conn(client_conn_a).await.unwrap(); + client_map.add_new_peer_conn(client_conn_b).await.unwrap(); + server_map.add_new_peer_conn(server_conn_a).await.unwrap(); + server_map.add_new_peer_conn(server_conn_b).await.unwrap(); + + let mut first_packet = ZCPacket::new_with_payload(b"first"); + first_packet.fill_peer_manager_hdr(1, 2, PacketType::Data as u8); + client_map.send_msg_directly(first_packet, 2).await.unwrap(); + let first_received = timeout(Duration::from_secs(1), server_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(first_received.payload(), b"first"); + + let first_conn_id = client_map.get_peer_default_conn_id(2).await.unwrap(); + assert_ne!(first_conn_id, PeerConnId::default()); + client_map.close_peer_conn(2, &first_conn_id).await.unwrap(); + timeout(Duration::from_secs(1), async { + while client_map.get_peer_default_conn_id(2).await == Some(first_conn_id) { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + let mut second_packet = ZCPacket::new_with_payload(b"second"); + second_packet.fill_peer_manager_hdr(1, 2, PacketType::Data as u8); + client_map + .send_msg_directly(second_packet, 2) + .await + .unwrap(); + let second_received = timeout(Duration::from_secs(1), server_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(second_received.payload(), b"second"); + assert_ne!( + client_map.get_peer_default_conn_id(2).await, + Some(first_conn_id) + ); +} diff --git a/easytier-core/src/tunnel/mpsc.rs b/easytier-core/src/tunnel/mpsc.rs index 18d8ab3d..72267c37 100644 --- a/easytier-core/src/tunnel/mpsc.rs +++ b/easytier-core/src/tunnel/mpsc.rs @@ -18,7 +18,13 @@ pub struct MpscTunnelSender(Sender); impl MpscTunnelSender { pub async fn send(&self, item: ZCPacket) -> Result<(), TunnelError> { - self.0.send(item).await.map_err(|_| TunnelError::Shutdown) + match self.0.try_send(item) { + Ok(()) => Ok(()), + Err(TrySendError::Full(item)) => { + self.0.send(item).await.map_err(|_| TunnelError::Shutdown) + } + Err(TrySendError::Closed(_)) => Err(TunnelError::Shutdown), + } } pub fn try_send(&self, item: ZCPacket) -> Result<(), TunnelError> { @@ -137,3 +143,33 @@ impl Drop for MpscTunnel { self.task.abort(); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn sender_falls_back_to_waiting_when_channel_is_full() { + let (tx, mut rx) = channel(1); + let sender = MpscTunnelSender(tx); + sender + .send(ZCPacket::new_with_payload(b"first")) + .await + .unwrap(); + + let blocked_send = sender.send(ZCPacket::new_with_payload(b"second")); + tokio::pin!(blocked_send); + assert!( + tokio::time::timeout(Duration::from_millis(10), &mut blocked_send) + .await + .is_err() + ); + + assert_eq!(rx.recv().await.unwrap().payload(), b"first"); + tokio::time::timeout(Duration::from_secs(1), &mut blocked_send) + .await + .unwrap() + .unwrap(); + assert_eq!(rx.recv().await.unwrap().payload(), b"second"); + } +}