From 1567bae16e81e3fb80c7887fa7b83e883c31e98d Mon Sep 17 00:00:00 2001 From: Mingwei Samuel Date: Sat, 29 Apr 2023 12:43:38 -0700 Subject: [PATCH] Update `FakeInstant`, token bucket test --- riven/Cargo.toml | 4 +- riven/src/req/token_bucket.test.rs | 103 +++++++++++++++++++++-------- 2 files changed, 79 insertions(+), 28 deletions(-) diff --git a/riven/Cargo.toml b/riven/Cargo.toml index f46a9de..3944fb9 100644 --- a/riven/Cargo.toml +++ b/riven/Cargo.toml @@ -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" ] } diff --git a/riven/src/req/token_bucket.test.rs b/riven/src/req/token_bucket.test.rs index e8d83cf..65aaa32 100644 --- a/riven/src/req/token_bucket.test.rs +++ b/riven/src/req/token_bucket.test.rs @@ -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." + ); } } }