diff --git a/easytier-core/src/foundation/token_bucket.rs b/easytier-core/src/foundation/token_bucket.rs index c692cd65..a22c441d 100644 --- a/easytier-core/src/foundation/token_bucket.rs +++ b/easytier-core/src/foundation/token_bucket.rs @@ -12,6 +12,11 @@ use crate::foundation::time; pub(crate) trait ByteLimiter: Send + Sync { fn try_consume(&self, bytes: u64) -> bool; + + /// Bytes charged by this limiter so far. + fn metered_bytes(&self) -> u64 { + 0 + } } impl ByteLimiter for () { @@ -37,6 +42,7 @@ struct BucketState { available_tokens: u64, last_refill: Instant, refill_remainder: u64, + metered_bytes: u64, } #[derive(Clone, Copy)] @@ -78,6 +84,7 @@ impl TokenBucket { available_tokens: config.capacity, last_refill: Instant::now(), refill_remainder: 0, + metered_bytes: 0, }), config, stop_notifier: Notify::new(), @@ -138,9 +145,15 @@ impl TokenBucket { } state.available_tokens -= tokens; + state.metered_bytes = state.metered_bytes.saturating_add(tokens); true } + /// Total bytes charged by this bucket so far. + pub fn metered_bytes(&self) -> u64 { + self.state.lock().metered_bytes + } + /// Consume tokens, sleeping until they become available. pub async fn consume(&self, tokens: u64) { let mut remaining = tokens; @@ -166,6 +179,7 @@ impl TokenBucket { self.refill(&mut state, Instant::now()); if state.available_tokens >= tokens { state.available_tokens -= tokens; + state.metered_bytes = state.metered_bytes.saturating_add(tokens); return; } @@ -195,6 +209,10 @@ impl ByteLimiter for TokenBucket { fn try_consume(&self, bytes: u64) -> bool { TokenBucket::try_consume(self, bytes) } + + fn metered_bytes(&self) -> u64 { + TokenBucket::metered_bytes(self) + } } pub struct TokenBucketManager { diff --git a/easytier-core/src/instance/test_utils.rs b/easytier-core/src/instance/test_utils.rs index c7eebe48..0cb318ec 100644 --- a/easytier-core/src/instance/test_utils.rs +++ b/easytier-core/src/instance/test_utils.rs @@ -87,4 +87,20 @@ where .get_peer_session_store() .evict_unused_sessions_idle(idle); } + + #[doc(hidden)] + pub fn foreign_forward_metered_bytes_for_test(&self, network_name: &str) -> Option { + self.peer_manager + .foreign_forward_metered_bytes(network_name) + } + + #[doc(hidden)] + pub fn network_name_for_test(&self) -> String { + self.peer_manager.network_name().to_owned() + } + + #[doc(hidden)] + pub fn instance_recv_metered_bytes_for_test(&self) -> Option { + self.peer_manager.instance_recv_metered_bytes() + } } diff --git a/easytier-core/src/peers/peer_manager.rs b/easytier-core/src/peers/peer_manager.rs index 7fc2f248..01a59713 100644 --- a/easytier-core/src/peers/peer_manager.rs +++ b/easytier-core/src/peers/peer_manager.rs @@ -1328,6 +1328,22 @@ impl PeerManagerCore { self.acl_filter.get_stats() } + /// Bytes charged by the foreign network forward limiter, if configured. + #[doc(hidden)] + pub fn foreign_forward_metered_bytes(&self, network_name: &str) -> Option { + self.context + .foreign_forward_limiter(network_name) + .map(|limiter| limiter.metered_bytes()) + } + + /// Bytes charged by the instance-wide receive limiter, if configured. + #[doc(hidden)] + pub fn instance_recv_metered_bytes(&self) -> Option { + self.context + .recv_limiter("", false) + .map(|limiter| limiter.metered_bytes()) + } + pub fn acl_filter(&self) -> Arc { self.acl_filter.clone() } diff --git a/easytier/src/tests/three_node.rs b/easytier/src/tests/three_node.rs index f926611c..7225f6dd 100644 --- a/easytier/src/tests/three_node.rs +++ b/easytier/src/tests/three_node.rs @@ -3384,18 +3384,26 @@ pub async fn relay_bps_limit_test(#[values(100, 200, 400, 800)] bps_limit: u64) let tcp_listener = core_tcp_listener("tcp://0.0.0.0:22223".parse().unwrap()); let tcp_connector = core_tcp_dialer("tcp://10.144.144.3:22223".parse().unwrap()); - let bps = _tunnel_bench_netns( + // The relay keys its forward limiter by the peer network name, which is + // the dialing instance's identity, not inst2's own "public" one. + let relay_network = insts[0].get_core_instance().network_name_for_test(); + let meter = || relay_forward_metered_bytes(&insts[1], &relay_network); + + let (bps, metered) = _tunnel_bench_netns( tcp_listener, tcp_connector, NetNS::new(Some(network.namespace("net_c").into())), NetNS::new(Some(network.namespace("net_a").into())), + Some(&meter), ) .await; + let metered = metered.expect("metered window should be sampled"); - println!("bps: {}", bps); - - let bps = bps as u64 / 1024; - assert_limited_payload_bps(bps, bps_limit); + println!( + "payload bps: {}, metered: {} over {:?}", + bps, metered.0, metered.1 + ); + assert_limited_bps(metered, bps, bps_limit); drop_insts(insts).await; } @@ -3423,33 +3431,67 @@ pub async fn instance_recv_bps_limit_test(#[values(100, 800)] bps_limit: u64) { let tcp_listener = core_tcp_listener("tcp://0.0.0.0:22223".parse().unwrap()); let tcp_connector = core_tcp_dialer("tcp://10.144.144.3:22223".parse().unwrap()); - let bps = _tunnel_bench_netns( + let meter = || instance_recv_metered_bytes(&insts[1]); + + let (bps, metered) = _tunnel_bench_netns( tcp_listener, tcp_connector, NetNS::new(Some(network.namespace("net_c").into())), NetNS::new(Some(network.namespace("net_a").into())), + Some(&meter), ) .await; + let metered = metered.expect("metered window should be sampled"); - println!("bps: {}", bps); - - let bps = bps as u64 / 1024; - assert_limited_payload_bps(bps, bps_limit); + println!( + "payload bps: {}, metered: {} over {:?}", + bps, metered.0, metered.1 + ); + assert_limited_bps(metered, bps, bps_limit); drop_insts(insts).await; } -fn assert_limited_payload_bps(bps: u64, bps_limit: u64) { - // The benchmark measures TCP application payload while the limiter counts - // EasyTier data payload, including the inner IP and transport headers. - let min_bps = bps_limit.saturating_sub((bps_limit / 10).max(50)); - let max_bps = bps_limit + 50; +fn relay_forward_metered_bytes(inst: &Instance, network_name: &str) -> u64 { + inst.get_core_instance() + .foreign_forward_metered_bytes_for_test(network_name) + .expect("relay forward limiter should be configured") +} + +fn instance_recv_metered_bytes(inst: &Instance) -> u64 { + inst.get_core_instance() + .instance_recv_metered_bytes_for_test() + .expect("instance receive limiter should be configured") +} + +fn assert_limited_bps(metered: (u64, Duration), payload_bps: usize, bps_limit: u64) { + // The limiter charges each packet with its EasyTier payload length, which + // includes the inner IP and transport headers, so its counter is the rate + // it actually enforced. The benchmark only sees TCP application payload, + // so comparing the two directly made this assertion flaky. The window + // starts once the initial bucket burst is drained. + let (bytes, window) = metered; + let expected = bps_limit as f64 * 1024.0 * window.as_secs_f64(); + let slack = expected * 0.05; assert!( - bps >= min_bps && bps <= max_bps, - "bps: {}, expected: {}..={}", - bps, - min_bps, - max_bps + (bytes as f64) >= expected - slack && (bytes as f64) <= expected + slack, + "metered: {} over {:?}, expected: {:.0}..={:.0}", + bytes, + window, + expected - slack, + expected + slack + ); + + // Encapsulation only shrinks the payload the receiver sees, so a limiter + // whose rejection the forwarding path ignores still shows up here. The + // benchmark rounds its window down to whole seconds, which inflates the + // payload rate by a few percent, hence the same 50 KiB/s headroom. + let payload_bps = payload_bps as u64 / 1024; + assert!( + payload_bps <= bps_limit + 50, + "payload bps: {}, expected: <= {}", + payload_bps, + bps_limit + 50 ); } diff --git a/easytier/src/tunnel/common.rs b/easytier/src/tunnel/common.rs index 32aebc3a..e25ebbbc 100644 --- a/easytier/src/tunnel/common.rs +++ b/easytier/src/tunnel/common.rs @@ -571,7 +571,8 @@ pub(crate) mod tests { connector: C, netns_l: NetNS, netns_c: NetNS, - ) -> usize + meter: Option<&(dyn Fn() -> u64 + Sync)>, + ) -> (usize, Option<(u64, std::time::Duration)>) where L: SocketListener> + Sync + 'static, C: TunnelDialer, @@ -618,11 +619,23 @@ pub(crate) mod tests { } let now = Instant::now(); + let mut meter_base = None; while now.elapsed().as_secs() < 10 { + // Skip the initial burst so the metered window reflects the + // sustained rate the limiter enforces. + if meter_base.is_none() && now.elapsed() >= std::time::Duration::from_secs(1) { + meter_base = meter.map(|meter| (meter(), now.elapsed())); + } // send.feed(item) let item = ZCPacket::new_with_payload(send_buf.as_ref()); send.feed(item).await.unwrap(); } + let metered = match (meter_base, meter) { + (Some((metered, started)), Some(meter)) => { + Some((meter().saturating_sub(metered), now.elapsed() - started)) + } + _ => None, + }; send.close().await.unwrap(); drop(send); @@ -634,7 +647,7 @@ pub(crate) mod tests { println!("bps: {}", bps); lis.abort(); - bps as usize + (bps as usize, metered) } pub async fn wait_for_condition(mut condition: F, timeout: std::time::Duration)