rusty_feeder/limit/

optimized_rate_limiter.rs

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

use parking_lot::Mutex;
use smallvec::SmallVec;
use tokio::sync::Semaphore;

use super::helpers::*;

/// Performance-optimized rate limiter for high-frequency trading applications
/// Uses a token bucket algorithm with precise timing and low-latency synchronization
#[repr(align(64))] // Cache-line alignment for better CPU cache efficiency
#[derive(Clone)]
pub struct OptimizedRateLimiter {
    /// Maximum number of concurrent requests
    concurrency_limit: Arc<Semaphore>,

    /// Rate limit state (protected by mutex)
    rate_state: Arc<Mutex<RateLimitState>>,

    /// Pre-configured rate limit window
    window: Duration,

    /// Maximum allowed requests per window
    max_tokens: usize,
}

/// Internal state for rate limiting
#[derive(Debug)]
struct RateLimitState {
    /// Current available tokens
    available_tokens: usize,

    /// Last token refresh time
    last_refresh: Instant,

    /// Timestamps of recent requests for adaptive rate limiting
    recent_timestamps: SmallVec<[Instant; 32]>,
}

impl OptimizedRateLimiter {
    /// Create a new optimized rate limiter
    #[must_use]
    pub fn new(max_concurrent: NonZeroUsize, window: Duration) -> Self {
        let max_tokens = max_concurrent.get();

        Self {
            concurrency_limit: Arc::new(Semaphore::new(max_tokens)),
            rate_state: Arc::new(Mutex::new(RateLimitState {
                available_tokens: max_tokens,
                last_refresh: Instant::now(),
                recent_timestamps: SmallVec::new(),
            })),
            window,
            max_tokens,
        }
    }

    /// Wait for a token to become available (async)
    pub async fn wait(&self) -> RateGuard {
        // First, acquire a permit from the semaphore
        let permit = self
            .concurrency_limit
            .clone()
            .acquire_owned()
            .await
            .expect("Semaphore closed");

        // Then check if we're rate limited
        let rate_state = self.rate_state.clone();
        let now = Instant::now();

        // Clone the necessary values from self to move into the closure
        let window = self.window;
        let max_tokens = self.max_tokens;

        let token_wait = tokio::task::spawn_blocking(move || {
            let mut state = rate_state.lock();

            // Check if we need to refresh tokens based on elapsed time
            let (new_tokens, new_last_refresh) = refill_tokens_with_time_adjustment(
                state.available_tokens,
                max_tokens,
                state.last_refresh,
                now,
                window,
            );

            // Update state with new values
            state.available_tokens = new_tokens;
            state.last_refresh = new_last_refresh;

            // Clean up old timestamps if a full reset occurred
            if should_perform_full_reset(now.duration_since(state.last_refresh), window) {
                state.recent_timestamps.clear();
            } else {
                // Clean up old timestamps
                let cutoff = now - window;
                state.recent_timestamps.retain(|ts| *ts >= cutoff);
            }

            // Wait if no tokens are available
            if state.available_tokens == 0 {
                // Calculate time until next token becomes available
                if let Some(&oldest) = state.recent_timestamps.first() {
                    let wait_duration = window
                        .checked_sub(now.duration_since(oldest))
                        .unwrap_or(Duration::from_millis(0));

                    // Return the wait duration
                    Some(wait_duration)
                } else {
                    None // This shouldn't happen if we track timestamps correctly
                }
            } else {
                // Consume a token
                state.available_tokens -= 1;
                state.recent_timestamps.push(now);

                // Sort recent timestamps to ensure oldest first
                state.recent_timestamps.sort_unstable();

                // No wait needed
                None
            }
        })
        .await
        .expect("Token wait task panicked");

        // If token_wait is Some, we need to wait before consuming a token
        if let Some(wait_duration) = token_wait {
            tokio::time::sleep(wait_duration).await;

            // After waiting, record that we've consumed a token
            let mut state = self.rate_state.lock();
            state.available_tokens = state.available_tokens.saturating_sub(1);
            state.recent_timestamps.push(Instant::now());

            // Sort recent timestamps to ensure oldest first
            state.recent_timestamps.sort_unstable();
        }

        // Return guard that will release the semaphore permit when dropped
        RateGuard { _permit: permit }
    }

    /// Try to acquire a token without waiting
    pub fn try_acquire(&self) -> Option<RateGuard> {
        // Try to acquire a permit from the semaphore
        let permit = self.concurrency_limit.clone().try_acquire_owned().ok()?;

        // Check if we're rate limited
        let mut state = self.rate_state.lock();
        let now = Instant::now();

        // Check if we need to refresh tokens based on elapsed time
        let (new_tokens, new_last_refresh) = refill_tokens_with_time_adjustment(
            state.available_tokens,
            self.max_tokens,
            state.last_refresh,
            now,
            self.window,
        );

        // Update state with new values
        state.available_tokens = new_tokens;
        state.last_refresh = new_last_refresh;

        // Clean up old timestamps if a full reset occurred
        if should_perform_full_reset(now.duration_since(state.last_refresh), self.window) {
            state.recent_timestamps.clear();
        } else {
            // Clean up old timestamps
            let cutoff = now - self.window;
            state.recent_timestamps.retain(|ts| *ts >= cutoff);
        }

        // Check if a token is available
        if state.available_tokens > 0 {
            // Consume a token
            state.available_tokens -= 1;
            state.recent_timestamps.push(now);

            // Sort recent timestamps to ensure oldest first
            state.recent_timestamps.sort_unstable();

            // Return the guard
            Some(RateGuard { _permit: permit })
        } else {
            // No token available, release the permit
            drop(permit);
            None
        }
    }
}

/// RAII guard that releases a rate limit token when dropped
pub struct RateGuard {
    _permit: tokio::sync::OwnedSemaphorePermit,
}