do-over 0.1.0

Async resilience policies for Rust inspired by Polly
Documentation
//! Rate limiter policy using a token bucket algorithm.
//!
//! The rate limiter controls the rate of operations by maintaining a bucket
//! of tokens that refill at regular intervals.
//!
//! # How It Works
//!
//! - Each operation consumes one token
//! - If no tokens are available, the operation is rejected
//! - Tokens refill to capacity after each interval
//!
//! # Examples
//!
//! ```rust
//! use do_over::{policy::Policy, rate_limit::RateLimiter, error::DoOverError};
//! use std::time::Duration;
//!
//! # async fn example() -> Result<(), DoOverError<std::io::Error>> {
//! // Allow 100 requests per second
//! let limiter = RateLimiter::new(100, Duration::from_secs(1));
//!
//! match limiter.execute(|| async {
//!     Ok::<_, DoOverError<std::io::Error>>("completed")
//! }).await {
//!     Ok(result) => println!("Success: {}", result),
//!     Err(DoOverError::BulkheadFull) => println!("Rate limit exceeded"),
//!     Err(e) => println!("Error: {:?}", e),
//! }
//! # Ok(())
//! # }
//! ```

use std::time::{Duration, Instant};
use std::sync::Arc;
use tokio::sync::Mutex;
use crate::{policy::Policy, error::DoOverError};

/// A token bucket rate limiter.
///
/// The rate limiter maintains a bucket of tokens. Each operation consumes one
/// token. When tokens are depleted, operations are rejected until the bucket
/// refills.
///
/// # Examples
///
/// ```rust
/// use do_over::{policy::Policy, rate_limit::RateLimiter, error::DoOverError};
/// use std::time::Duration;
///
/// # async fn example() {
/// // 10 requests per second
/// let limiter = RateLimiter::new(10, Duration::from_secs(1));
///
/// // 1000 requests per minute
/// let limiter = RateLimiter::new(1000, Duration::from_secs(60));
/// # }
/// ```
pub struct RateLimiter {
    capacity: u64,
    interval: Duration,
    state: Arc<Mutex<(u64, Instant)>>,
}

impl Clone for RateLimiter {
    fn clone(&self) -> Self {
        Self {
            capacity: self.capacity,
            interval: self.interval,
            state: Arc::clone(&self.state),
        }
    }
}

impl RateLimiter {
    /// Create a new rate limiter.
    ///
    /// # Arguments
    ///
    /// * `capacity` - Number of tokens (requests) allowed per interval
    /// * `interval` - Duration after which tokens refill
    ///
    /// # Examples
    ///
    /// ```rust
    /// use do_over::rate_limit::RateLimiter;
    /// use std::time::Duration;
    ///
    /// // Allow 100 requests per second
    /// let limiter = RateLimiter::new(100, Duration::from_secs(1));
    ///
    /// // Allow 5 requests per 100ms (burst limiting)
    /// let limiter = RateLimiter::new(5, Duration::from_millis(100));
    /// ```
    pub fn new(capacity: u64, interval: Duration) -> Self {
        Self {
            capacity,
            interval,
            state: Arc::new(Mutex::new((capacity, Instant::now()))),
        }
    }
}

#[async_trait::async_trait]
impl<E> Policy<DoOverError<E>> for RateLimiter
where
    E: Send + Sync,
{
    async fn execute<F, Fut, T>(&self, f: F) -> Result<T, DoOverError<E>>
    where
        F: Fn() -> Fut + Send + Sync,
        Fut: std::future::Future<Output = Result<T, DoOverError<E>>> + Send,
        T: Send,
    {
        let mut state = self.state.lock().await;
        if state.1.elapsed() >= self.interval {
            state.0 = self.capacity;
            state.1 = Instant::now();
        }

        if state.0 == 0 {
            return Err(DoOverError::BulkheadFull);
        }

        state.0 -= 1;
        drop(state);
        f().await
    }
}