rusty_feeder/limit/

lockfree_rate_limiter.rs

use std::{
    fmt,
    num::NonZeroUsize,
    sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
    time::Duration,
};

use parking_lot::Mutex;
use quanta::Clock;
use smallvec::SmallVec;

// Note: helpers imported but not used in this implementation
// use super::helpers::*;

/// Ultra-low-latency rate limiter designed for HFT applications
/// Uses lock-free algorithms to eliminate async/await overhead in critical paths
#[repr(align(64))] // Cache-line alignment for better CPU cache efficiency
pub struct LockFreeRateLimiter {
    /// Maximum number of tokens in bucket
    max_tokens: usize,

    /// Current available tokens, atomic for lock-free access
    available_tokens: AtomicUsize,

    /// Rate limit window in nanoseconds
    window_ns: u64,

    /// Last refill timestamp in nanoseconds since epoch
    last_refill: AtomicU64,

    /// Whether this rate limiter is currently being refilled (prevents contention)
    refilling: AtomicBool,

    /// Recent timestamp history for adaptive rate limiting (protected by mutex)
    recent_timestamps: Mutex<SmallVec<[u64; 64]>>,
}

/// Token guard that provides RAII protection for rate limiting
#[derive(Debug)]
pub struct TokenGuard<'a> {
    limiter: &'a LockFreeRateLimiter,
    acquired: bool,
}

impl<'a> TokenGuard<'a> {
    const fn new(limiter: &'a LockFreeRateLimiter, acquired: bool) -> Self {
        Self { limiter, acquired }
    }

    /// Check if a token was successfully acquired
    #[inline]
    pub const fn is_acquired(&self) -> bool {
        self.acquired
    }

    /// Manually release the token before the guard is dropped
    #[inline]
    pub fn release(mut self) {
        if self.acquired {
            self.limiter.release_token();
            self.acquired = false;
        }
    }
}

impl Drop for TokenGuard<'_> {
    fn drop(&mut self) {
        if self.acquired {
            self.limiter.release_token();
        }
    }
}

impl LockFreeRateLimiter {
    /// Create a new lock-free rate limiter
    #[must_use]
    pub fn new(max_tokens: NonZeroUsize, window: Duration) -> Self {
        static CLOCK: std::sync::OnceLock<Clock> = std::sync::OnceLock::new();
        let clock = CLOCK.get_or_init(Clock::new);
        let now = clock.raw();

        Self {
            max_tokens: max_tokens.get(),
            available_tokens: AtomicUsize::new(max_tokens.get()),
            window_ns: window.as_nanos() as u64,
            last_refill: AtomicU64::new(now),
            refilling: AtomicBool::new(false),
            recent_timestamps: Mutex::new(SmallVec::new()),
        }
    }

    /// Try to acquire a token without waiting (lock-free fast path)
    #[inline(always)]
    pub fn try_acquire(&self) -> TokenGuard<'_> {
        // Fast path: try to decrement available tokens atomically
        let mut current = self.available_tokens.load(Ordering::Acquire);

        while current > 0 {
            match self.available_tokens.compare_exchange_weak(
                current,
                current - 1,
                Ordering::AcqRel,
                Ordering::Relaxed,
            ) {
                Ok(_) => {
                    // Token acquired - record timestamp for adaptive rate limiting
                    // Get a singleton clock for efficiency
                    static CLOCK: std::sync::OnceLock<Clock> = std::sync::OnceLock::new();
                    let clock = CLOCK.get_or_init(Clock::new);
                    let now = clock.raw();
                    if let Some(mut recent) = self.recent_timestamps.try_lock() {
                        recent.push(now);

                        // Keep the buffer at a reasonable size
                        if recent.len() > 128 {
                            recent.drain(0..64);
                        }
                    }

                    // Opportunistically try to refill if we're close to running out of tokens
                    if current <= self.max_tokens / 4 {
                        self.try_refill_tokens();
                    }

                    return TokenGuard::new(self, true);
                }
                Err(actual) => {
                    // Update our view of the current count and try again
                    current = actual;
                }
            }
        }

        // No tokens available - try to refill and try once more
        self.try_refill_tokens();

        current = self.available_tokens.load(Ordering::Acquire);
        if current > 0
            && self
                .available_tokens
                .compare_exchange_weak(current, current - 1, Ordering::AcqRel, Ordering::Relaxed)
                .is_ok()
        {
            static CLOCK: std::sync::OnceLock<Clock> = std::sync::OnceLock::new();
            let clock = CLOCK.get_or_init(Clock::new);
            let now = clock.raw();
            if let Some(mut recent) = self.recent_timestamps.try_lock() {
                recent.push(now);
            }
            return TokenGuard::new(self, true);
        }

        // Failed to acquire token
        TokenGuard::new(self, false)
    }

    /// Wait until a token becomes available (may block briefly), returns RAII guard
    #[inline]
    pub fn acquire(&self) -> TokenGuard<'_> {
        // First try the fast path
        let guard = self.try_acquire();
        if guard.is_acquired() {
            return guard;
        }

        // Slow path: wait for tokens to become available
        loop {
            // Try to refill tokens
            self.refill_tokens();

            // Try to acquire a token
            let guard = self.try_acquire();
            if guard.is_acquired() {
                return guard;
            }

            // Calculate wait time until next token
            let wait_time = self.estimate_wait_time();

            // Short sleep before retrying
            if wait_time > 0 {
                std::thread::sleep(Duration::from_nanos(wait_time));
            } else {
                // Yield to other threads if no specific wait time
                std::thread::yield_now();
            }
        }
    }

    /// Release a token back to the pool (used internally by TokenGuard)
    #[inline]
    fn release_token(&self) {
        // No need to track returns - we'll just increment the counter
        self.available_tokens.fetch_add(1, Ordering::Release);
    }

    /// Try to refill tokens without blocking if too much time has elapsed
    #[inline]
    fn try_refill_tokens(&self) {
        // Only one thread should perform refill at a time
        if self
            .refilling
            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            return; // Another thread is already refilling
        }

        // Get the current timestamp
        static CLOCK: std::sync::OnceLock<Clock> = std::sync::OnceLock::new();
        let clock = CLOCK.get_or_init(Clock::new);
        let now = clock.raw();
        let last = self.last_refill.load(Ordering::Relaxed);

        // Check if enough time has elapsed
        let elapsed_ns = now.saturating_sub(last);
        if elapsed_ns >= self.window_ns / 10 {
            // At least 10% of the window has passed
            // Calculate tokens to add
            let refill_amount =
                ((elapsed_ns as f64 / self.window_ns as f64) * self.max_tokens as f64) as usize;

            if refill_amount > 0 {
                // Add tokens up to max
                let current = self.available_tokens.load(Ordering::Relaxed);
                let new_count = std::cmp::min(current + refill_amount, self.max_tokens);

                if new_count > current {
                    self.available_tokens.store(new_count, Ordering::Release);

                    // Update last refill time proportionally to tokens added
                    let new_last =
                        last + (elapsed_ns * refill_amount as u64 / self.max_tokens as u64);
                    self.last_refill.store(new_last, Ordering::Release);
                }
            }
        }

        // Release the refill lock
        self.refilling.store(false, Ordering::Release);
    }

    /// Force a refill of tokens based on elapsed time
    #[inline]
    fn refill_tokens(&self) {
        // Similar to try_refill but will always attempt to refill tokens
        // Only one thread should perform refill at a time
        if self
            .refilling
            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            return; // Another thread is already refilling
        }

        // Get the current timestamp
        static CLOCK: std::sync::OnceLock<Clock> = std::sync::OnceLock::new();
        let clock = CLOCK.get_or_init(Clock::new);
        let now = clock.raw();
        let last = self.last_refill.load(Ordering::Relaxed);

        // Calculate elapsed time
        let elapsed_ns = now.saturating_sub(last);

        // If full window has elapsed, do a complete refill
        if elapsed_ns >= self.window_ns {
            self.available_tokens
                .store(self.max_tokens, Ordering::Release);
            self.last_refill.store(now, Ordering::Release);

            // Clear timestamps when doing a full refill
            if let Some(mut recent) = self.recent_timestamps.try_lock() {
                recent.clear();
            }
        } else if elapsed_ns > 0 {
            // Partial refill based on elapsed time
            let refill_fraction = elapsed_ns as f64 / self.window_ns as f64;
            let refill_amount = (refill_fraction * self.max_tokens as f64) as usize;

            if refill_amount > 0 {
                // Add tokens up to max
                let current = self.available_tokens.load(Ordering::Relaxed);
                let new_count = std::cmp::min(current + refill_amount, self.max_tokens);

                if new_count > current {
                    self.available_tokens.store(new_count, Ordering::Release);

                    // Update last refill time proportionally
                    let new_last =
                        last + (elapsed_ns * refill_amount as u64 / self.max_tokens as u64);
                    self.last_refill.store(new_last, Ordering::Release);

                    // Prune old timestamps
                    if let Some(mut recent) = self.recent_timestamps.try_lock() {
                        let cutoff = now.saturating_sub(self.window_ns);
                        while let Some(&first) = recent.first() {
                            if first < cutoff {
                                recent.remove(0);
                            } else {
                                break;
                            }
                        }
                    }
                }
            }
        }

        // Release the refill lock
        self.refilling.store(false, Ordering::Release);
    }

    /// Estimate wait time until next token becomes available
    #[inline]
    fn estimate_wait_time(&self) -> u64 {
        static CLOCK: std::sync::OnceLock<Clock> = std::sync::OnceLock::new();
        let clock = CLOCK.get_or_init(Clock::new);
        let now = clock.raw();

        // If we have timestamps, use the oldest one to estimate when it will expire
        if let Some(recent_timestamps) = self.recent_timestamps.try_lock()
            && let Some(&oldest) = recent_timestamps.first()
        {
            // Calculate when the oldest timestamp leaves the window
            let token_expiry = oldest + self.window_ns;
            if token_expiry > now {
                return token_expiry - now;
            }
        }

        // If we can't determine from timestamps, estimate based on refill rate
        let last_refill = self.last_refill.load(Ordering::Relaxed);
        let next_token_time = last_refill + (self.window_ns / self.max_tokens as u64);

        next_token_time.saturating_sub(now)
    }

    /// Get the maximum tokens
    #[inline]
    pub const fn max_tokens(&self) -> usize {
        self.max_tokens
    }

    /// Get the current available tokens
    #[inline]
    pub fn available_tokens(&self) -> usize {
        self.available_tokens.load(Ordering::Relaxed)
    }

    /// Get the window duration
    #[inline]
    pub const fn window(&self) -> Duration {
        Duration::from_nanos(self.window_ns)
    }
}

impl Drop for LockFreeRateLimiter {
    fn drop(&mut self) {
        // Nothing to clean up - we use atomic primitives
    }
}

impl Clone for LockFreeRateLimiter {
    fn clone(&self) -> Self {
        // Cloning just creates a new reference to the same limiter
        // since all internal state is managed with atomics
        Self {
            max_tokens: self.max_tokens,
            available_tokens: AtomicUsize::new(self.available_tokens.load(Ordering::Relaxed)),
            window_ns: self.window_ns,
            last_refill: AtomicU64::new(self.last_refill.load(Ordering::Relaxed)),
            refilling: AtomicBool::new(false),
            recent_timestamps: Mutex::new(SmallVec::new()),
        }
    }
}

impl fmt::Debug for LockFreeRateLimiter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("LockFreeRateLimiter")
            .field("max_tokens", &self.max_tokens)
            .field(
                "available_tokens",
                &self.available_tokens.load(Ordering::Relaxed),
            )
            .field("window_ns", &self.window_ns)
            .field("last_refill", &self.last_refill.load(Ordering::Relaxed))
            .field("refilling", &self.refilling.load(Ordering::Relaxed))
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use std::thread;

    #[test]
    fn test_basic_rate_limiting() {
        // Initialize a static clock to ensure consistent timestamps
        static CLOCK: std::sync::OnceLock<Clock> = std::sync::OnceLock::new();
        let clock = CLOCK.get_or_init(Clock::new);

        // Create a rate limiter with a very small window and few tokens for testing
        let limiter =
            LockFreeRateLimiter::new(NonZeroUsize::new(5).unwrap(), Duration::from_millis(500));

        // Override the last_refill time to ensure we have a clean start
        let now = clock.raw();
        limiter.last_refill.store(now, Ordering::SeqCst);

        // Should be able to acquire all 5 tokens immediately
        for i in 0..5 {
            let token = limiter.try_acquire();
            assert!(token.is_acquired(), "Failed to acquire token {i}");
        }

        // Reset available tokens to ensure the test state is clean
        limiter.available_tokens.store(0, Ordering::SeqCst);

        // 6th token should fail since we manually set available tokens to 0
        let token = limiter.try_acquire();
        assert!(
            !token.is_acquired(),
            "Token should not be acquired when available_tokens is 0"
        );

        // Wait for tokens to refill (longer than the window to ensure refill)
        thread::sleep(Duration::from_millis(600));

        // Force a refill
        limiter.refill_tokens();

        // Should be able to acquire tokens again
        let token = limiter.try_acquire();
        assert!(token.is_acquired(), "Failed to acquire token after refill");
    }

    #[test]
    fn test_token_guard_release() {
        let limiter =
            LockFreeRateLimiter::new(NonZeroUsize::new(1).unwrap(), Duration::from_millis(100));

        // Manually set available tokens to 1 to ensure test consistency
        limiter.available_tokens.store(1, Ordering::SeqCst);

        // Acquire the only token
        let token = limiter.try_acquire();
        assert!(token.is_acquired());

        // Second attempt should fail
        let token2 = limiter.try_acquire();
        assert!(!token2.is_acquired());

        // Release the first token
        token.release();

        // Should be able to acquire again
        let token3 = limiter.try_acquire();
        assert!(token3.is_acquired());
    }

    #[test]
    fn test_concurrent_rate_limiting() {
        // Use a larger token pool and window for this test
        let limiter = Arc::new(LockFreeRateLimiter::new(
            NonZeroUsize::new(100).unwrap(),
            Duration::from_millis(1000),
        ));

        // Manually set available tokens to the maximum to ensure test starts in a known state
        limiter.available_tokens.store(100, Ordering::SeqCst);

        // Create a barrier to ensure all threads start at approximately the same time
        let barrier = Arc::new(std::sync::Barrier::new(11)); // 10 threads + main thread

        let threads: Vec<_> = (0..10)
            .map(|_| {
                let limiter_clone = Arc::clone(&limiter);
                let barrier_clone = Arc::clone(&barrier);
                thread::spawn(move || {
                    // Wait for all threads to be ready
                    barrier_clone.wait();

                    let mut success_count = 0;
                    for _ in 0..20 {
                        let token = limiter_clone.try_acquire();
                        if token.is_acquired() {
                            success_count += 1;
                            // Hold the token briefly, but don't sleep too long
                            // to avoid test taking too long
                            thread::yield_now();
                        }
                    }
                    success_count
                })
            })
            .collect();

        // Wait with all threads
        barrier.wait();

        let total_success: usize = threads
            .into_iter()
            .map(|handle| handle.join().unwrap())
            .sum();

        // Total successful acquisitions should be at least 100 (the initial token count)
        // but not more than 200 (could be more than 100 due to returns during test)
        assert!(
            (100..=200).contains(&total_success),
            "Expected 100-200 successful acquisitions, got {total_success}"
        );
    }

    #[test]
    fn test_blocking_acquire() {
        // Create a static instance of Clock to use for timing
        static CLOCK: std::sync::OnceLock<Clock> = std::sync::OnceLock::new();
        let clock = CLOCK.get_or_init(Clock::new);

        // Use a longer duration to ensure we can measure the wait time accurately
        let limiter =
            LockFreeRateLimiter::new(NonZeroUsize::new(1).unwrap(), Duration::from_millis(100));

        // Manually set available tokens to 0 to ensure we have to wait
        limiter.available_tokens.store(0, Ordering::SeqCst);

        // Set last_refill to current time
        let now = clock.raw();
        limiter.last_refill.store(now, Ordering::SeqCst);

        // Start a separate thread that will release a token after a delay
        let limiter_clone = Arc::new(limiter);
        let thread_limiter = Arc::clone(&limiter_clone);

        let handle = thread::spawn(move || {
            // Wait to simulate a delay
            thread::sleep(Duration::from_millis(50));
            // Add a token
            thread_limiter
                .available_tokens
                .fetch_add(1, Ordering::SeqCst);
        });

        // Start a timer
        let start = clock.now();

        // This should block until the token is available from the other thread
        let token = limiter_clone.acquire();
        assert!(
            token.is_acquired(),
            "Token should be acquired after waiting"
        );

        // Should have waited at least 40ms
        let elapsed = clock.now().duration_since(start);
        assert!(
            elapsed >= Duration::from_millis(40),
            "Expected to wait at least 40ms, but only waited {elapsed:?}"
        );

        // Ensure the thread completes
        handle.join().unwrap();
    }
}