diff --git a/Cargo.lock b/Cargo.lock index 797f07b..0127f0f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -197,7 +197,7 @@ dependencies = [ [[package]] name = "gcra" -version = "0.3.2" +version = "0.3.3" dependencies = [ "chrono", "thingvellir", diff --git a/Cargo.toml b/Cargo.toml index b4d15d4..84ac6d7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "gcra" -version = "0.3.2" +version = "0.3.3" edition = "2021" authors = ["Sam Shih "] license = "MIT" diff --git a/src/gcra.rs b/src/gcra.rs index e2babe0..4724b11 100644 --- a/src/gcra.rs +++ b/src/gcra.rs @@ -87,6 +87,23 @@ impl GcraState { } } } + + pub fn remaining_resources(&self, rate_limit: &RateLimit, now: Instant) -> u32 { + if rate_limit.period.is_zero() { + return 0; + } + + let time_to_tat = match self.tat.and_then(|tat| tat.checked_duration_since(now)) { + Some(duration_until) => duration_until, + None => return rate_limit.resource_limit, + }; + + // Logically this makes more sense as: + // consumed_resources = time_to_tat * (resource_limit/period) + // but we run it this way because of Duration's arithmetic functions + let consumed_resources = (time_to_tat * rate_limit.resource_limit).div_duration_f32(rate_limit.period); + rate_limit.resource_limit - consumed_resources.ceil() as u32 + } } #[cfg(test)] @@ -95,6 +112,29 @@ mod tests { use super::*; + #[test] + fn test_rate_limit_unused_counts() { + let base_tat = Instant::now(); + let rate_limit = RateLimit::new(10, Duration::from_secs(1)); + + assert_eq!( + 4, + GcraState { tat: Some(base_tat+Duration::from_millis(550))}.remaining_resources(&rate_limit, base_tat), + "Remaining count should ceiled" + ); + assert_eq!( + 0, + GcraState { tat: Some(base_tat+Duration::from_millis(950))}.remaining_resources(&rate_limit, base_tat), + "Remaining count should ceiled, thus preventing any additional requests" + ); + + assert_eq!( + 9, + GcraState { tat: Some(base_tat+Duration::from_millis(100))}.remaining_resources(&rate_limit, base_tat), + "Remaining count is based on max_period timeout" + ); + } + #[test] fn gcra_basics() { let mut gcra = GcraState::default(); diff --git a/src/lib.rs b/src/lib.rs index 1e90861..25aaff9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ +#![feature(div_duration)] + //! Library which implements the core //! [GCRA](https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm) functionality in rust. //! diff --git a/src/rate_limit.rs b/src/rate_limit.rs index 2f1e677..77189ad 100644 --- a/src/rate_limit.rs +++ b/src/rate_limit.rs @@ -32,6 +32,7 @@ impl RateLimit { self.emission_interval * cost } } + #[cfg(test)] mod tests { use super::*;