ostium_rust_sdk/

rate_limit.rs

//! Rate limiting functionality for the Ostium SDK
//!
//! This module provides rate limiting capabilities to ensure we don't exceed
//! API rate limits. It uses a token bucket algorithm for flexible rate limiting.

use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;
use tokio::time::sleep;
use tracing::{debug, warn};

/// Configuration for rate limiting
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    /// Maximum number of requests per time window
    pub max_requests: u32,
    /// Time window for the rate limit
    pub time_window: Duration,
    /// Maximum burst size (number of requests that can be made instantly)
    pub burst_size: u32,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            max_requests: 100,
            time_window: Duration::from_secs(60), // 100 requests per minute
            burst_size: 10,
        }
    }
}

impl RateLimitConfig {
    /// Create a config for GraphQL API rate limiting
    pub fn graphql() -> Self {
        Self {
            max_requests: 60,
            time_window: Duration::from_secs(60), // 60 requests per minute
            burst_size: 5,
        }
    }

    /// Create a config for REST API rate limiting  
    pub fn rest_api() -> Self {
        Self {
            max_requests: 120,
            time_window: Duration::from_secs(60), // 120 requests per minute
            burst_size: 10,
        }
    }

    /// Create a config for blockchain RPC rate limiting
    pub fn blockchain() -> Self {
        Self {
            max_requests: 200,
            time_window: Duration::from_secs(60), // 200 requests per minute
            burst_size: 20,
        }
    }

    /// Create a conservative rate limit config
    pub fn conservative() -> Self {
        Self {
            max_requests: 30,
            time_window: Duration::from_secs(60), // 30 requests per minute
            burst_size: 3,
        }
    }
}

/// Token bucket rate limiter
#[derive(Debug)]
pub struct RateLimiter {
    config: RateLimitConfig,
    state: Arc<Mutex<RateLimiterState>>,
}

#[derive(Debug)]
struct RateLimiterState {
    tokens: f64,
    last_refill: Instant,
}

impl RateLimiter {
    /// Create a new rate limiter with the given configuration
    pub fn new(config: RateLimitConfig) -> Self {
        let burst_size = config.burst_size;
        Self {
            state: Arc::new(Mutex::new(RateLimiterState {
                tokens: burst_size as f64,
                last_refill: Instant::now(),
            })),
            config,
        }
    }

    /// Wait for permission to make a request
    /// Returns immediately if a token is available, otherwise waits
    pub async fn acquire(&self) -> Result<(), RateLimitError> {
        loop {
            {
                let mut state = self.state.lock().await;
                self.refill_tokens(&mut state);

                if state.tokens >= 1.0 {
                    state.tokens -= 1.0;
                    debug!(
                        "Rate limit token acquired, {} tokens remaining",
                        state.tokens
                    );
                    return Ok(());
                }
            }

            // Calculate how long to wait for the next token
            let wait_time = self.calculate_wait_time().await;
            debug!(
                "Rate limit exceeded, waiting {:?} for next token",
                wait_time
            );

            if wait_time > Duration::from_secs(30) {
                warn!("Rate limit wait time exceeds 30 seconds, rejecting request");
                return Err(RateLimitError::ExcessiveWait(wait_time));
            }

            sleep(wait_time).await;
        }
    }

    /// Try to acquire a token without waiting
    /// Returns true if successful, false if rate limited
    pub async fn try_acquire(&self) -> bool {
        let mut state = self.state.lock().await;
        self.refill_tokens(&mut state);

        if state.tokens >= 1.0 {
            state.tokens -= 1.0;
            debug!(
                "Rate limit token acquired, {} tokens remaining",
                state.tokens
            );
            true
        } else {
            debug!("Rate limit exceeded, no tokens available");
            false
        }
    }

    /// Get the current number of available tokens
    pub async fn available_tokens(&self) -> f64 {
        let mut state = self.state.lock().await;
        self.refill_tokens(&mut state);
        state.tokens
    }

    /// Calculate how long until the next token is available
    async fn calculate_wait_time(&self) -> Duration {
        let state = self.state.lock().await;
        let tokens_per_second =
            self.config.max_requests as f64 / self.config.time_window.as_secs_f64();
        let time_per_token = Duration::from_secs_f64(1.0 / tokens_per_second);

        // If we have no tokens, wait for one token duration
        if state.tokens <= 0.0 {
            time_per_token
        } else {
            // Otherwise wait a shorter time
            Duration::from_millis(100)
        }
    }

    /// Refill tokens based on elapsed time
    fn refill_tokens(&self, state: &mut RateLimiterState) {
        let now = Instant::now();
        let elapsed = now.duration_since(state.last_refill);

        let tokens_per_second =
            self.config.max_requests as f64 / self.config.time_window.as_secs_f64();
        let tokens_to_add = elapsed.as_secs_f64() * tokens_per_second;

        state.tokens = (state.tokens + tokens_to_add).min(self.config.burst_size as f64);
        state.last_refill = now;
    }
}

/// Rate limiting errors
#[derive(Debug, thiserror::Error)]
pub enum RateLimitError {
    /// Wait time would be excessive
    #[error("Rate limit wait time would be excessive: {0:?}")]
    ExcessiveWait(Duration),
}

/// Rate limiter manager that handles multiple rate limiters
#[derive(Debug)]
pub struct RateLimiterManager {
    graphql_limiter: Option<RateLimiter>,
    rest_limiter: Option<RateLimiter>,
    blockchain_limiter: Option<RateLimiter>,
}

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

impl RateLimiterManager {
    /// Create a new rate limiter manager with no rate limiting
    pub fn new() -> Self {
        Self {
            graphql_limiter: None,
            rest_limiter: None,
            blockchain_limiter: None,
        }
    }

    /// Enable GraphQL rate limiting
    pub fn with_graphql_rate_limit(mut self, config: RateLimitConfig) -> Self {
        self.graphql_limiter = Some(RateLimiter::new(config));
        self
    }

    /// Enable REST API rate limiting  
    pub fn with_rest_rate_limit(mut self, config: RateLimitConfig) -> Self {
        self.rest_limiter = Some(RateLimiter::new(config));
        self
    }

    /// Enable blockchain RPC rate limiting
    pub fn with_blockchain_rate_limit(mut self, config: RateLimitConfig) -> Self {
        self.blockchain_limiter = Some(RateLimiter::new(config));
        self
    }

    /// Enable all rate limiters with default configs
    pub fn with_default_limits(mut self) -> Self {
        self.graphql_limiter = Some(RateLimiter::new(RateLimitConfig::graphql()));
        self.rest_limiter = Some(RateLimiter::new(RateLimitConfig::rest_api()));
        self.blockchain_limiter = Some(RateLimiter::new(RateLimitConfig::blockchain()));
        self
    }

    /// Wait for permission to make a GraphQL request
    pub async fn acquire_graphql(&self) -> Result<(), RateLimitError> {
        if let Some(limiter) = &self.graphql_limiter {
            limiter.acquire().await
        } else {
            Ok(())
        }
    }

    /// Wait for permission to make a REST API request
    pub async fn acquire_rest(&self) -> Result<(), RateLimitError> {
        if let Some(limiter) = &self.rest_limiter {
            limiter.acquire().await
        } else {
            Ok(())
        }
    }

    /// Wait for permission to make a blockchain RPC request
    pub async fn acquire_blockchain(&self) -> Result<(), RateLimitError> {
        if let Some(limiter) = &self.blockchain_limiter {
            limiter.acquire().await
        } else {
            Ok(())
        }
    }

    /// Try to acquire permission for GraphQL without waiting
    pub async fn try_acquire_graphql(&self) -> bool {
        if let Some(limiter) = &self.graphql_limiter {
            limiter.try_acquire().await
        } else {
            true
        }
    }

    /// Try to acquire permission for REST API without waiting  
    pub async fn try_acquire_rest(&self) -> bool {
        if let Some(limiter) = &self.rest_limiter {
            limiter.try_acquire().await
        } else {
            true
        }
    }

    /// Try to acquire permission for blockchain RPC without waiting
    pub async fn try_acquire_blockchain(&self) -> bool {
        if let Some(limiter) = &self.blockchain_limiter {
            limiter.try_acquire().await
        } else {
            true
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::time::{sleep, Duration};

    #[tokio::test]
    async fn test_rate_limiter_basic() {
        let config = RateLimitConfig {
            max_requests: 5,
            time_window: Duration::from_secs(1),
            burst_size: 2,
        };

        let limiter = RateLimiter::new(config);

        // Should be able to make 2 requests immediately (burst)
        assert!(limiter.try_acquire().await);
        assert!(limiter.try_acquire().await);

        // Third request should be rate limited
        assert!(!limiter.try_acquire().await);
    }

    #[tokio::test]
    async fn test_rate_limiter_refill() {
        let config = RateLimitConfig {
            max_requests: 10,
            time_window: Duration::from_secs(1),
            burst_size: 1,
        };

        let limiter = RateLimiter::new(config);

        // Use up the token
        assert!(limiter.try_acquire().await);
        assert!(!limiter.try_acquire().await);

        // Wait for refill (should get a token every 100ms with 10 req/sec)
        sleep(Duration::from_millis(150)).await;

        // Should have a token now
        assert!(limiter.try_acquire().await);
    }

    #[tokio::test]
    async fn test_rate_limiter_manager() {
        let manager = RateLimiterManager::new().with_graphql_rate_limit(RateLimitConfig {
            max_requests: 2,
            time_window: Duration::from_secs(1),
            burst_size: 1,
        });

        // Should work for GraphQL
        assert!(manager.acquire_graphql().await.is_ok());

        // REST should work without limits
        assert!(manager.try_acquire_rest().await);
        assert!(manager.try_acquire_rest().await);
    }
}