rate_limiters 0.1.7

Flexible and easy-to-use Rust library for implementing rate limits.
Documentation
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, UNIX_EPOCH};

use crate::token_bucket::r#impl::{RateLimiter, RateLimiterShared};

// *** LEAKY BUCKET ***
pub struct LeakyBucket {
    capacity: u32,
    leak_rate: f64,
    water: f64,
    last_check: Instant,
}

impl LeakyBucket {
    pub fn new(capacity: u32, leak_rate: f64) -> Self {
        Self {
            capacity,
            leak_rate,
            water: 0.0,
            last_check: Instant::now(),
        }
    }
}

impl RateLimiter for LeakyBucket {
    fn refresh(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_check).as_secs_f64();
        let leaked = elapsed * self.leak_rate;

        if leaked > 0.0 {
            self.water = (self.water - leaked).max(0.0);
            self.last_check = now;
        }
    }

    fn try_acquire(&mut self, amount: u32) -> bool {
        self.refresh();
        if self.water + amount as f64 <= self.capacity as f64 {
            self.water += amount as f64;
            true
        } else {
            false
        }
    }

    fn get_limit(&self) -> u32 {
        self.capacity
    }

    fn get_remaining(&self) -> u32 {
        (self.capacity as f64 - self.water.round()) as u32
    }

    fn get_used(&self) -> u32 {
        self.water.round() as u32
    }

    fn get_reset(&self) -> u64 {
        let now = std::time::SystemTime::now();
        let seconds = self.water / self.leak_rate;
        let reset_time = now + Duration::from_secs_f64(seconds);
        reset_time.duration_since(UNIX_EPOCH).unwrap().as_secs()
    }
}

// *** LEAKY BUCKET SHARED ***
pub struct LeakyBucketShared {
    inner: Arc<Mutex<LeakyBucket>>,
}

impl LeakyBucketShared {
    pub fn new(capacity: u32, leak_rate: f64) -> Self {
        Self {
            inner: Arc::new(Mutex::new(LeakyBucket::new(capacity, leak_rate))),
        }
    }
}

impl RateLimiterShared for LeakyBucketShared {
    fn refresh(&self) {
        let mut bucket = self.inner.lock().unwrap();
        bucket.refresh()
    }

    fn try_acquire(&self, amount: u32) -> bool {
        let mut bucket = self.inner.lock().unwrap();
        bucket.try_acquire(amount)
    }

    fn get_limit(&self) -> u32 {
        let bucket = self.inner.lock().unwrap();
        bucket.get_limit()
    }

    fn get_remaining(&self) -> u32 {
        let bucket = self.inner.lock().unwrap();
        bucket.get_remaining()
    }

    fn get_used(&self) -> u32 {
        let bucket = self.inner.lock().unwrap();
        bucket.get_used()
    }

    fn get_reset(&self) -> u64 {
        let bucket = self.inner.lock().unwrap();
        bucket.get_reset()
    }
}