mirror of
https://github.com/EasyTier/EasyTier.git
synced 2026-08-05 20:19:45 +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:
Chenx Dust
@@ -1,13 +1,14 @@
|
||||
use atomic_shim::AtomicU64;
|
||||
use dashmap::DashMap;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use parking_lot::Mutex;
|
||||
use std::sync::{
|
||||
Arc, Mutex as StdMutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
};
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::sync::Notify;
|
||||
use tokio_util::task::AbortOnDropHandle;
|
||||
|
||||
use crate::foundation::time;
|
||||
use easytier_proto::common::LimiterConfig;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
pub(crate) trait ByteLimiter: Send + Sync {
|
||||
@@ -27,49 +28,46 @@ impl ByteLimiter for () {
|
||||
|
||||
pub(crate) type ArcByteLimiter = Arc<dyn ByteLimiter>;
|
||||
|
||||
/// Token Bucket rate limiter using atomic operations
|
||||
pub struct TokenBucket {
|
||||
available_tokens: AtomicU64, // Current token count (atomic)
|
||||
last_refill_time: AtomicU64, // Last refill time as micros since epoch
|
||||
config: BucketConfig, // Immutable configuration
|
||||
refill_task: Mutex<Option<AbortOnDropHandle<()>>>, // Background refill task
|
||||
start_time: Instant, // Bucket creation time
|
||||
const MIN_FILL_RATE: u64 = 8_196;
|
||||
const NANOS_PER_SECOND: u128 = 1_000_000_000;
|
||||
|
||||
refill_notifier: Arc<Notify>,
|
||||
/// Token Bucket rate limiter with on-demand refill.
|
||||
pub struct TokenBucket {
|
||||
state: Mutex<BucketState>,
|
||||
config: BucketConfig,
|
||||
stop_notifier: Notify,
|
||||
stopped: AtomicBool,
|
||||
}
|
||||
|
||||
struct BucketState {
|
||||
available_tokens: u64,
|
||||
last_refill: Instant,
|
||||
refill_remainder: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct BucketConfig {
|
||||
capacity: u64, // Maximum token capacity
|
||||
fill_rate: u64, // Tokens added per second
|
||||
refill_interval: Duration, // Time between refill operations
|
||||
capacity: u64,
|
||||
fill_rate: u64,
|
||||
}
|
||||
|
||||
impl From<LimiterConfig> for BucketConfig {
|
||||
fn from(cfg: LimiterConfig) -> Self {
|
||||
let burst_rate = 1.max(cfg.burst_rate.unwrap_or(1));
|
||||
let fill_rate = 8196.max(cfg.bps.unwrap_or(u64::MAX / burst_rate));
|
||||
let refill_interval = cfg
|
||||
.fill_duration_ms
|
||||
.map(|x| Duration::from_millis(1.max(x)))
|
||||
.unwrap_or(Duration::from_millis(10));
|
||||
BucketConfig {
|
||||
capacity: burst_rate * fill_rate,
|
||||
impl BucketConfig {
|
||||
pub fn new(capacity: u64, fill_rate: u64) -> Self {
|
||||
Self {
|
||||
capacity,
|
||||
fill_rate,
|
||||
refill_interval,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_default_capacity(fill_rate: u64) -> Self {
|
||||
let fill_rate = fill_rate.max(MIN_FILL_RATE);
|
||||
Self::new(fill_rate, fill_rate)
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenBucket {
|
||||
pub fn new(capacity: u64, bps: u64, refill_interval: Duration) -> Arc<Self> {
|
||||
let config = BucketConfig {
|
||||
capacity,
|
||||
fill_rate: bps,
|
||||
refill_interval,
|
||||
};
|
||||
Self::new_from_cfg(config)
|
||||
pub fn new(capacity: u64, bps: u64) -> Arc<Self> {
|
||||
Self::new_from_cfg(BucketConfig::new(capacity, bps))
|
||||
}
|
||||
|
||||
/// Creates a new Token Bucket rate limiter
|
||||
@@ -77,79 +75,50 @@ impl TokenBucket {
|
||||
/// # Arguments
|
||||
/// * `capacity` - Bucket capacity in bytes
|
||||
/// * `bps` - Bandwidth limit in bytes per second
|
||||
/// * `refill_interval` - Refill interval (recommended 10-50ms)
|
||||
pub fn new_from_cfg(config: BucketConfig) -> Arc<Self> {
|
||||
// Create Arc instance with placeholder task
|
||||
let arc_self = Arc::new(Self {
|
||||
available_tokens: AtomicU64::new(config.capacity),
|
||||
last_refill_time: AtomicU64::new(0),
|
||||
pub fn new_from_cfg(mut config: BucketConfig) -> Arc<Self> {
|
||||
config.capacity = config.capacity.max(1);
|
||||
config.fill_rate = config.fill_rate.max(1);
|
||||
|
||||
Arc::new(Self {
|
||||
state: Mutex::new(BucketState {
|
||||
available_tokens: config.capacity,
|
||||
last_refill: Instant::now(),
|
||||
refill_remainder: 0,
|
||||
}),
|
||||
config,
|
||||
refill_task: Mutex::new(None),
|
||||
start_time: std::time::Instant::now(),
|
||||
refill_notifier: Arc::new(Notify::new()),
|
||||
stop_notifier: Notify::new(),
|
||||
stopped: AtomicBool::new(false),
|
||||
});
|
||||
|
||||
// Start background refill task
|
||||
let weak_bucket = Arc::downgrade(&arc_self);
|
||||
let refill_interval = arc_self.config.refill_interval;
|
||||
let refill_notifer = arc_self.refill_notifier.clone();
|
||||
let refill_task = tokio::spawn(async move {
|
||||
let mut interval = time::interval(refill_interval);
|
||||
loop {
|
||||
interval.tick().await;
|
||||
let Some(bucket) = weak_bucket.upgrade() else {
|
||||
break;
|
||||
};
|
||||
bucket.refill();
|
||||
refill_notifer.notify_waiters();
|
||||
}
|
||||
});
|
||||
|
||||
// Replace placeholder task with actual one
|
||||
arc_self
|
||||
.refill_task
|
||||
.lock()
|
||||
.unwrap()
|
||||
.replace(AbortOnDropHandle::new(refill_task));
|
||||
arc_self
|
||||
})
|
||||
}
|
||||
|
||||
/// Internal refill method (called only by background task)
|
||||
fn refill(&self) {
|
||||
let now_micros = self.elapsed_micros();
|
||||
let prev_time = self.last_refill_time.swap(now_micros, Ordering::Acquire);
|
||||
|
||||
// Calculate elapsed time in seconds
|
||||
let elapsed_secs = (now_micros.saturating_sub(prev_time)) as f64 / 1_000_000.0;
|
||||
|
||||
// Calculate tokens to add
|
||||
let tokens_to_add = (self.config.fill_rate as f64 * elapsed_secs) as u64;
|
||||
if tokens_to_add == 0 {
|
||||
/// Refill tokens based on elapsed time since last refill.
|
||||
/// Called while holding the bucket state lock.
|
||||
fn refill(&self, state: &mut BucketState, now: Instant) {
|
||||
let elapsed_nanos = now.saturating_duration_since(state.last_refill).as_nanos();
|
||||
if elapsed_nanos == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
// Add tokens without exceeding capacity
|
||||
let mut current = self.available_tokens.load(Ordering::Relaxed);
|
||||
loop {
|
||||
let new = current
|
||||
.saturating_add(tokens_to_add)
|
||||
.min(self.config.capacity);
|
||||
match self.available_tokens.compare_exchange_weak(
|
||||
current,
|
||||
new,
|
||||
Ordering::Release,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => break,
|
||||
Err(actual) => current = actual,
|
||||
}
|
||||
state.last_refill = now;
|
||||
if state.available_tokens == self.config.capacity {
|
||||
state.refill_remainder = 0;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
/// Calculate microseconds since bucket creation
|
||||
fn elapsed_micros(&self) -> u64 {
|
||||
self.start_time.elapsed().as_micros() as u64
|
||||
let generated = (self.config.fill_rate as u128)
|
||||
.saturating_mul(elapsed_nanos)
|
||||
.saturating_add(state.refill_remainder as u128);
|
||||
let tokens_to_add = generated / NANOS_PER_SECOND;
|
||||
let refill_remainder = generated % NANOS_PER_SECOND;
|
||||
let available_capacity = self.config.capacity - state.available_tokens;
|
||||
|
||||
if tokens_to_add >= available_capacity as u128 {
|
||||
state.available_tokens = self.config.capacity;
|
||||
state.refill_remainder = 0;
|
||||
} else {
|
||||
state.available_tokens += tokens_to_add as u64;
|
||||
state.refill_remainder = refill_remainder as u64;
|
||||
}
|
||||
}
|
||||
|
||||
/// Attempt to consume tokens without blocking
|
||||
@@ -165,44 +134,63 @@ impl TokenBucket {
|
||||
return false;
|
||||
}
|
||||
|
||||
let mut current = self.available_tokens.load(Ordering::Relaxed);
|
||||
loop {
|
||||
if current < tokens {
|
||||
return false;
|
||||
}
|
||||
let mut state = self.state.lock();
|
||||
self.refill(&mut state, Instant::now());
|
||||
if state.available_tokens < tokens {
|
||||
return false;
|
||||
}
|
||||
|
||||
let new = current - tokens;
|
||||
match self.available_tokens.compare_exchange_weak(
|
||||
current,
|
||||
new,
|
||||
Ordering::AcqRel,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => return true,
|
||||
Err(actual) => current = actual,
|
||||
state.available_tokens -= tokens;
|
||||
true
|
||||
}
|
||||
|
||||
/// Consume tokens, sleeping until they become available.
|
||||
pub async fn consume(&self, tokens: u64) {
|
||||
let mut remaining = tokens;
|
||||
while remaining > 0 {
|
||||
if self.stopped.load(Ordering::Acquire) {
|
||||
return;
|
||||
}
|
||||
let chunk = remaining.min(self.config.capacity);
|
||||
self.consume_chunk(chunk).await;
|
||||
remaining -= chunk;
|
||||
}
|
||||
}
|
||||
|
||||
/// Consume tokens, blocking if not available
|
||||
pub async fn consume(&self, tokens: u64) {
|
||||
async fn consume_chunk(&self, tokens: u64) {
|
||||
loop {
|
||||
let notified = self.refill_notifier.notified();
|
||||
if self.try_consume(tokens) {
|
||||
let stopped = self.stop_notifier.notified();
|
||||
if self.stopped.load(Ordering::Acquire) {
|
||||
return;
|
||||
}
|
||||
notified.await;
|
||||
|
||||
let sleep_dur = {
|
||||
let mut state = self.state.lock();
|
||||
self.refill(&mut state, Instant::now());
|
||||
if state.available_tokens >= tokens {
|
||||
state.available_tokens -= tokens;
|
||||
return;
|
||||
}
|
||||
|
||||
let deficit = tokens - state.available_tokens;
|
||||
let required = deficit as u128 * NANOS_PER_SECOND;
|
||||
let remaining = required - state.refill_remainder as u128;
|
||||
let sleep_nanos = remaining
|
||||
.div_ceil(self.config.fill_rate as u128)
|
||||
.min(u64::MAX as u128) as u64;
|
||||
Duration::from_nanos(sleep_nanos.max(1_000_000))
|
||||
};
|
||||
|
||||
tokio::select! {
|
||||
_ = time::sleep(sleep_dur) => {}
|
||||
_ = stopped => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn stop(&self) {
|
||||
self.stopped.store(true, Ordering::Release);
|
||||
self.refill_notifier.notify_waiters();
|
||||
let task = self.refill_task.lock().unwrap().take();
|
||||
if let Some(task) = task {
|
||||
task.abort();
|
||||
let _ = task.await;
|
||||
}
|
||||
self.stop_notifier.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -219,7 +207,7 @@ impl ByteLimiter for TokenBucket {
|
||||
|
||||
pub struct TokenBucketManager {
|
||||
buckets: Arc<DashMap<String, Arc<TokenBucket>>>,
|
||||
retain_task: Mutex<Option<AbortOnDropHandle<()>>>,
|
||||
retain_task: StdMutex<Option<AbortOnDropHandle<()>>>,
|
||||
}
|
||||
|
||||
impl Default for TokenBucketManager {
|
||||
@@ -252,7 +240,7 @@ impl TokenBucketManager {
|
||||
|
||||
Self {
|
||||
buckets,
|
||||
retain_task: Mutex::new(Some(AbortOnDropHandle::new(retain_task))),
|
||||
retain_task: StdMutex::new(Some(AbortOnDropHandle::new(retain_task))),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -285,12 +273,36 @@ impl TokenBucketManager {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::time::{Duration, sleep};
|
||||
use tokio::time::{Duration, sleep, timeout};
|
||||
|
||||
#[test]
|
||||
fn bucket_config_uses_one_second_default_capacity() {
|
||||
let config = BucketConfig::with_default_capacity(100_000);
|
||||
|
||||
assert_eq!(config.capacity, 100_000);
|
||||
assert_eq!(config.fill_rate, 100_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bucket_config_default_capacity_fits_regular_packets() {
|
||||
let config = BucketConfig::with_default_capacity(8_196);
|
||||
|
||||
assert_eq!(config.capacity, 8_196);
|
||||
assert_eq!(config.fill_rate, 8_196);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bucket_config_preserves_explicit_capacity() {
|
||||
let config = BucketConfig::new(200_000, 100_000);
|
||||
|
||||
assert_eq!(config.capacity, 200_000);
|
||||
assert_eq!(config.fill_rate, 100_000);
|
||||
}
|
||||
|
||||
/// Test initial state after creation
|
||||
#[tokio::test]
|
||||
async fn test_initial_state() {
|
||||
let bucket = TokenBucket::new(1000, 1000, Duration::from_millis(10));
|
||||
let bucket = TokenBucket::new(1000, 1000);
|
||||
|
||||
// Should have full capacity initially
|
||||
assert!(bucket.try_consume(1000));
|
||||
@@ -300,7 +312,7 @@ mod tests {
|
||||
/// Test token consumption behavior
|
||||
#[tokio::test]
|
||||
async fn test_consumption() {
|
||||
let bucket = TokenBucket::new(1500, 1000, Duration::from_millis(10));
|
||||
let bucket = TokenBucket::new(1500, 1000);
|
||||
|
||||
// First packet should succeed
|
||||
assert!(bucket.try_consume(1000));
|
||||
@@ -314,7 +326,7 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn stop_releases_waiting_consumers() {
|
||||
let bucket = TokenBucket::new(1, 1, Duration::from_secs(60));
|
||||
let bucket = TokenBucket::new(1, 1);
|
||||
assert!(bucket.try_consume(1));
|
||||
let waiting = tokio::spawn({
|
||||
let bucket = bucket.clone();
|
||||
@@ -332,27 +344,85 @@ mod tests {
|
||||
assert!(bucket.try_consume(u64::MAX));
|
||||
}
|
||||
|
||||
/// Test background refill functionality
|
||||
/// Test lazy refill functionality
|
||||
#[tokio::test]
|
||||
async fn test_refill() {
|
||||
let bucket = TokenBucket::new(1000, 1000, Duration::from_millis(10));
|
||||
let bucket = TokenBucket::new(1_000_000, 10_000);
|
||||
|
||||
// Drain the bucket
|
||||
assert!(bucket.try_consume(1000));
|
||||
assert!(!bucket.try_consume(1));
|
||||
assert!(bucket.try_consume(1_000_000));
|
||||
|
||||
// Wait for refill (1 refill interval + buffer)
|
||||
// Wait for time to pass (tokens accumulate lazily on next consume)
|
||||
sleep(Duration::from_millis(25)).await;
|
||||
let tokens = {
|
||||
let mut state = bucket.state.lock();
|
||||
bucket.refill(&mut state, Instant::now());
|
||||
state.available_tokens
|
||||
};
|
||||
assert!(tokens > 0, "Expected some refilled tokens");
|
||||
assert!(
|
||||
tokens < bucket.config.capacity,
|
||||
"Bucket unexpectedly refilled to capacity: {}",
|
||||
tokens
|
||||
);
|
||||
}
|
||||
|
||||
// Should have approximately 20 tokens (1000 tokens/s * 0.02s)
|
||||
assert!(bucket.try_consume(15));
|
||||
assert!(!bucket.try_consume(10)); // But not full capacity
|
||||
#[test]
|
||||
fn test_refill_preserves_fractional_tokens() {
|
||||
let bucket = TokenBucket::new(100, 100);
|
||||
let start = Instant::now();
|
||||
let mut state = bucket.state.lock();
|
||||
state.available_tokens = 0;
|
||||
state.last_refill = start;
|
||||
state.refill_remainder = 0;
|
||||
|
||||
for step in 1..=10 {
|
||||
bucket.refill(&mut state, start + Duration::from_millis(step * 15));
|
||||
}
|
||||
|
||||
assert_eq!(state.available_tokens, 15);
|
||||
assert_eq!(state.refill_remainder, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_refill_preserves_submicrosecond_time() {
|
||||
let bucket = TokenBucket::new(100, 1_000_000);
|
||||
let start = Instant::now();
|
||||
let mut state = bucket.state.lock();
|
||||
state.available_tokens = 0;
|
||||
state.last_refill = start;
|
||||
state.refill_remainder = 0;
|
||||
|
||||
for step in 1..=10 {
|
||||
bucket.refill(&mut state, start + Duration::from_nanos(step * 1_500));
|
||||
}
|
||||
|
||||
assert_eq!(state.available_tokens, 15);
|
||||
assert_eq!(state.refill_remainder, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_refill_discards_excess_credit_at_capacity() {
|
||||
let bucket = TokenBucket::new(10, 10);
|
||||
let start = Instant::now();
|
||||
let mut state = bucket.state.lock();
|
||||
state.available_tokens = 0;
|
||||
state.last_refill = start;
|
||||
|
||||
let refill_time = start + Duration::from_secs(2);
|
||||
bucket.refill(&mut state, refill_time);
|
||||
assert_eq!(state.available_tokens, 10);
|
||||
assert_eq!(state.refill_remainder, 0);
|
||||
|
||||
state.available_tokens = 0;
|
||||
bucket.refill(&mut state, refill_time);
|
||||
assert_eq!(state.available_tokens, 0);
|
||||
}
|
||||
|
||||
/// Test capacity enforcement
|
||||
#[tokio::test]
|
||||
async fn test_capacity_limit() {
|
||||
let bucket = TokenBucket::new(500, 1000, Duration::from_millis(10));
|
||||
let bucket = TokenBucket::new(500, 1000);
|
||||
|
||||
// Wait longer than refill interval
|
||||
sleep(Duration::from_millis(50)).await;
|
||||
@@ -363,39 +433,39 @@ mod tests {
|
||||
}
|
||||
|
||||
/// Test high load with concurrent access
|
||||
#[tokio::test]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn test_concurrent_access() {
|
||||
let bucket = TokenBucket::new(10_000, 1_000_000, Duration::from_millis(10));
|
||||
let bucket = TokenBucket::new(10_000, 1);
|
||||
let mut handles = vec![];
|
||||
|
||||
// Spawn 100 tasks to consume tokens concurrently
|
||||
for _ in 0..100 {
|
||||
let bucket = bucket.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
let mut consumed = 0;
|
||||
for _ in 0..100 {
|
||||
let _ = bucket.try_consume(10);
|
||||
if bucket.try_consume(10) {
|
||||
consumed += 10;
|
||||
}
|
||||
}
|
||||
consumed
|
||||
}));
|
||||
}
|
||||
|
||||
// Wait for all tasks to complete
|
||||
let mut consumed = 0;
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
consumed += handle.await.unwrap();
|
||||
}
|
||||
|
||||
// Verify we didn't exceed capacity
|
||||
let tokens_left = bucket.available_tokens.load(Ordering::Relaxed);
|
||||
assert!(
|
||||
tokens_left <= 10_000,
|
||||
"Tokens exceeded capacity: {}",
|
||||
tokens_left
|
||||
);
|
||||
assert_eq!(consumed, 10_000);
|
||||
assert_eq!(bucket.state.lock().available_tokens, 0);
|
||||
}
|
||||
|
||||
/// Test behavior when packet size exceeds capacity
|
||||
#[tokio::test]
|
||||
async fn test_oversized_packet() {
|
||||
let bucket = TokenBucket::new(1500, 1000, Duration::from_millis(10));
|
||||
let bucket = TokenBucket::new(1500, 1000);
|
||||
|
||||
// Packet larger than capacity should be rejected
|
||||
assert!(!bucket.try_consume(1600));
|
||||
@@ -404,23 +474,35 @@ mod tests {
|
||||
assert!(bucket.try_consume(1000));
|
||||
}
|
||||
|
||||
/// Test refill precision with small intervals
|
||||
#[tokio::test]
|
||||
async fn test_refill_precision() {
|
||||
let bucket = TokenBucket::new(10_000, 10_000, Duration::from_micros(100)); // 100μs interval
|
||||
async fn test_zero_fill_rate_is_normalized() {
|
||||
let bucket = TokenBucket::new(1000, 0);
|
||||
|
||||
// Drain most tokens
|
||||
assert!(bucket.try_consume(9900));
|
||||
assert_eq!(bucket.config.fill_rate, 1);
|
||||
}
|
||||
|
||||
// Wait for multiple refills
|
||||
sleep(Duration::from_millis(1)).await;
|
||||
#[tokio::test]
|
||||
async fn test_consume_oversized_packet_in_chunks() {
|
||||
let bucket = TokenBucket::new(10, 1_000_000);
|
||||
|
||||
// Should have accumulated about 100 tokens (10,000 tokens/s * 0.001s)
|
||||
let tokens = bucket.available_tokens.load(Ordering::Relaxed);
|
||||
assert!(
|
||||
(100..=200).contains(&tokens),
|
||||
"Unexpected token count: {}",
|
||||
tokens
|
||||
);
|
||||
timeout(Duration::from_millis(100), bucket.consume(25))
|
||||
.await
|
||||
.expect("oversized consume should be split into capacity-sized chunks");
|
||||
}
|
||||
|
||||
/// Test refill precision after elapsed time.
|
||||
#[test]
|
||||
fn test_refill_precision() {
|
||||
let bucket = TokenBucket::new(10_000, 10_000);
|
||||
let start = Instant::now();
|
||||
let mut state = bucket.state.lock();
|
||||
state.available_tokens = 100;
|
||||
state.last_refill = start;
|
||||
state.refill_remainder = 0;
|
||||
|
||||
bucket.refill(&mut state, start + Duration::from_micros(1_234));
|
||||
|
||||
assert_eq!(state.available_tokens, 112);
|
||||
assert_eq!(state.refill_remainder, 340_000_000);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -569,11 +569,7 @@ impl ZCPacket {
|
||||
ret.mut_payload()[foreign_network_hdr.get_header_len()..]
|
||||
.copy_from_slice(foreign_zc_packet.tunnel_payload());
|
||||
|
||||
let hdr = ret.mut_peer_manager_header().unwrap();
|
||||
hdr.from_peer_id = 0.into();
|
||||
hdr.to_peer_id = 0.into();
|
||||
hdr.packet_type = PacketType::ForeignNetworkPacket as u8;
|
||||
hdr.len.set(total_payload_len as u32);
|
||||
ret.fill_peer_manager_hdr(0, 0, PacketType::ForeignNetworkPacket as u8);
|
||||
|
||||
ret
|
||||
}
|
||||
@@ -689,6 +685,7 @@ impl ZCPacket {
|
||||
hdr.packet_type = packet_type;
|
||||
hdr.flags = 0;
|
||||
hdr.forward_counter = 1;
|
||||
hdr.reserved = 0;
|
||||
hdr.len.set(payload_len as u32);
|
||||
}
|
||||
|
||||
@@ -783,6 +780,11 @@ impl ZCPacket {
|
||||
}
|
||||
|
||||
pub fn foreign_network_inner_packet_type(&self) -> Option<u8> {
|
||||
self.foreign_network_inner_packet_info()
|
||||
.map(|(hdr, _)| hdr.packet_type)
|
||||
}
|
||||
|
||||
pub fn foreign_network_inner_packet_info(&self) -> Option<(&PeerManagerHeader, usize)> {
|
||||
if self.peer_manager_header()?.packet_type != PacketType::ForeignNetworkPacket as u8 {
|
||||
return None;
|
||||
}
|
||||
@@ -790,7 +792,9 @@ impl ZCPacket {
|
||||
let payload = self.payload();
|
||||
let hdr = ForeignNetworkPacketHeader::ref_from_prefix(payload)?;
|
||||
let inner_packet = payload.get(hdr.get_header_len()..)?;
|
||||
PeerManagerHeader::ref_from_prefix(inner_packet).map(|hdr| hdr.packet_type)
|
||||
let peer_manager_header = PeerManagerHeader::ref_from_prefix(inner_packet)?;
|
||||
let payload_len = inner_packet.len() - PEER_MANAGER_HEADER_SIZE;
|
||||
Some((peer_manager_header, payload_len))
|
||||
}
|
||||
|
||||
pub fn foreign_network_packet(mut self) -> Self {
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -250,14 +250,6 @@ message PortForwardConfigPb {
|
||||
|
||||
message ProxyDstInfo { SocketAddr dst_addr = 1; }
|
||||
|
||||
message LimiterConfig {
|
||||
optional uint64 burst_rate =
|
||||
1; // default 1 means no burst (capacity is same with bps)
|
||||
optional uint64 bps = 2; // default 0 means no limit (unit is B/s)
|
||||
optional uint64 fill_duration_ms =
|
||||
3; // default 10ms, the period to fill the bucket
|
||||
}
|
||||
|
||||
message SecureModeConfig {
|
||||
bool enabled = 1;
|
||||
|
||||
|
||||
@@ -2296,13 +2296,7 @@ pub async fn relay_bps_limit_test(#[values(100, 200, 400, 800)] bps_limit: u64)
|
||||
println!("bps: {}", bps);
|
||||
|
||||
let bps = bps as u64 / 1024;
|
||||
// allow 50kb jitter
|
||||
assert!(
|
||||
bps >= bps_limit - 50 && bps <= bps_limit + 50,
|
||||
"bps: {}, bps_limit: {}",
|
||||
bps,
|
||||
bps_limit
|
||||
);
|
||||
assert_limited_payload_bps(bps, bps_limit);
|
||||
|
||||
drop_insts(insts).await;
|
||||
}
|
||||
@@ -2339,16 +2333,25 @@ pub async fn instance_recv_bps_limit_test(#[values(100, 800)] bps_limit: u64) {
|
||||
println!("bps: {}", bps);
|
||||
|
||||
let bps = bps as u64 / 1024;
|
||||
assert!(
|
||||
bps >= bps_limit - 50 && bps <= bps_limit + 50,
|
||||
"bps: {}, bps_limit: {}",
|
||||
bps,
|
||||
bps_limit
|
||||
);
|
||||
assert_limited_payload_bps(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;
|
||||
assert!(
|
||||
bps >= min_bps && bps <= max_bps,
|
||||
"bps: {}, expected: {}..={}",
|
||||
bps,
|
||||
min_bps,
|
||||
max_bps
|
||||
);
|
||||
}
|
||||
|
||||
async fn assert_peer_admission_blocked(inst: &Instance, url: url::Url) {
|
||||
let ip = url
|
||||
.host_str()
|
||||
|
||||
Reference in New Issue
Block a user