#![allow(dead_code)]
use std::time::{Duration, Instant};
#[non_exhaustive]
#[derive(Copy, Clone, Debug)]
pub enum RateClass {
Regular,
Moderator,
Known,
Verified,
}
impl Default for RateClass {
fn default() -> Self {
Self::Regular
}
}
impl RateClass {
pub fn tickets(self) -> u64 {
match self {
Self::Regular => 20,
Self::Moderator => 100,
Self::Known => 50,
Self::Verified => 7500,
}
}
pub const fn period() -> Duration {
Duration::from_secs(30)
}
}
#[derive(Debug, Clone)]
pub struct RateLimit {
cap: u64,
bucket: Bucket,
}
impl Default for RateLimit {
fn default() -> Self {
Self::from_class(<_>::default())
}
}
impl RateLimit {
pub fn set_cap(&mut self, cap: u64) {
self.cap = cap
}
pub fn set_period(&mut self, period: Duration) {
self.bucket.period = period;
}
pub fn get_cap(&self) -> u64 {
self.cap
}
pub fn get_period(&self) -> Duration {
self.bucket.period
}
pub fn from_class(rate_class: RateClass) -> Self {
Self::full(rate_class.tickets(), RateClass::period())
}
pub fn new(cap: u64, initial: u64, period: Duration) -> Self {
Self {
cap,
bucket: Bucket::new(cap, initial, period),
}
}
pub fn full(cap: u64, period: Duration) -> Self {
Self {
cap,
bucket: Bucket::new(cap, cap, period),
}
}
pub fn empty(cap: u64, period: Duration) -> Self {
Self {
cap,
bucket: Bucket::new(cap, 0, period),
}
}
pub fn get_available_tokens(&self) -> u64 {
self.bucket.tokens
}
pub fn get_current_rate_class(&self) -> Option<RateClass> {
const DUR: Duration = Duration::from_secs(30);
let class = match (self.get_cap(), self.get_period()) {
(20, DUR) => RateClass::Regular,
(50, DUR) => RateClass::Known,
(100, DUR) => RateClass::Moderator,
(7500, DUR) => RateClass::Verified,
_ => return None,
};
Some(class)
}
pub fn consume(&mut self, tokens: u64) -> Result<u64, Duration> {
let Self { bucket, .. } = self;
let now = Instant::now();
if let Some(n) = bucket.refill(now) {
bucket.tokens = std::cmp::min(bucket.tokens + n, self.cap);
}
if tokens <= bucket.tokens {
bucket.tokens -= tokens;
bucket.backoff = 0;
return Ok(bucket.tokens);
}
let prev = bucket.tokens;
Err(bucket.estimate(tokens - prev, now))
}
}
#[derive(Debug, Clone)]
struct Bucket {
tokens: u64,
backoff: u32,
next: Instant,
last: Instant,
quantum: u64,
period: Duration,
}
impl Bucket {
fn new(tokens: u64, initial: u64, period: Duration) -> Self {
let now = Instant::now();
Self {
tokens: initial,
backoff: 0,
next: now + period,
last: now,
quantum: tokens,
period,
}
}
fn refill(&mut self, now: Instant) -> Option<u64> {
if now < self.next {
return None;
}
let last = now.duration_since(self.last);
let periods = last.as_nanos().checked_div(self.period.as_nanos())? as u64;
self.last += self.period * (periods as u32);
self.next = self.last + self.period;
(periods * self.quantum).into()
}
fn estimate(&mut self, tokens: u64, now: Instant) -> Duration {
let until = self.next.duration_since(now);
let periods = (tokens.checked_add(self.quantum).unwrap() - 1) / self.quantum;
until + self.period * (periods as u32 - 1)
}
}