diff --git a/easytier-core/src/gateway/dataplane/flow.rs b/easytier-core/src/gateway/dataplane/flow.rs index d51e5de3..8ed78ff3 100644 --- a/easytier-core/src/gateway/dataplane/flow.rs +++ b/easytier-core/src/gateway/dataplane/flow.rs @@ -252,7 +252,7 @@ impl FlowTable { fn increment_count(&self) -> FlowCountChange { let previous = self .count - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| { + .try_update(Ordering::AcqRel, Ordering::Acquire, |count| { count.checked_add(1) }) .expect("flow count overflow"); @@ -294,7 +294,7 @@ impl FlowTable { let previous = self .count - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| { + .try_update(Ordering::AcqRel, Ordering::Acquire, |count| { count.checked_sub(delta) }) .expect("flow count underflow"); diff --git a/easytier-core/src/gateway/proxy/tcp_proxy_engine.rs b/easytier-core/src/gateway/proxy/tcp_proxy_engine.rs index 298fe9e1..150c3b73 100644 --- a/easytier-core/src/gateway/proxy/tcp_proxy_engine.rs +++ b/easytier-core/src/gateway/proxy/tcp_proxy_engine.rs @@ -186,7 +186,7 @@ impl TcpProxyEngine { for _ in 0..u16::MAX { let translated_port = self .next_translated_port - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |port| { + .try_update(Ordering::Relaxed, Ordering::Relaxed, |port| { Some(if port == u16::MAX { 1 } else { port + 1 }) }) .expect("translated port counter update cannot fail"); diff --git a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs index 24850674..6184d43e 100644 --- a/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs +++ b/easytier-core/src/gateway/smoltcp/tokio_smoltcp/mod.rs @@ -126,7 +126,7 @@ impl Net { } pub fn get_port(&self) -> u16 { self.from_port - .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |x| { + .try_update(Ordering::SeqCst, Ordering::SeqCst, |x| { Some(if x > 60000 { 10000 } else { x + 1 }) }) .unwrap() @@ -188,3 +188,36 @@ impl Drop for Net { self.stopper.notify_waiters() } } + +#[cfg(test)] +mod tests { + use super::*; + use smoltcp::{phy::Medium, wire::HardwareAddress}; + + #[allow(clippy::unused_async)] + #[tokio::test] + async fn port_allocation_returns_previous_value_and_wraps() { + let mut capabilities = device::DeviceCapabilities::default(); + capabilities.medium = Medium::Ip; + capabilities.max_transmission_unit = 1280; + let (device, _ingress, _egress) = channel_device::ChannelDevice::new(capabilities); + let net = Net::new( + device, + NetConfig::new( + Config::new(HardwareAddress::Ip), + "10.126.126.1/24".parse().unwrap(), + Vec::new(), + None, + ), + ); + + assert_eq!(net.get_port(), 10001); + assert_eq!(net.get_port(), 10002); + + net.from_port.store(60000, Ordering::SeqCst); + assert_eq!(net.get_port(), 60000); + assert_eq!(net.get_port(), 60001); + assert_eq!(net.get_port(), 10000); + assert_eq!(net.get_port(), 10001); + } +} diff --git a/easytier/src/common/log/tracing_backend.rs b/easytier/src/common/log/tracing_backend.rs index 4d2842ff..cc71e31e 100644 --- a/easytier/src/common/log/tracing_backend.rs +++ b/easytier/src/common/log/tracing_backend.rs @@ -70,7 +70,7 @@ impl tracing::Subscriber for EventSubscriber { fn new_span(&self, _attributes: &Attributes<'_>) -> Id { let id = self .next_span_id - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| { + .try_update(Ordering::Relaxed, Ordering::Relaxed, |id| { Some(if id == usize::MAX { 1 } else { id + 1 }) }) .expect("span id update always succeeds");