From 80eeb184a073db7de1480c997b64b4b9c0e00f2d Mon Sep 17 00:00:00 2001 From: Mingwei Samuel Date: Fri, 23 Jul 2021 18:03:55 -0700 Subject: [PATCH] Fix rounding bug in token_bucket for ~1 or 2 second buckets --- src/req/token_bucket.rs | 20 ++++++++++--- src/req/token_bucket.test.rs | 54 +++++++++++++++++------------------- tests/tests_jp.rs | 21 +++++++------- 3 files changed, 52 insertions(+), 43 deletions(-) diff --git a/src/req/token_bucket.rs b/src/req/token_bucket.rs index 9f40989..a37f4ad 100644 --- a/src/req/token_bucket.rs +++ b/src/req/token_bucket.rs @@ -68,9 +68,7 @@ impl VectorTokenBucket { // Effective duration. let d_eff = duration + duration_overhead; - let burst_duration = Duration::new( - (d_eff.as_secs() as f32 * burst_pct).ceil() as u64, - (d_eff.subsec_nanos() as f32 * burst_pct).ceil() as u32); + let burst_duration = d_eff.mul_f32(burst_pct); let burst_limit = std::cmp::max(1, (total_limit as f32 * burst_pct).floor() as usize); debug_assert!(burst_limit <= total_limit); @@ -137,7 +135,21 @@ impl TokenBucket for VectorTokenBucket { for _ in 0..n { timestamps.push_front(now); } - timestamps.len() <= self.total_limit + + // Check total limit. + if self.total_limit < timestamps.len() { + return false; + } + + // Check burst limit. + if let Some(burst_time) = timestamps.get(self.burst_limit) { + let duration_since = now.duration_since(*burst_time); // `now` before `burst_time` will panic. + if duration_since < self.burst_duration { + return false; + } + } + + return true; } fn get_bucket_duration(&self) -> Duration { diff --git a/src/req/token_bucket.test.rs b/src/req/token_bucket.test.rs index ce78ba4..de6f220 100644 --- a/src/req/token_bucket.test.rs +++ b/src/req/token_bucket.test.rs @@ -32,46 +32,42 @@ mod token_bucket { assert_eq!(1, bucket.burst_limit); } - #[test] - fn test_saturated_100_burst() { - let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 1.00); - - Instant::set_time(50_000); - 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_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.50); + let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.95); Instant::set_time(50_000); 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); // Total 951. + Instant::advance_time(475); assert_ne!(None, bucket.get_delay(), "Bucket should have delay."); + Instant::advance_time(476); // Total 951. Extra buffer for Duration(0). - Instant::advance_time(476); // Extra buffer for Duration(0). assert!(bucket.get_tokens(5), "Last 5 tokens should be available."); assert_ne!(None, bucket.get_delay(), "Bucket should have delay."); - Instant::advance_time(51); - assert!(bucket.get_tokens(95), "95 tokens should be available."); + Instant::advance_time(51); // Total 1002. + assert!(bucket.get_tokens(90), "90 tokens should be available."); assert_ne!(None, bucket.get_delay(), "Bucket should have delay."); Instant::advance_time(951); - assert!(bucket.get_tokens(5), "Last 5 tokens should be available."); + assert!(bucket.get_tokens(10), "Last 10 tokens should be available."); + assert_ne!(None, bucket.get_delay(), "Bucket should have delay."); + } + + #[test] + fn test_violated_50_burst() { + let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.50); + + Instant::set_time(50_000); + assert!(!bucket.get_tokens(90), "Burst should be violated."); assert_ne!(None, bucket.get_delay(), "Bucket should have delay."); } #[test] fn test_saturated_50_burst() { - let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.5); + let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.50); Instant::set_time(50_000); assert!(bucket.get_tokens(50), "Half the tokens should be immediately available."); @@ -93,18 +89,18 @@ mod token_bucket { #[test] fn test_many() { Instant::set_time(50_000); - let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.95); - assert!(bucket.get_tokens(50), "Should have not violated limit."); - assert_eq!(None, bucket.get_delay(), "Should not be blocked."); - for _ in 0..20_000 { + let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, *D00, 0.5); + 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."); - assert_ne!(None, bucket.get_delay(), "Bucket should have delay."); + 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."); - assert_ne!(None, bucket.get_delay(), "Bucket should have delay."); + 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, "Check memory leak."); + assert!(bucket.timestamps.lock().len() < 110, "Should not memory leak."); } } } diff --git a/tests/tests_jp.rs b/tests/tests_jp.rs index 9b2996b..c678618 100644 --- a/tests/tests_jp.rs +++ b/tests/tests_jp.rs @@ -47,16 +47,17 @@ async_tests!{ Ok(()) }, - // tft-league-v1.getLeagueEntriesForSummoner - // https://github.com/MingweiSamuel/Riven/issues/25 - tft_league_getleagueentriesforsummoner: async { - let sp = RIOT_API.summoner_v4().get_by_summoner_name(Region::JP, "Caihonbbt"); - let sr = sp.await.map_err(|e| e.to_string())?.ok_or("Failed to get \"Caihonbbt\"".to_owned())?; - let lp = RIOT_API.tft_league_v1().get_league_entries_for_summoner(Region::JP, &sr.id); - let lr = lp.await.map_err(|e| e.to_string())?; - rassert!(0 < lr.len()); - Ok(()) - }, + // Disabled: Caihonbbt no longer ranked. + // // tft-league-v1.getLeagueEntriesForSummoner + // // https://github.com/MingweiSamuel/Riven/issues/25 + // tft_league_getleagueentriesforsummoner: async { + // let sp = RIOT_API.summoner_v4().get_by_summoner_name(Region::JP, "Caihonbbt"); + // let sr = sp.await.map_err(|e| e.to_string())?.ok_or("Failed to get \"Caihonbbt\"".to_owned())?; + // let lp = RIOT_API.tft_league_v1().get_league_entries_for_summoner(Region::JP, &sr.id); + // let lr = lp.await.map_err(|e| e.to_string())?; + // rassert!(0 < lr.len()); + // Ok(()) + // }, // tft-league-v1.getTopRatedLadder // https://github.com/MingweiSamuel/Riven/issues/24 tft_league_gettopratedladder: async {