Zynh0722/Riven
1
0
Fork 0
forked from mirror/Riven
Code Pull requests Activity

Update FakeInstant, token bucket test

Browse source
This commit is contained in:
Mingwei Samuel 2023-04-29 12:43:38 -07:00
parent 112f6bceae
commit 1567bae16e
2 changed files with 79 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
riven/Cargo.toml
View file

@ -55,7 +55,7 @@ tracing = { version = "0.1", optional = true }
[dev-dependencies]
colored = "2"
env_logger = "0.9"
fake_instant = "0.4"
env_logger = "0.10.0"
fake_instant = "0.5.0"
futures = "0.3"
tokio = { version = "1", default-features = false, features = [ "rt-multi-thread" ] }

103
riven/src/req/token_bucket.test.rs
View file

@ -1,8 +1,8 @@
#![cfg(test)]
use fake_instant::FakeClock as Instant;
use fake_instant::FakeInstant as Instant;
/// This is a hack to test token bucket, substituting FakeClock for Instant.
/// This is a hack to test token bucket, substituting `FakeInstant` in place of `Instant`.
mod token_bucket {
include!("token_bucket.rs");
@ -11,13 +11,13 @@ mod token_bucket {
use lazy_static::lazy_static;
lazy_static! {
pub static ref D00: Duration = Duration::new(0, 0);
pub static ref ZERO: Duration = Duration::new(0, 0);
}
#[test]
fn test_basic() {
Instant::set_time(50_000);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.95, 1.0);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *ZERO, 0.95, 1.0);
assert!(bucket.get_tokens(50), "Should have not violated limit.");
assert_eq!(None, bucket.get_delay(), "Can get stuff.");
assert!(!bucket.get_tokens(51), "Should have violated limit.");
@ -25,36 +25,45 @@ mod token_bucket {
#[test]
fn test_internal_constructor() {
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 1.0, 1.0);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *ZERO, 1.0, 1.0);
assert_eq!(100, bucket.burst_limit);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 1e-6, 1.0);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *ZERO, 1e-6, 1.0);
assert_eq!(1, bucket.burst_limit);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 1.0, 1e-6);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *ZERO, 1.0, 1e-6);
assert_eq!(1, bucket.total_limit);
assert_eq!(1, bucket.burst_limit);
}
#[test]
fn test_saturated_100_burst() {
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 1.00, 1.0);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *ZERO, 1.00, 1.0);
Instant::set_time(50_000);
assert!(bucket.get_tokens(100), "All tokens should be immediately available.");
assert!(
bucket.get_tokens(100),
"All tokens should be immediately available."
);
assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
Instant::advance_time(1001); // Extra buffer for Duration(0).
assert!(bucket.get_tokens(100), "All tokens should be available after a bucket duration.");
assert!(
bucket.get_tokens(100),
"All tokens should be available after a bucket duration."
);
assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
}
#[test]
fn test_saturated_95_burst() {
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.95, 1.0);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *ZERO, 0.95, 1.0);
Instant::set_time(50_000);
assert!(bucket.get_tokens(95), "95 tokens should be immediately available.");
assert!(
bucket.get_tokens(95),
"95 tokens should be immediately available."
);
assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
Instant::advance_time(475);
@ -75,7 +84,7 @@ mod token_bucket {
#[test]
fn test_violated_50_burst() {
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.50, 1.0);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *ZERO, 0.50, 1.0);
Instant::set_time(50_000);
assert!(!bucket.get_tokens(90), "Burst should be violated.");
@ -84,18 +93,27 @@ mod token_bucket {
#[test]
fn test_saturated_50_burst() {
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.50, 1.0);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *ZERO, 0.50, 1.0);
Instant::set_time(50_000);
assert!(bucket.get_tokens(50), "Half the tokens should be immediately available.");
assert!(
bucket.get_tokens(50),
"Half the tokens should be immediately available."
);
assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
Instant::advance_time(501); // Extra buffer for Duration(0).
assert!(bucket.get_tokens(50), "Half the tokens should be available after a half bucket duration.");
assert!(
bucket.get_tokens(50),
"Half the tokens should be available after a half bucket duration."
);
assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
Instant::advance_time(501);
assert!(bucket.get_tokens(50), "Half the tokens should be available after a full bucket duration.");
assert!(
bucket.get_tokens(50),
"Half the tokens should be available after a full bucket duration."
);
assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
Instant::advance_time(501);
@ -106,10 +124,19 @@ mod token_bucket {
#[test]
fn test_saturated_90_burst_rate_usage_factor_50() {
let rate_usage_factor = 0.5;
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.90, rate_usage_factor);
let bucket = VectorTokenBucket::new(
Duration::from_millis(1000),
100,
*ZERO,
0.90,
rate_usage_factor,
);
Instant::set_time(50_000);
assert!(bucket.get_tokens(45), "45 tokens should be immediately available.");
assert!(
bucket.get_tokens(45),
"45 tokens should be immediately available."
);
assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
Instant::advance_time(475);
@ -131,18 +158,42 @@ mod token_bucket {
#[test]
fn test_many() {
Instant::set_time(50_000);
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.5, 1.0);
assert!(bucket.get_tokens(50), "Should have not violated limit. i=-1.");
let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *ZERO, 0.5, 1.0);
assert!(
bucket.get_tokens(50),
"Should have not violated limit. i=-1."
);
assert_ne!(None, bucket.get_delay(), "Bucket should have delay. i=-1.");
for i in 0..20_000 {
Instant::advance_time(501);
assert!(bucket.get_tokens(50), "Should have not violated limit. i={}.", i);
assert_ne!(None, bucket.get_delay(), "Bucket should have delay. i={}.", i);
assert!(
bucket.get_tokens(50),
"Should have not violated limit. i={}.",
i
);
assert_ne!(
None,
bucket.get_delay(),
"Bucket should have delay. i={}.",
i
);
Instant::advance_time(501);
assert!(bucket.get_tokens(50), "Should have not violated limit. i={}.", i);
assert_ne!(None, bucket.get_delay(), "Bucket should have delay. i={}.", i);
assert!(
bucket.get_tokens(50),
"Should have not violated limit. i={}.",
i
);
assert_ne!(
None,
bucket.get_delay(),
"Bucket should have delay. i={}.",
i
);
}
assert!(bucket.timestamps.lock().len() < 110, "Should not memory leak.");
assert!(
bucket.timestamps.lock().len() < 110,
"Should not memory leak."
);
}
}
}