refactor(core): use linearizable lazy token bucket (#2421)

Replace periodic refill tasks with on-demand accounting to avoid waking
idle token buckets.

Keep balance, refill time, and fractional credit in one locked state so
concurrent consumers cannot observe partially published refills or
exceed the configured burst capacity. Track credit in nanoseconds and
discard excess credit at capacity to preserve precise limiter behavior.

Use a one-second default burst capacity to preserve the existing
limiter behavior while supporting explicit capacity configuration. Keep
limiter capacity and fill rate in a local config instead of an unused
protobuf message.

Charge only logical EasyTier data payload, unwrap foreign network
packets before accounting, and leave control traffic outside the
limiter. Reject forged payload lengths by accounting from actual packet
boundaries.

Split oversized blocking consumes into capacity-sized chunks and cover
concurrency, refill precision, burst caps, payload accounting, and
bandwidth integration behavior.
This commit is contained in:
Chenx Dust
2026-08-04 10:08:37 +08:00
committed by GitHub
parent 8d475dc3fc
commit df874b85be
9 changed files with 387 additions and 216 deletions
+1 -5
View File
@@ -786,11 +786,7 @@ impl AclProcessor {
panic!("Rate limit bucket not found");
}
RateLimitValue {
token_bucket: TokenBucket::new(
burst as u64,
rate as u64,
Duration::from_millis(10),
),
token_bucket: TokenBucket::new(burst as u64, rate as u64),
last_update: Instant::now(),
}
});
+6 -2
View File
@@ -36,6 +36,7 @@ use crate::peers::{
PacketRecvChan,
context::{ArcPeerContext, NetworkIdentity, NetworkSecretDigest},
send_packet_to_chan,
traffic_metrics::data_packet_payload_len,
};
use crate::{
config::PeerId,
@@ -1290,6 +1291,7 @@ impl PeerConn {
let mut zc_packet = ret.unwrap();
let buf_len = zc_packet.buf_len() as u64;
let limited_payload_len = data_packet_payload_len(&zc_packet);
let Some(peer_mgr_hdr) = zc_packet.mut_peer_manager_header() else {
tracing::error!(
"unexpected packet: {:?}, cannot decode peer manager hdr",
@@ -1315,8 +1317,10 @@ impl PeerConn {
break;
}
if let Some(limiter) = recv_limiter.as_ref() {
limiter.consume(buf_len).await;
if let Some(payload_len) = limited_payload_len
&& let Some(limiter) = recv_limiter.as_ref()
{
limiter.consume(payload_len).await;
}
}
+3 -15
View File
@@ -13,9 +13,7 @@ use cidr::{Ipv4Cidr, Ipv4Inet, Ipv6Cidr, Ipv6Inet};
use dashmap::DashMap;
use easytier_proto::{
acl::Acl,
common::{
FlagsInConfig, LimiterConfig, PeerFeatureFlag, SecureModeConfig, StunInfo, TunnelInfo,
},
common::{FlagsInConfig, PeerFeatureFlag, SecureModeConfig, StunInfo, TunnelInfo},
peer_rpc::{PeerGroupInfo, TrustedCredentialPubkeyProof},
};
use hmac::{Hmac, Mac};
@@ -31,7 +29,7 @@ use crate::{
},
events::{CoreEvent, CoreEventSink},
foundation::stats::{LabelSet, LabelType, MetricName, StatsManager},
foundation::token_bucket::{ArcByteLimiter, TokenBucketManager},
foundation::token_bucket::{ArcByteLimiter, BucketConfig, TokenBucketManager},
peers::{
credential_manager::{CredentialManager, CredentialStorage},
util::shrink_dashmap,
@@ -362,17 +360,7 @@ impl CorePeerContext {
return None;
}
let manager = state.manager.get_or_insert_with(TokenBucketManager::new);
Some(
manager.get_or_create(
key,
LimiterConfig {
burst_rate: None,
bps: Some(bps),
fill_duration_ms: None,
}
.into(),
),
)
Some(manager.get_or_create(key, BucketConfig::with_default_capacity(bps)))
}
pub(crate) async fn stop(&self) {
@@ -51,7 +51,8 @@ use super::{
relay_peer_map::RelayPeerMap,
route::{NextHopPolicy, Route, RouteInterface, peer_ospf_route::PeerRoute},
traffic_metrics::{
TrafficKind, TrafficMetricRecorder, is_relay_data_packet_type, traffic_kind,
TrafficKind, TrafficMetricRecorder, data_packet_payload_len, is_relay_data_packet_type,
traffic_kind,
},
util::shrink_dashmap,
whitelist::check_network_in_relay_whitelist,
@@ -1347,8 +1348,9 @@ impl ForeignNetworkPacketRouter {
);
continue;
}
if let Some(bps_limiter) = bps_limiter.as_ref()
&& !bps_limiter.try_consume(len.into())
if let Some(payload_len) = data_packet_payload_len(&zc_packet)
&& let Some(bps_limiter) = bps_limiter.as_ref()
&& !bps_limiter.try_consume(payload_len)
{
continue;
}
+101 -1
View File
@@ -5,7 +5,7 @@ use futures::future::BoxFuture;
use crate::config::PeerId;
use crate::foundation::stats::{CounterHandle, LabelSet, LabelType, MetricName, StatsManager};
use crate::packet::PacketType;
use crate::packet::{PacketType, ZCPacket};
use crate::peers::util::shrink_dashmap;
use crate::proto::peer_rpc::RoutePeerInfo;
@@ -200,6 +200,26 @@ pub fn is_relay_data_packet_type(packet_type: u8) -> bool {
|| packet_type == PacketType::ForeignNetworkPacket as u8
}
pub(crate) fn data_packet_payload_len(packet: &ZCPacket) -> Option<u64> {
let header = packet.peer_manager_header()?;
if header.packet_type == PacketType::ForeignNetworkPacket as u8 {
if header.is_encrypted() {
return Some(packet.payload_len() as u64);
}
return match packet.foreign_network_inner_packet_info() {
Some((inner_header, payload_len))
if traffic_kind(inner_header.packet_type) == TrafficKind::Data =>
{
Some(payload_len as u64)
}
Some(_) => None,
None => Some(packet.payload_len() as u64),
};
}
(traffic_kind(header.packet_type) == TrafficKind::Data).then(|| packet.payload_len() as u64)
}
#[derive(Clone)]
struct TrafficMetricGroup {
data: Arc<LogicalTrafficMetrics>,
@@ -327,6 +347,86 @@ mod tests {
.with_label_type(LabelType::ToInstanceId(instance_id.to_string()))
}
#[test]
fn relay_limiter_classifies_only_data_packets() {
for packet_type in [
PacketType::Data,
PacketType::KcpSrc,
PacketType::KcpDst,
PacketType::QuicSrc,
PacketType::QuicDst,
PacketType::DataWithKcpSrcModified,
PacketType::DataWithQuicSrcModified,
PacketType::ForeignNetworkPacket,
] {
assert!(is_relay_data_packet_type(packet_type as u8));
}
for packet_type in [
PacketType::Ping,
PacketType::Pong,
PacketType::RpcReq,
PacketType::RpcResp,
PacketType::RelayHandshake,
PacketType::RelayHandshakeAck,
] {
assert!(!is_relay_data_packet_type(packet_type as u8));
}
}
fn packet_with_type(packet_type: PacketType, payload_len: usize) -> ZCPacket {
let mut packet = ZCPacket::new_with_payload(&vec![0; payload_len]);
packet.fill_peer_manager_hdr(1, 2, packet_type as u8);
packet
}
#[test]
fn data_packet_payload_len_uses_easytier_payload() {
let mut data_packet = packet_with_type(PacketType::Data, 1_024);
data_packet.mut_peer_manager_header().unwrap().len.set(0);
assert_eq!(data_packet_payload_len(&data_packet), Some(1_024));
let control_packet = packet_with_type(PacketType::RpcReq, 128);
assert_eq!(data_packet_payload_len(&control_packet), None);
}
#[test]
fn data_packet_payload_len_unwraps_foreign_network_packet() {
let network_name = "foreign".to_string();
let mut inner_data = packet_with_type(PacketType::Data, 1_024);
inner_data.mut_peer_manager_header().unwrap().len.set(0);
let foreign_data = ZCPacket::new_for_foreign_network(&network_name, 2, &inner_data);
assert!(!foreign_data.peer_manager_header().unwrap().is_encrypted());
assert_eq!(data_packet_payload_len(&foreign_data), Some(1_024));
let inner_control = packet_with_type(PacketType::RpcReq, 128);
let foreign_control = ZCPacket::new_for_foreign_network(&network_name, 2, &inner_control);
assert_eq!(data_packet_payload_len(&foreign_control), None);
let malformed_foreign = packet_with_type(PacketType::ForeignNetworkPacket, 7);
assert_eq!(
data_packet_payload_len(&malformed_foreign),
Some(malformed_foreign.payload_len() as u64)
);
}
#[test]
fn data_packet_payload_len_charges_encrypted_foreign_payload() {
let network_name = "foreign".to_string();
let inner_control = packet_with_type(PacketType::RpcReq, 128);
let mut encrypted_foreign =
ZCPacket::new_for_foreign_network(&network_name, 2, &inner_control);
encrypted_foreign
.mut_peer_manager_header()
.unwrap()
.set_encrypted(true);
assert_eq!(
data_packet_payload_len(&encrypted_foreign),
Some(encrypted_foreign.payload_len() as u64)
);
}
#[tokio::test]
async fn logical_traffic_metrics_upgrade_unknown_instance_label() {
let stats_mgr = Arc::new(StatsManager::new());