mirror of
https://github.com/EasyTier/EasyTier.git
synced 2026-09-02 09:09:17 +00:00
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:
@@ -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(),
|
||||
}
|
||||
});
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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());
|
||||
|
||||
Reference in New Issue
Block a user