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:
@@ -1,13 +1,14 @@
|
|||||||
use atomic_shim::AtomicU64;
|
|
||||||
use dashmap::DashMap;
|
use dashmap::DashMap;
|
||||||
use std::sync::atomic::{AtomicBool, Ordering};
|
use parking_lot::Mutex;
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{
|
||||||
|
Arc, Mutex as StdMutex,
|
||||||
|
atomic::{AtomicBool, Ordering},
|
||||||
|
};
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
use tokio::sync::Notify;
|
use tokio::sync::Notify;
|
||||||
use tokio_util::task::AbortOnDropHandle;
|
use tokio_util::task::AbortOnDropHandle;
|
||||||
|
|
||||||
use crate::foundation::time;
|
use crate::foundation::time;
|
||||||
use easytier_proto::common::LimiterConfig;
|
|
||||||
|
|
||||||
#[async_trait::async_trait]
|
#[async_trait::async_trait]
|
||||||
pub(crate) trait ByteLimiter: Send + Sync {
|
pub(crate) trait ByteLimiter: Send + Sync {
|
||||||
@@ -27,49 +28,46 @@ impl ByteLimiter for () {
|
|||||||
|
|
||||||
pub(crate) type ArcByteLimiter = Arc<dyn ByteLimiter>;
|
pub(crate) type ArcByteLimiter = Arc<dyn ByteLimiter>;
|
||||||
|
|
||||||
/// Token Bucket rate limiter using atomic operations
|
const MIN_FILL_RATE: u64 = 8_196;
|
||||||
pub struct TokenBucket {
|
const NANOS_PER_SECOND: u128 = 1_000_000_000;
|
||||||
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
|
|
||||||
|
|
||||||
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,
|
stopped: AtomicBool,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct BucketState {
|
||||||
|
available_tokens: u64,
|
||||||
|
last_refill: Instant,
|
||||||
|
refill_remainder: u64,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Clone, Copy)]
|
#[derive(Clone, Copy)]
|
||||||
pub struct BucketConfig {
|
pub struct BucketConfig {
|
||||||
capacity: u64, // Maximum token capacity
|
capacity: u64,
|
||||||
fill_rate: u64, // Tokens added per second
|
fill_rate: u64,
|
||||||
refill_interval: Duration, // Time between refill operations
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl From<LimiterConfig> for BucketConfig {
|
impl BucketConfig {
|
||||||
fn from(cfg: LimiterConfig) -> Self {
|
pub fn new(capacity: u64, fill_rate: u64) -> Self {
|
||||||
let burst_rate = 1.max(cfg.burst_rate.unwrap_or(1));
|
Self {
|
||||||
let fill_rate = 8196.max(cfg.bps.unwrap_or(u64::MAX / burst_rate));
|
capacity,
|
||||||
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,
|
|
||||||
fill_rate,
|
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 {
|
impl TokenBucket {
|
||||||
pub fn new(capacity: u64, bps: u64, refill_interval: Duration) -> Arc<Self> {
|
pub fn new(capacity: u64, bps: u64) -> Arc<Self> {
|
||||||
let config = BucketConfig {
|
Self::new_from_cfg(BucketConfig::new(capacity, bps))
|
||||||
capacity,
|
|
||||||
fill_rate: bps,
|
|
||||||
refill_interval,
|
|
||||||
};
|
|
||||||
Self::new_from_cfg(config)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Creates a new Token Bucket rate limiter
|
/// Creates a new Token Bucket rate limiter
|
||||||
@@ -77,79 +75,50 @@ impl TokenBucket {
|
|||||||
/// # Arguments
|
/// # Arguments
|
||||||
/// * `capacity` - Bucket capacity in bytes
|
/// * `capacity` - Bucket capacity in bytes
|
||||||
/// * `bps` - Bandwidth limit in bytes per second
|
/// * `bps` - Bandwidth limit in bytes per second
|
||||||
/// * `refill_interval` - Refill interval (recommended 10-50ms)
|
pub fn new_from_cfg(mut config: BucketConfig) -> Arc<Self> {
|
||||||
pub fn new_from_cfg(config: BucketConfig) -> Arc<Self> {
|
config.capacity = config.capacity.max(1);
|
||||||
// Create Arc instance with placeholder task
|
config.fill_rate = config.fill_rate.max(1);
|
||||||
let arc_self = Arc::new(Self {
|
|
||||||
available_tokens: AtomicU64::new(config.capacity),
|
Arc::new(Self {
|
||||||
last_refill_time: AtomicU64::new(0),
|
state: Mutex::new(BucketState {
|
||||||
|
available_tokens: config.capacity,
|
||||||
|
last_refill: Instant::now(),
|
||||||
|
refill_remainder: 0,
|
||||||
|
}),
|
||||||
config,
|
config,
|
||||||
refill_task: Mutex::new(None),
|
stop_notifier: Notify::new(),
|
||||||
start_time: std::time::Instant::now(),
|
|
||||||
refill_notifier: Arc::new(Notify::new()),
|
|
||||||
stopped: AtomicBool::new(false),
|
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)
|
/// Refill tokens based on elapsed time since last refill.
|
||||||
fn refill(&self) {
|
/// Called while holding the bucket state lock.
|
||||||
let now_micros = self.elapsed_micros();
|
fn refill(&self, state: &mut BucketState, now: Instant) {
|
||||||
let prev_time = self.last_refill_time.swap(now_micros, Ordering::Acquire);
|
let elapsed_nanos = now.saturating_duration_since(state.last_refill).as_nanos();
|
||||||
|
if elapsed_nanos == 0 {
|
||||||
// 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 {
|
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Add tokens without exceeding capacity
|
state.last_refill = now;
|
||||||
let mut current = self.available_tokens.load(Ordering::Relaxed);
|
if state.available_tokens == self.config.capacity {
|
||||||
loop {
|
state.refill_remainder = 0;
|
||||||
let new = current
|
return;
|
||||||
.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,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Calculate microseconds since bucket creation
|
let generated = (self.config.fill_rate as u128)
|
||||||
fn elapsed_micros(&self) -> u64 {
|
.saturating_mul(elapsed_nanos)
|
||||||
self.start_time.elapsed().as_micros() as u64
|
.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
|
/// Attempt to consume tokens without blocking
|
||||||
@@ -165,44 +134,63 @@ impl TokenBucket {
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut current = self.available_tokens.load(Ordering::Relaxed);
|
let mut state = self.state.lock();
|
||||||
loop {
|
self.refill(&mut state, Instant::now());
|
||||||
if current < tokens {
|
if state.available_tokens < tokens {
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
let new = current - tokens;
|
state.available_tokens -= tokens;
|
||||||
match self.available_tokens.compare_exchange_weak(
|
true
|
||||||
current,
|
|
||||||
new,
|
|
||||||
Ordering::AcqRel,
|
|
||||||
Ordering::Relaxed,
|
|
||||||
) {
|
|
||||||
Ok(_) => return true,
|
|
||||||
Err(actual) => current = actual,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 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
|
async fn consume_chunk(&self, tokens: u64) {
|
||||||
pub async fn consume(&self, tokens: u64) {
|
|
||||||
loop {
|
loop {
|
||||||
let notified = self.refill_notifier.notified();
|
let stopped = self.stop_notifier.notified();
|
||||||
if self.try_consume(tokens) {
|
if self.stopped.load(Ordering::Acquire) {
|
||||||
return;
|
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) {
|
async fn stop(&self) {
|
||||||
self.stopped.store(true, Ordering::Release);
|
self.stopped.store(true, Ordering::Release);
|
||||||
self.refill_notifier.notify_waiters();
|
self.stop_notifier.notify_waiters();
|
||||||
let task = self.refill_task.lock().unwrap().take();
|
|
||||||
if let Some(task) = task {
|
|
||||||
task.abort();
|
|
||||||
let _ = task.await;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -219,7 +207,7 @@ impl ByteLimiter for TokenBucket {
|
|||||||
|
|
||||||
pub struct TokenBucketManager {
|
pub struct TokenBucketManager {
|
||||||
buckets: Arc<DashMap<String, Arc<TokenBucket>>>,
|
buckets: Arc<DashMap<String, Arc<TokenBucket>>>,
|
||||||
retain_task: Mutex<Option<AbortOnDropHandle<()>>>,
|
retain_task: StdMutex<Option<AbortOnDropHandle<()>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for TokenBucketManager {
|
impl Default for TokenBucketManager {
|
||||||
@@ -252,7 +240,7 @@ impl TokenBucketManager {
|
|||||||
|
|
||||||
Self {
|
Self {
|
||||||
buckets,
|
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)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
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
|
/// Test initial state after creation
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_initial_state() {
|
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
|
// Should have full capacity initially
|
||||||
assert!(bucket.try_consume(1000));
|
assert!(bucket.try_consume(1000));
|
||||||
@@ -300,7 +312,7 @@ mod tests {
|
|||||||
/// Test token consumption behavior
|
/// Test token consumption behavior
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_consumption() {
|
async fn test_consumption() {
|
||||||
let bucket = TokenBucket::new(1500, 1000, Duration::from_millis(10));
|
let bucket = TokenBucket::new(1500, 1000);
|
||||||
|
|
||||||
// First packet should succeed
|
// First packet should succeed
|
||||||
assert!(bucket.try_consume(1000));
|
assert!(bucket.try_consume(1000));
|
||||||
@@ -314,7 +326,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn stop_releases_waiting_consumers() {
|
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));
|
assert!(bucket.try_consume(1));
|
||||||
let waiting = tokio::spawn({
|
let waiting = tokio::spawn({
|
||||||
let bucket = bucket.clone();
|
let bucket = bucket.clone();
|
||||||
@@ -332,27 +344,85 @@ mod tests {
|
|||||||
assert!(bucket.try_consume(u64::MAX));
|
assert!(bucket.try_consume(u64::MAX));
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Test background refill functionality
|
/// Test lazy refill functionality
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_refill() {
|
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
|
// Drain the bucket
|
||||||
assert!(bucket.try_consume(1000));
|
assert!(bucket.try_consume(1_000_000));
|
||||||
assert!(!bucket.try_consume(1));
|
|
||||||
|
|
||||||
// Wait for refill (1 refill interval + buffer)
|
// Wait for time to pass (tokens accumulate lazily on next consume)
|
||||||
sleep(Duration::from_millis(25)).await;
|
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)
|
#[test]
|
||||||
assert!(bucket.try_consume(15));
|
fn test_refill_preserves_fractional_tokens() {
|
||||||
assert!(!bucket.try_consume(10)); // But not full capacity
|
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
|
/// Test capacity enforcement
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_capacity_limit() {
|
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
|
// Wait longer than refill interval
|
||||||
sleep(Duration::from_millis(50)).await;
|
sleep(Duration::from_millis(50)).await;
|
||||||
@@ -363,39 +433,39 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Test high load with concurrent access
|
/// Test high load with concurrent access
|
||||||
#[tokio::test]
|
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||||
async fn test_concurrent_access() {
|
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![];
|
let mut handles = vec![];
|
||||||
|
|
||||||
// Spawn 100 tasks to consume tokens concurrently
|
// Spawn 100 tasks to consume tokens concurrently
|
||||||
for _ in 0..100 {
|
for _ in 0..100 {
|
||||||
let bucket = bucket.clone();
|
let bucket = bucket.clone();
|
||||||
handles.push(tokio::spawn(async move {
|
handles.push(tokio::spawn(async move {
|
||||||
|
let mut consumed = 0;
|
||||||
for _ in 0..100 {
|
for _ in 0..100 {
|
||||||
let _ = bucket.try_consume(10);
|
if bucket.try_consume(10) {
|
||||||
|
consumed += 10;
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
consumed
|
||||||
}));
|
}));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Wait for all tasks to complete
|
// Wait for all tasks to complete
|
||||||
|
let mut consumed = 0;
|
||||||
for handle in handles {
|
for handle in handles {
|
||||||
handle.await.unwrap();
|
consumed += handle.await.unwrap();
|
||||||
}
|
}
|
||||||
|
|
||||||
// Verify we didn't exceed capacity
|
assert_eq!(consumed, 10_000);
|
||||||
let tokens_left = bucket.available_tokens.load(Ordering::Relaxed);
|
assert_eq!(bucket.state.lock().available_tokens, 0);
|
||||||
assert!(
|
|
||||||
tokens_left <= 10_000,
|
|
||||||
"Tokens exceeded capacity: {}",
|
|
||||||
tokens_left
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Test behavior when packet size exceeds capacity
|
/// Test behavior when packet size exceeds capacity
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_oversized_packet() {
|
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
|
// Packet larger than capacity should be rejected
|
||||||
assert!(!bucket.try_consume(1600));
|
assert!(!bucket.try_consume(1600));
|
||||||
@@ -404,23 +474,35 @@ mod tests {
|
|||||||
assert!(bucket.try_consume(1000));
|
assert!(bucket.try_consume(1000));
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Test refill precision with small intervals
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_refill_precision() {
|
async fn test_zero_fill_rate_is_normalized() {
|
||||||
let bucket = TokenBucket::new(10_000, 10_000, Duration::from_micros(100)); // 100μs interval
|
let bucket = TokenBucket::new(1000, 0);
|
||||||
|
|
||||||
// Drain most tokens
|
assert_eq!(bucket.config.fill_rate, 1);
|
||||||
assert!(bucket.try_consume(9900));
|
}
|
||||||
|
|
||||||
// Wait for multiple refills
|
#[tokio::test]
|
||||||
sleep(Duration::from_millis(1)).await;
|
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)
|
timeout(Duration::from_millis(100), bucket.consume(25))
|
||||||
let tokens = bucket.available_tokens.load(Ordering::Relaxed);
|
.await
|
||||||
assert!(
|
.expect("oversized consume should be split into capacity-sized chunks");
|
||||||
(100..=200).contains(&tokens),
|
}
|
||||||
"Unexpected token count: {}",
|
|
||||||
tokens
|
/// 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()..]
|
ret.mut_payload()[foreign_network_hdr.get_header_len()..]
|
||||||
.copy_from_slice(foreign_zc_packet.tunnel_payload());
|
.copy_from_slice(foreign_zc_packet.tunnel_payload());
|
||||||
|
|
||||||
let hdr = ret.mut_peer_manager_header().unwrap();
|
ret.fill_peer_manager_hdr(0, 0, PacketType::ForeignNetworkPacket as u8);
|
||||||
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
|
ret
|
||||||
}
|
}
|
||||||
@@ -689,6 +685,7 @@ impl ZCPacket {
|
|||||||
hdr.packet_type = packet_type;
|
hdr.packet_type = packet_type;
|
||||||
hdr.flags = 0;
|
hdr.flags = 0;
|
||||||
hdr.forward_counter = 1;
|
hdr.forward_counter = 1;
|
||||||
|
hdr.reserved = 0;
|
||||||
hdr.len.set(payload_len as u32);
|
hdr.len.set(payload_len as u32);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -783,6 +780,11 @@ impl ZCPacket {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn foreign_network_inner_packet_type(&self) -> Option<u8> {
|
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 {
|
if self.peer_manager_header()?.packet_type != PacketType::ForeignNetworkPacket as u8 {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
@@ -790,7 +792,9 @@ impl ZCPacket {
|
|||||||
let payload = self.payload();
|
let payload = self.payload();
|
||||||
let hdr = ForeignNetworkPacketHeader::ref_from_prefix(payload)?;
|
let hdr = ForeignNetworkPacketHeader::ref_from_prefix(payload)?;
|
||||||
let inner_packet = payload.get(hdr.get_header_len()..)?;
|
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 {
|
pub fn foreign_network_packet(mut self) -> Self {
|
||||||
|
|||||||
@@ -786,11 +786,7 @@ impl AclProcessor {
|
|||||||
panic!("Rate limit bucket not found");
|
panic!("Rate limit bucket not found");
|
||||||
}
|
}
|
||||||
RateLimitValue {
|
RateLimitValue {
|
||||||
token_bucket: TokenBucket::new(
|
token_bucket: TokenBucket::new(burst as u64, rate as u64),
|
||||||
burst as u64,
|
|
||||||
rate as u64,
|
|
||||||
Duration::from_millis(10),
|
|
||||||
),
|
|
||||||
last_update: Instant::now(),
|
last_update: Instant::now(),
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -36,6 +36,7 @@ use crate::peers::{
|
|||||||
PacketRecvChan,
|
PacketRecvChan,
|
||||||
context::{ArcPeerContext, NetworkIdentity, NetworkSecretDigest},
|
context::{ArcPeerContext, NetworkIdentity, NetworkSecretDigest},
|
||||||
send_packet_to_chan,
|
send_packet_to_chan,
|
||||||
|
traffic_metrics::data_packet_payload_len,
|
||||||
};
|
};
|
||||||
use crate::{
|
use crate::{
|
||||||
config::PeerId,
|
config::PeerId,
|
||||||
@@ -1290,6 +1291,7 @@ impl PeerConn {
|
|||||||
|
|
||||||
let mut zc_packet = ret.unwrap();
|
let mut zc_packet = ret.unwrap();
|
||||||
let buf_len = zc_packet.buf_len() as u64;
|
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 {
|
let Some(peer_mgr_hdr) = zc_packet.mut_peer_manager_header() else {
|
||||||
tracing::error!(
|
tracing::error!(
|
||||||
"unexpected packet: {:?}, cannot decode peer manager hdr",
|
"unexpected packet: {:?}, cannot decode peer manager hdr",
|
||||||
@@ -1315,8 +1317,10 @@ impl PeerConn {
|
|||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some(limiter) = recv_limiter.as_ref() {
|
if let Some(payload_len) = limited_payload_len
|
||||||
limiter.consume(buf_len).await;
|
&& 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 dashmap::DashMap;
|
||||||
use easytier_proto::{
|
use easytier_proto::{
|
||||||
acl::Acl,
|
acl::Acl,
|
||||||
common::{
|
common::{FlagsInConfig, PeerFeatureFlag, SecureModeConfig, StunInfo, TunnelInfo},
|
||||||
FlagsInConfig, LimiterConfig, PeerFeatureFlag, SecureModeConfig, StunInfo, TunnelInfo,
|
|
||||||
},
|
|
||||||
peer_rpc::{PeerGroupInfo, TrustedCredentialPubkeyProof},
|
peer_rpc::{PeerGroupInfo, TrustedCredentialPubkeyProof},
|
||||||
};
|
};
|
||||||
use hmac::{Hmac, Mac};
|
use hmac::{Hmac, Mac};
|
||||||
@@ -31,7 +29,7 @@ use crate::{
|
|||||||
},
|
},
|
||||||
events::{CoreEvent, CoreEventSink},
|
events::{CoreEvent, CoreEventSink},
|
||||||
foundation::stats::{LabelSet, LabelType, MetricName, StatsManager},
|
foundation::stats::{LabelSet, LabelType, MetricName, StatsManager},
|
||||||
foundation::token_bucket::{ArcByteLimiter, TokenBucketManager},
|
foundation::token_bucket::{ArcByteLimiter, BucketConfig, TokenBucketManager},
|
||||||
peers::{
|
peers::{
|
||||||
credential_manager::{CredentialManager, CredentialStorage},
|
credential_manager::{CredentialManager, CredentialStorage},
|
||||||
util::shrink_dashmap,
|
util::shrink_dashmap,
|
||||||
@@ -362,17 +360,7 @@ impl CorePeerContext {
|
|||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
let manager = state.manager.get_or_insert_with(TokenBucketManager::new);
|
let manager = state.manager.get_or_insert_with(TokenBucketManager::new);
|
||||||
Some(
|
Some(manager.get_or_create(key, BucketConfig::with_default_capacity(bps)))
|
||||||
manager.get_or_create(
|
|
||||||
key,
|
|
||||||
LimiterConfig {
|
|
||||||
burst_rate: None,
|
|
||||||
bps: Some(bps),
|
|
||||||
fill_duration_ms: None,
|
|
||||||
}
|
|
||||||
.into(),
|
|
||||||
),
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn stop(&self) {
|
pub(crate) async fn stop(&self) {
|
||||||
|
|||||||
@@ -51,7 +51,8 @@ use super::{
|
|||||||
relay_peer_map::RelayPeerMap,
|
relay_peer_map::RelayPeerMap,
|
||||||
route::{NextHopPolicy, Route, RouteInterface, peer_ospf_route::PeerRoute},
|
route::{NextHopPolicy, Route, RouteInterface, peer_ospf_route::PeerRoute},
|
||||||
traffic_metrics::{
|
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,
|
util::shrink_dashmap,
|
||||||
whitelist::check_network_in_relay_whitelist,
|
whitelist::check_network_in_relay_whitelist,
|
||||||
@@ -1347,8 +1348,9 @@ impl ForeignNetworkPacketRouter {
|
|||||||
);
|
);
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
if let Some(bps_limiter) = bps_limiter.as_ref()
|
if let Some(payload_len) = data_packet_payload_len(&zc_packet)
|
||||||
&& !bps_limiter.try_consume(len.into())
|
&& let Some(bps_limiter) = bps_limiter.as_ref()
|
||||||
|
&& !bps_limiter.try_consume(payload_len)
|
||||||
{
|
{
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ use futures::future::BoxFuture;
|
|||||||
|
|
||||||
use crate::config::PeerId;
|
use crate::config::PeerId;
|
||||||
use crate::foundation::stats::{CounterHandle, LabelSet, LabelType, MetricName, StatsManager};
|
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::peers::util::shrink_dashmap;
|
||||||
use crate::proto::peer_rpc::RoutePeerInfo;
|
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
|
|| 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)]
|
#[derive(Clone)]
|
||||||
struct TrafficMetricGroup {
|
struct TrafficMetricGroup {
|
||||||
data: Arc<LogicalTrafficMetrics>,
|
data: Arc<LogicalTrafficMetrics>,
|
||||||
@@ -327,6 +347,86 @@ mod tests {
|
|||||||
.with_label_type(LabelType::ToInstanceId(instance_id.to_string()))
|
.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]
|
#[tokio::test]
|
||||||
async fn logical_traffic_metrics_upgrade_unknown_instance_label() {
|
async fn logical_traffic_metrics_upgrade_unknown_instance_label() {
|
||||||
let stats_mgr = Arc::new(StatsManager::new());
|
let stats_mgr = Arc::new(StatsManager::new());
|
||||||
|
|||||||
@@ -250,14 +250,6 @@ message PortForwardConfigPb {
|
|||||||
|
|
||||||
message ProxyDstInfo { SocketAddr dst_addr = 1; }
|
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 {
|
message SecureModeConfig {
|
||||||
bool enabled = 1;
|
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);
|
println!("bps: {}", bps);
|
||||||
|
|
||||||
let bps = bps as u64 / 1024;
|
let bps = bps as u64 / 1024;
|
||||||
// allow 50kb jitter
|
assert_limited_payload_bps(bps, bps_limit);
|
||||||
assert!(
|
|
||||||
bps >= bps_limit - 50 && bps <= bps_limit + 50,
|
|
||||||
"bps: {}, bps_limit: {}",
|
|
||||||
bps,
|
|
||||||
bps_limit
|
|
||||||
);
|
|
||||||
|
|
||||||
drop_insts(insts).await;
|
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);
|
println!("bps: {}", bps);
|
||||||
|
|
||||||
let bps = bps as u64 / 1024;
|
let bps = bps as u64 / 1024;
|
||||||
assert!(
|
assert_limited_payload_bps(bps, bps_limit);
|
||||||
bps >= bps_limit - 50 && bps <= bps_limit + 50,
|
|
||||||
"bps: {}, bps_limit: {}",
|
|
||||||
bps,
|
|
||||||
bps_limit
|
|
||||||
);
|
|
||||||
|
|
||||||
drop_insts(insts).await;
|
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) {
|
async fn assert_peer_admission_blocked(inst: &Instance, url: url::Url) {
|
||||||
let ip = url
|
let ip = url
|
||||||
.host_str()
|
.host_str()
|
||||||
|
|||||||
Reference in New Issue
Block a user