Skip to content

Commit

Permalink
Merge pull request #1 from rob-maron/panic-checking
Browse files Browse the repository at this point in the history
More panic checking
  • Loading branch information
rob-maron authored Oct 29, 2024
2 parents 7db48c6 + afaea52 commit 4b40f18
Show file tree
Hide file tree
Showing 4 changed files with 98 additions and 78 deletions.
2 changes: 1 addition & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
[package]
name = "gcr"
description = "A fast, simple, and small Generic Cell Rate (GCR) algorithm implementation with zero dependencies"
version = "0.1.2"
version = "0.1.3"
edition = "2021"
license = "MIT"
readme = "README.md"
Expand Down
55 changes: 39 additions & 16 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
//! use std::time::Duration;
//!
//! let mut rate = Gcr::new(10, Duration::from_secs(1), Some(30)).unwrap();
//!
//!
//! rate.adjust(20, Duration::from_secs(1), Some(30)).unwrap();
//! // Double the allowed rate while preserving the current capacity
//! ```
Expand Down Expand Up @@ -98,7 +98,7 @@ impl Gcr {
/// * `rate` - The number of units to "refill" per `period`
/// * `period` - The amount of time between each "refill"
/// * `max_burst` - The maximum number of units to allow in a single request. If
/// not specified, this will be set to the rate.
/// not specified, this will be set to the rate.
///
/// Returns a new [`Gcr`] instance on success.
///
Expand All @@ -114,15 +114,17 @@ impl Gcr {
period
.checked_div(rate)
.ok_or(GcrCreationError::ParametersOutOfRange(
"duration division failed: supplied rate was zero".to_string(),
"Supplied rate was zero".to_string(),
))?;

// If not set, the max burst is the rate
let max_burst = max_burst.unwrap_or(rate);

// The delay tolerance is the time between the theoretical arrival time and the
// allow at time
let delay_tolerance = emission_interval * max_burst;
let delay_tolerance = emission_interval.checked_mul(max_burst).ok_or(
GcrCreationError::ParametersOutOfRange("Period / rate was too large".to_string()),
)?;

// This is set to the current time so we can instantly have our full burst
let theoretical_arrival_time = Instant::now();
Expand All @@ -131,8 +133,7 @@ impl Gcr {
let allow_at = theoretical_arrival_time
.checked_sub(delay_tolerance)
.ok_or(GcrCreationError::ParametersOutOfRange(
"interval subtraction failed: max_burst * (period / rate) was too large"
.to_string(),
"Period / rate was too large".to_string(),
))?;

Ok(Self {
Expand All @@ -145,7 +146,7 @@ impl Gcr {
}

/// Get the capacity of the rate limiter at a given time.
///
///
/// Note: this function calculates the capacity on the fly
fn capacity_at(&self, now: Instant) -> u32 {
// Get the duration since the allow at time
Expand All @@ -162,7 +163,7 @@ impl Gcr {
}

/// Get the current capacity of the rate limiter
///
///
/// Note: this function calculates the capacity on the fly
pub fn capacity(&self) -> u32 {
self.capacity_at(Instant::now())
Expand All @@ -185,12 +186,22 @@ impl Gcr {
// This is the canonical request time
let now = Instant::now();

// Calculate how long it would take to allow the request
let required_duration =
self.emission_interval
.checked_mul(n)
.ok_or(GcrRequestError::ParametersOutOfRange(
"Period / rate was too large".to_string(),
))?;

// If the request exceeds capacity, deny it
if n > self.capacity_at(now) {
// If we are not past the virtual theoretical arrival time, disallow the request

// Calculate the time at which all units would have been allowed
let allow_time = self.allow_at + (n * self.emission_interval);
let allow_time = self.allow_at.checked_add(required_duration).ok_or(
GcrRequestError::ParametersOutOfRange("Period / rate was too large".to_string()),
)?;

// See how far it is from the current time
let denied_for = allow_time.checked_duration_since(now);
Expand All @@ -202,15 +213,18 @@ impl Gcr {
// We are past the virtual theoretical arrival time, so allow the request

// Update the theoretical arrival time to account for the new units consumed
self.theoretical_arrival_time =
max(self.theoretical_arrival_time, now) + (n * self.emission_interval);
self.theoretical_arrival_time = max(self.theoretical_arrival_time, now)
.checked_add(required_duration)
.ok_or(GcrRequestError::ParametersOutOfRange(
"Period / rate was too large".to_string(),
))?;

// Update the `allow_at` time to account for the new units consumed
self.allow_at = self
.theoretical_arrival_time
.checked_sub(self.delay_tolerance)
.ok_or(GcrRequestError::ParametersOutOfRange(
"interval subtraction failed: delay_tolerance was too large".to_string(),
"(Period / rate * max_burst) was too large".to_string(),
))?;

Ok(())
Expand All @@ -237,15 +251,24 @@ impl Gcr {
// Update the allow at time to account for the new rate
new_rate.allow_at = now
.checked_sub(
time_since.div_duration_f64(self.emission_interval) as u32
* new_rate.emission_interval,
new_rate
.emission_interval
.checked_mul(time_since.div_duration_f64(self.emission_interval) as u32)
.ok_or(GcrCreationError::ParametersOutOfRange(
"Period / rate was too large".to_string(),
))?,
)
.ok_or(GcrCreationError::ParametersOutOfRange(
"interval subtraction failed: emission_interval was too large".to_string(),
"Period / rate was too large".to_string(),
))?;

// Update the theoretical arrival time to account for the new rate
new_rate.theoretical_arrival_time = new_rate.allow_at + new_rate.delay_tolerance;
new_rate.theoretical_arrival_time = new_rate
.allow_at
.checked_add(new_rate.delay_tolerance)
.ok_or(GcrCreationError::ParametersOutOfRange(
"Delay tolerance was too large".to_string(),
))?;
}

// Replace ourselves with the new rate
Expand Down
117 changes: 57 additions & 60 deletions src/test.rs
Original file line number Diff line number Diff line change
@@ -1,61 +1,58 @@
#[cfg(test)]
mod test {
use std::{thread::sleep, time::Duration};

use crate::{Gcr, GcrRequestError};

#[test]
fn test_request() {
let mut rate: Gcr = Gcr::new(100, Duration::from_millis(100), Some(500))
.expect("Failed to create GCR instance");

// Make sure we can't request more than the max burst, even if we wait
sleep(Duration::from_millis(100));
assert!(matches!(
rate.request(501),
Err(GcrRequestError::RequestTooLarge)
));
assert!(rate.capacity() == 500);

// Make sure we can request up to the burst
rate.request(500).expect("Failed to request burst");
assert!(rate.capacity() == 0 && rate.request(1).is_err());
assert!(rate.allow_at.elapsed().as_secs() == 0);

// Make sure the rate is consistent
sleep(Duration::from_millis(100));
assert!(rate.capacity() / 10 == 10);

// Make sure we are denied for the correct amount of time
sleep(Duration::from_millis(100));
let Err(GcrRequestError::DeniedFor(duration)) = rate.request(500) else {
panic!("Expected a denied for error");
};
assert!(
duration.as_millis() / 10 == 29
|| duration.as_millis() / 10 == 30
|| duration.as_millis() / 10 == 28
);
}

#[test]
fn test_adjust() {
let mut rate: Gcr = Gcr::new(100, Duration::from_millis(100), Some(500))
.expect("Failed to create GCR instance");

// Make sure the capacity stays the same when we adjust the parameters
let mut rate2 = rate.clone();
rate2
.adjust(200, Duration::from_millis(100), Some(1000))
.expect("Failed to adjust GCR");
assert!(rate.capacity() == rate2.capacity());

// Make sure we respect the new rate and burst
rate.request(200).expect("Failed to request 200 units");
rate.adjust(200, Duration::from_millis(100), Some(1000))
.expect("Failed to adjust GCR");
assert!(rate.capacity() == 300);
sleep(Duration::from_millis(200));
assert!(rate.capacity() / 100 == 7);
}
use std::{thread::sleep, time::Duration};

use crate::{Gcr, GcrRequestError};

#[test]
fn test_request() {
let mut rate: Gcr = Gcr::new(100, Duration::from_millis(100), Some(500))
.expect("Failed to create GCR instance");

// Make sure we can't request more than the max burst, even if we wait
sleep(Duration::from_millis(100));
assert!(matches!(
rate.request(501),
Err(GcrRequestError::RequestTooLarge)
));
assert!(rate.capacity() == 500);

// Make sure we can request up to the burst
rate.request(500).expect("Failed to request burst");
assert!(rate.capacity() == 0 && rate.request(1).is_err());
assert!(rate.allow_at.elapsed().as_secs() == 0);

// Make sure the rate is consistent
sleep(Duration::from_millis(100));
assert!(rate.capacity() / 10 == 10);

// Make sure we are denied for the correct amount of time
sleep(Duration::from_millis(100));
let Err(GcrRequestError::DeniedFor(duration)) = rate.request(500) else {
panic!("Expected a denied for error");
};
assert!(
duration.as_millis() / 10 == 29
|| duration.as_millis() / 10 == 30
|| duration.as_millis() / 10 == 28
);
}

#[test]
fn test_adjust() {
let mut rate: Gcr = Gcr::new(100, Duration::from_millis(100), Some(500))
.expect("Failed to create GCR instance");

// Make sure the capacity stays the same when we adjust the parameters
let mut rate2 = rate.clone();
rate2
.adjust(200, Duration::from_millis(100), Some(1000))
.expect("Failed to adjust GCR");
assert!(rate.capacity() == rate2.capacity());

// Make sure we respect the new rate and burst
rate.request(200).expect("Failed to request 200 units");
rate.adjust(200, Duration::from_millis(100), Some(1000))
.expect("Failed to adjust GCR");
assert!(rate.capacity() == 300);
sleep(Duration::from_millis(200));
assert!(rate.capacity() / 100 == 7);
}

0 comments on commit 4b40f18

Please sign in to comment.