#![warn(clippy::cargo, clippy::nursery, clippy::pedantic, clippy::restriction)]
#![allow(
clippy::blanket_clippy_restriction_lints,
clippy::missing_inline_in_public_items,
clippy::implicit_return,
clippy::shadow_same,
clippy::separated_literal_suffix
)]
use std::{
num::NonZeroU64,
time::{Duration, Instant},
};
use dashmap::DashMap;
#[must_use]
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
pub struct Limit {
duration: Duration,
count: u16,
}
impl Limit {
pub const fn new(duration: Duration, count: u16) -> Self {
Self { duration, count }
}
}
#[must_use]
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
struct Usage {
time: Instant,
count: u16,
}
impl Usage {
fn new() -> Self {
Self {
time: Instant::now(),
count: 1,
}
}
}
#[must_use]
#[derive(Debug)]
pub struct Bucket {
limit: Limit,
usages: DashMap<NonZeroU64, Usage>,
}
impl Bucket {
pub fn new(limit: Limit) -> Self {
Self {
limit,
usages: DashMap::new(),
}
}
#[allow(clippy::unwrap_used, clippy::integer_arithmetic)]
pub fn register(&self, id: u64) {
let id_non_zero = id.try_into().unwrap();
match self.usages.get_mut(&id_non_zero) {
Some(mut usage) => {
let now = Instant::now();
usage.count = if now - usage.time > self.limit.duration {
1
} else {
usage.count + 1
};
usage.time = now;
}
None => {
self.usages.insert(id_non_zero, Usage::new());
}
}
}
#[must_use]
#[allow(clippy::unwrap_in_result, clippy::unwrap_used)]
pub fn limit_duration(&self, id: u64) -> Option<Duration> {
let usage = self.usages.get(&id.try_into().unwrap())?;
let elapsed = Instant::now() - usage.time;
(usage.count >= self.limit.count && self.limit.duration > elapsed)
.then(|| self.limit.duration - elapsed)
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use tokio::time::sleep;
use crate::{Bucket, Limit};
#[allow(clippy::unwrap_used)]
#[tokio::test]
async fn limit_count_1() {
let bucket = Bucket::new(Limit::new(Duration::from_secs(2), 1));
let id = 123;
assert!(bucket.limit_duration(id).is_none());
bucket.register(id);
assert!(
bucket.limit_duration(id).unwrap()
> bucket.limit.duration - Duration::from_secs_f32(0.1)
);
sleep(bucket.limit.duration).await;
assert!(bucket.limit_duration(id).is_none());
}
#[allow(clippy::unwrap_used)]
#[tokio::test]
async fn limit_count_5() {
let bucket = Bucket::new(Limit::new(Duration::from_secs(5), 5));
let id = 123;
for _ in 0_u8..5 {
assert!(bucket.limit_duration(id).is_none());
bucket.register(id);
}
assert!(
bucket.limit_duration(id).unwrap()
> bucket.limit.duration - Duration::from_secs_f32(0.1)
);
sleep(bucket.limit.duration).await;
assert!(bucket.limit_duration(id).is_none());
}
}