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"); + } +}