use {
std::sync::Arc,
futures::Future,
};
#[cfg(feature = "rate-limit")]
use tokio::{
sync::{Mutex, MutexGuard},
time::{sleep_until, Duration, Instant},
};
const REQ_COOLDOWN_DURATION: Duration = Duration::from_millis(600);
pub trait Limiter {
async fn check<F, R>(self, fut: F) -> R
where
F: Future<Output = R>,
Self: Sized;
}
#[cfg(feature = "rate-limit")]
#[derive(Debug, Clone, Default)]
pub struct RateLimit {
deadline: Arc<Mutex<Option<Instant>>>,
}
#[cfg(feature = "rate-limit")]
struct Guard<'a>(MutexGuard<'a, Option<Instant>>);
#[cfg(feature = "rate-limit")]
impl<'a> Drop for Guard<'a> {
fn drop(&mut self) {
*self.0 = Some(Instant::now() + REQ_COOLDOWN_DURATION);
}
}
#[cfg(feature = "rate-limit")]
impl Limiter for RateLimit {
async fn check<F, R>(self, fut: F) -> R
where
F: Future<Output = R>,
{
let guard = self.lock().await;
let result = fut.await;
drop(guard);
result
}
}
#[cfg(feature = "rate-limit")]
impl RateLimit {
async fn lock(&self) -> Guard {
loop {
let now = Instant::now();
let deadline = {
let guard = self.deadline.lock().await;
match &*guard {
None => return Guard(guard),
Some(deadline) if now >= *deadline => return Guard(guard),
Some(deadline) => *deadline,
}
};
sleep_until(deadline).await;
}
}
}
#[derive(Debug, Clone, Default)]
pub struct DummyRateLimit {}
impl Limiter for DummyRateLimit {
async fn check<F, R>(self, fut: F) -> R
where
F: Future<Output = R>,
Self: Sized
{
fut.await
}
}
#[cfg(feature = "rate-limit")]
pub type DefaultRateLimit = RateLimit;
#[cfg(not(feature = "rate-limit"))]
pub type DefaultRateLimit = DummyRateLimit;