Skip to content

Commit

Permalink
add test for current access of RateLimiter (#11)
Browse files Browse the repository at this point in the history
add test for current access
  • Loading branch information
lytefast authored Jul 3, 2024
1 parent 0c878b9 commit 68fad0a
Show file tree
Hide file tree
Showing 6 changed files with 155 additions and 13 deletions.
105 changes: 105 additions & 0 deletions Cargo.lock

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

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ thiserror = "1.0.60"
[dev-dependencies]
chrono = "0.4.38"
tokio = { version = "1.37.0", features = ["full"] }
futures = "0.3.30"

[[example]]
name = "rate_limiter"
Expand Down
9 changes: 5 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -34,17 +34,18 @@ fn check_rate_limit() {
### With `rate-limiter`

```rust
use std::sync::Arc;
use gcra::{GcraError, RateLimit, RateLimiter};

#[tokio::main]
async fn main() -> Result<(), GcraError> {
let rate_limit = RateLimit::per_sec(2);
let rl = RateLimiter::new(4);
let rate_limiter = Arc::new(RateLimiter::new(4));

rl.check("key", rate_limit.clone(), 1).await?;
rl.check("key", rate_limit.clone(), 1).await?;
rate_limiter.check("key", &rate_limit, 1).await?;
rate_limiter.check("key", &rate_limit, 1).await?;

match rl.check("key", rate_limit.clone(), 1).await {
match rate_limiter.check("key", rate_limit.clone(), 1).await {
Err(GcraError::DeniedUntil { next_allowed_at }) => {
print!("Denied: Request next at {:?}", next_allowed_at);
Ok(())
Expand Down
10 changes: 6 additions & 4 deletions examples/rate_limiter.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::sync::Arc;

use gcra::{GcraError, RateLimit, RateLimiter};

const CACHE_CAPACITY: usize = 4;
Expand All @@ -6,12 +8,12 @@ const WORKER_SHARD_COUNT: usize = 2;
#[tokio::main]
async fn main() -> Result<(), GcraError> {
let rate_limit = RateLimit::per_sec(2);
let rl = RateLimiter::with_shards(CACHE_CAPACITY, WORKER_SHARD_COUNT);
let rate_limiter = Arc::new(RateLimiter::with_shards(CACHE_CAPACITY, WORKER_SHARD_COUNT));

rl.check("key", &rate_limit, 1).await?;
rl.check("key", &rate_limit, 1).await?;
rate_limiter.check("key", &rate_limit, 1).await?;
rate_limiter.check("key", &rate_limit, 1).await?;

match rl.check("key", &rate_limit, 1).await {
match rate_limiter.check("key", &rate_limit, 1).await {
Err(GcraError::DeniedUntil { next_allowed_at }) => {
print!("Denied: Request next at {:?}", next_allowed_at);
Ok(())
Expand Down
9 changes: 5 additions & 4 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,17 +29,18 @@
//! ## With `rate-limiter`
//!
//! ```rust
//! use std::sync::Arc;
//! use gcra::{GcraError, RateLimit, RateLimiter};
//!
//! #[tokio::main]
//! async fn main() -> Result<(), GcraError> {
//! let rate_limit = RateLimit::per_sec(2);
//! let rl = RateLimiter::new(4);
//! let rate_limiter = Arc::new(RateLimiter::new(4));
//!
//! rl.check("key", &rate_limit, 1).await?;
//! rl.check("key", &rate_limit, 1).await?;
//! rate_limiter.check("key", &rate_limit, 1).await?;
//! rate_limiter.check("key", &rate_limit, 1).await?;
//!
//! match rl.check("key", &rate_limit, 1).await {
//! match rate_limiter.check("key", &rate_limit, 1).await {
//! Err(GcraError::DeniedUntil { next_allowed_at }) => {
//! print!("Denied: Request next at {:?}", next_allowed_at);
//! Ok(())
Expand Down
34 changes: 33 additions & 1 deletion src/rate_limiter/rate_limiter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -134,9 +134,14 @@ where

#[cfg(test)]
mod tests {
use futures::stream::{self, StreamExt};

use crate::clock::tests::FakeClock;
use core::panic;
use std::time::{Duration, Instant};
use std::{
sync::Arc,
time::{Duration, Instant},
};

use super::*;

Expand All @@ -161,6 +166,33 @@ mod tests {
}
}

#[tokio::test]
async fn rate_limiter_run_until_denied_concurrent_access() {
let rate_limit = RateLimit::new(3, Duration::from_secs(3));
let rate_limiter = Arc::new(RateLimiter::with_shards(4, 2));

let all_checked = stream::iter(0..rate_limit.resource_limit)
.then(|_| async {
let rate_limiter = rate_limiter.clone();
rate_limiter.check("key", &rate_limit, 1).await
})
.all(|result| async move { result.is_ok() })
.await;

assert!(
all_checked,
"All checks should have passed and not rate limited"
);

match rate_limiter.check("key", &rate_limit, 1).await {
Ok(_) => panic!("We should be rate limited"),
Err(GcraError::DeniedUntil { next_allowed_at }) => {
assert!(next_allowed_at > Instant::now())
}
Err(_) => panic!("Unexpected error"),
}
}

#[tokio::test]
async fn rate_limiter_indefinitly_denied() {
let rate_limit = RateLimit::new(3, Duration::from_secs(3));
Expand Down

0 comments on commit 68fad0a

Please sign in to comment.