foxhunt/crates/ctrader-openapi/src/rate_limiter.rs

//! Token-bucket rate limiter for cTrader Open API.
//!
//! cTrader enforces two rate limits:
//! - **Non-historical**: 50 requests/second (general API calls)
//! - **Historical**: 5 requests/second (tick/candle data requests)

use std::sync::Arc;
use std::time::{Duration, Instant};

use tokio::sync::Mutex;

use crate::error::{CTraderError, RateLimitBucket, Result};

/// Internal state for a single token bucket.
struct Bucket {
    /// Maximum tokens (burst capacity).
    capacity: f64,
    /// Current available tokens.
    tokens: f64,
    /// Token refill rate (tokens per second).
    rate: f64,
    /// Last time tokens were refilled.
    last_refill: Instant,
}

impl Bucket {
    fn new(capacity: f64, rate: f64) -> Self {
        Self {
            capacity,
            tokens: capacity,
            rate,
            last_refill: Instant::now(),
        }
    }

    /// Refill tokens based on elapsed time and try to acquire one.
    /// Returns the wait duration if no token is available.
    fn try_acquire(&mut self) -> Option<Duration> {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
        self.tokens = (self.tokens + elapsed * self.rate).min(self.capacity);
        self.last_refill = now;

        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            None // acquired
        } else {
            // How long until we have 1 token?
            let deficit = 1.0 - self.tokens;
            let wait_secs = deficit / self.rate;
            Some(Duration::from_secs_f64(wait_secs))
        }
    }
}

/// Rate limiter with non-historical and historical buckets.
#[derive(Clone)]
pub struct RateLimiter {
    non_historical: Arc<Mutex<Bucket>>,
    historical: Arc<Mutex<Bucket>>,
    /// Maximum time to wait for a token before giving up.
    max_wait: Duration,
}

impl RateLimiter {
    /// Create a new rate limiter with cTrader's default limits.
    pub fn new() -> Self {
        Self {
            non_historical: Arc::new(Mutex::new(Bucket::new(50.0, 50.0))),
            historical: Arc::new(Mutex::new(Bucket::new(5.0, 5.0))),
            max_wait: Duration::from_secs(5),
        }
    }

    /// Acquire a token from the specified bucket.
    ///
    /// Blocks until a token is available or the max wait time is exceeded.
    pub async fn acquire(&self, bucket: RateLimitBucket) -> Result<()> {
        let bucket_mutex = match bucket {
            RateLimitBucket::NonHistorical => &self.non_historical,
            RateLimitBucket::Historical => &self.historical,
        };

        let deadline = Instant::now() + self.max_wait;

        loop {
            let wait = {
                let mut b = bucket_mutex.lock().await;
                b.try_acquire()
            };

            match wait {
                None => return Ok(()),
                Some(duration) => {
                    if Instant::now() + duration > deadline {
                        return Err(CTraderError::RateLimitExceeded { bucket });
                    }
                    tokio::time::sleep(duration).await;
                }
            }
        }
    }
}

impl Default for RateLimiter {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn acquire_within_limit_succeeds() {
        let limiter = RateLimiter::new();
        for _ in 0..50 {
            limiter
                .acquire(RateLimitBucket::NonHistorical)
                .await
                .expect("should acquire");
        }
    }

    #[tokio::test]
    async fn historical_has_lower_capacity() {
        let limiter = RateLimiter::new();
        for _ in 0..5 {
            limiter
                .acquire(RateLimitBucket::Historical)
                .await
                .expect("should acquire");
        }
    }

    #[test]
    fn bucket_refills_over_time() {
        let mut bucket = Bucket::new(5.0, 5.0);
        for _ in 0..5 {
            assert!(bucket.try_acquire().is_none());
        }
        let wait = bucket.try_acquire();
        assert!(wait.is_some());
        assert!(wait.expect("wait").as_millis() <= 200);
    }
}