rusty_ems/

throttle.rs

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use anyhow::Result;
use crossbeam::queue::SegQueue;
use quanta::Clock;
use tokio::time;

/// A rate limiter that enforces rate limits for exchange API calls
#[derive(Debug)]
pub struct RateLimiter {
    /// The maximum number of requests allowed in the window
    max_requests_per_window: usize,

    /// The time window in milliseconds
    window_milliseconds: u64,

    /// Request timestamps for sliding window enforcement
    /// Using a lock-free queue for better performance under contention
    request_times: Arc<SegQueue<u64>>,

    /// Current number of active requests in the window
    /// This is an approximation and may be higher than actual due to expired requests
    active_count: AtomicUsize,

    /// High-precision clock
    clock: Clock,
}

impl RateLimiter {
    /// Creates a new rate limiter
    ///
    /// # Parameters
    /// * `max_requests_per_window` - Maximum number of requests allowed per time window
    /// * `window_milliseconds` - Time window in milliseconds
    #[must_use]
    pub fn new(max_requests_per_window: usize, window_milliseconds: u64) -> Self {
        Self {
            max_requests_per_window,
            window_milliseconds,
            request_times: Arc::new(SegQueue::new()),
            active_count: AtomicUsize::new(0),
            clock: Clock::new(),
        }
    }

    /// Acquires permission to make a request, waiting if necessary
    ///
    /// # Returns
    /// `Ok(())` when permission is granted, or an error if the wait was interrupted
    pub async fn acquire(&self) -> Result<()> {
        loop {
            let now = self.clock.raw() / 1_000_000; // Convert ns to ms

            // Count active requests and find the oldest timestamp
            let mut active_count = 0;
            let mut oldest_active_timestamp = now;
            let mut timestamps = Vec::new();

            // Drain the queue temporarily to scan all timestamps
            while let Some(timestamp) = self.request_times.pop() {
                timestamps.push(timestamp);

                // Check if this timestamp is still active
                if now - timestamp <= self.window_milliseconds {
                    active_count += 1;
                    if timestamp < oldest_active_timestamp {
                        oldest_active_timestamp = timestamp;
                    }
                }
            }

            // Put all timestamps back in the queue
            for timestamp in timestamps {
                self.request_times.push(timestamp);
            }

            // Update the atomic counter with the accurate count
            self.active_count.store(active_count, Ordering::SeqCst);

            if active_count < self.max_requests_per_window {
                // We have capacity, add a new timestamp
                self.request_times.push(now);
                self.active_count.fetch_add(1, Ordering::SeqCst);
                return Ok(());
            }

            // We need to wait until the oldest active request expires
            let elapsed = now - oldest_active_timestamp;
            if elapsed >= self.window_milliseconds {
                // The oldest request has already expired (this shouldn't happen with our counting logic)
                self.request_times.push(now);
                self.active_count.fetch_add(1, Ordering::SeqCst);
                return Ok(());
            }

            // Wait until the oldest request expires
            let wait_time_ms = self.window_milliseconds - elapsed + 1; // Add 1ms buffer

            // Wait for the calculated time
            time::sleep(Duration::from_millis(wait_time_ms)).await;
        }
    }

    /// Tries to acquire permission to make a request without waiting
    ///
    /// # Returns
    /// `true` if permission was granted, `false` if the rate limit would be exceeded
    pub fn try_acquire(&self) -> bool {
        let now = self.clock.raw() / 1_000_000; // Convert ns to ms

        // First, try a fast path using the atomic counter
        // This is an approximation and may reject some requests unnecessarily
        let current_count = self.active_count.load(Ordering::Relaxed);
        if current_count < self.max_requests_per_window {
            // Optimistically assume we have capacity and try to increment
            if self
                .active_count
                .compare_exchange(
                    current_count,
                    current_count + 1,
                    Ordering::SeqCst,
                    Ordering::Relaxed,
                )
                .is_ok()
            {
                // We successfully claimed a slot
                self.request_times.push(now);
                return true;
            }
        }

        // Slow path: count active requests accurately
        let mut active_count = 0;
        let mut timestamps = Vec::new();

        // Drain the queue temporarily to scan all timestamps
        while let Some(timestamp) = self.request_times.pop() {
            timestamps.push(timestamp);

            // Check if this timestamp is still active
            if now - timestamp <= self.window_milliseconds {
                active_count += 1;
            }
        }

        // Put all timestamps back in the queue
        for timestamp in timestamps {
            self.request_times.push(timestamp);
        }

        // Update the atomic counter with the accurate count
        self.active_count.store(active_count, Ordering::SeqCst);

        if active_count < self.max_requests_per_window {
            // We have capacity
            self.request_times.push(now);
            self.active_count.fetch_add(1, Ordering::SeqCst);
            true
        } else {
            // Rate limit would be exceeded
            false
        }
    }

    /// Marks a request as completed, useful for tracking successful API calls
    pub fn mark_request(&self) {
        let now = self.clock.raw() / 1_000_000; // Convert ns to ms

        // Add the timestamp to the queue
        self.request_times.push(now);

        // Increment the active count
        self.active_count.fetch_add(1, Ordering::SeqCst);

        // Note: We don't need to remove old timestamps here since they'll be
        // filtered out when counting active requests in acquire/try_acquire
    }

    /// Returns the number of remaining requests that can be made immediately
    pub fn remaining_requests(&self) -> usize {
        // Fast path: use the atomic counter for an approximate count
        let current_count = self.active_count.load(Ordering::Relaxed);

        // If we're well below the limit, we can return immediately
        if current_count < self.max_requests_per_window / 2 {
            return self.max_requests_per_window - current_count;
        }

        // Slow path: count active requests accurately
        let now = self.clock.raw() / 1_000_000; // Convert ns to ms
        let mut active_count = 0;
        let mut timestamps = Vec::new();

        // Drain the queue temporarily to scan all timestamps
        while let Some(timestamp) = self.request_times.pop() {
            timestamps.push(timestamp);

            // Check if this timestamp is still active
            if now - timestamp <= self.window_milliseconds {
                active_count += 1;
            }
        }

        // Put all timestamps back in the queue
        for timestamp in timestamps {
            self.request_times.push(timestamp);
        }

        // Update the atomic counter with the accurate count
        self.active_count.store(active_count, Ordering::SeqCst);

        self.max_requests_per_window.saturating_sub(active_count)
    }

    /// Returns the time in milliseconds until the next slot becomes available
    pub fn time_until_next_available_ms(&self) -> u64 {
        let now = self.clock.raw() / 1_000_000; // Convert ns to ms

        // Fast path: check if we're below capacity using the atomic counter
        let current_count = self.active_count.load(Ordering::Relaxed);
        if current_count < self.max_requests_per_window {
            // We have capacity now
            return 0;
        }

        // Slow path: scan all timestamps to find the oldest active one
        let mut active_count = 0;
        let mut oldest_active_timestamp = now;
        let mut timestamps = Vec::new();

        // Drain the queue temporarily to scan all timestamps
        while let Some(timestamp) = self.request_times.pop() {
            timestamps.push(timestamp);

            // Check if this timestamp is still active
            if now - timestamp <= self.window_milliseconds {
                active_count += 1;
                if timestamp < oldest_active_timestamp {
                    oldest_active_timestamp = timestamp;
                }
            }
        }

        // Put all timestamps back in the queue
        for timestamp in timestamps {
            self.request_times.push(timestamp);
        }

        // Update the atomic counter with the accurate count
        self.active_count.store(active_count, Ordering::SeqCst);

        if active_count < self.max_requests_per_window {
            // We have capacity now
            return 0;
        }

        // Calculate time until the oldest active request expires
        let elapsed = now - oldest_active_timestamp;
        if elapsed >= self.window_milliseconds {
            // The oldest request has already expired
            return 0;
        }

        // Return time until the oldest request expires
        self.window_milliseconds - elapsed
    }
}

/// A weight-based rate limiter for more complex exchange rate limiting rules
#[derive(Debug)]
pub struct WeightBasedRateLimiter {
    /// The maximum weight allowed in the window
    max_weight: usize,

    /// The time window in milliseconds
    window_milliseconds: u64,

    /// Request weights and timestamps for sliding window enforcement
    /// Using a lock-free queue for better performance under contention
    requests: Arc<SegQueue<(u64, usize)>>, // (timestamp_ms, weight)

    /// Current total weight of active requests in the window
    /// This is an approximation and may be higher than actual due to expired requests
    active_weight: AtomicUsize,

    /// High-precision clock
    clock: Clock,
}

impl WeightBasedRateLimiter {
    /// Creates a new weight-based rate limiter
    ///
    /// # Parameters
    /// * `max_weight` - Maximum weight allowed per time window
    /// * `window_milliseconds` - Time window in milliseconds
    #[must_use]
    pub fn new(max_weight: usize, window_milliseconds: u64) -> Self {
        Self {
            max_weight,
            window_milliseconds,
            requests: Arc::new(SegQueue::new()),
            active_weight: AtomicUsize::new(0),
            clock: Clock::new(),
        }
    }

    /// Acquires permission to make a request with a specific weight, waiting if necessary
    ///
    /// # Parameters
    /// * `weight` - The weight of the request
    ///
    /// # Returns
    /// `Ok(())` when permission is granted, or an error if the wait was interrupted
    pub async fn acquire(&self, weight: usize) -> Result<()> {
        loop {
            let now = self.clock.raw() / 1_000_000; // Convert ns to ms

            // Count active weight and find the oldest timestamp
            let mut active_weight = 0;
            let mut oldest_active_timestamp = now;
            let mut entries = Vec::new();

            // Drain the queue temporarily to scan all entries
            while let Some(entry) = self.requests.pop() {
                entries.push(entry);

                // Check if this entry is still active
                let (timestamp, entry_weight) = entry;
                if now - timestamp <= self.window_milliseconds {
                    active_weight += entry_weight;
                    if timestamp < oldest_active_timestamp {
                        oldest_active_timestamp = timestamp;
                    }
                }
            }

            // Put all entries back in the queue
            for entry in entries {
                self.requests.push(entry);
            }

            // Update the atomic counter with the accurate weight
            self.active_weight.store(active_weight, Ordering::SeqCst);

            if active_weight + weight <= self.max_weight {
                // We have capacity, add a new entry
                self.requests.push((now, weight));
                self.active_weight.fetch_add(weight, Ordering::SeqCst);
                return Ok(());
            }

            // We need to wait until the oldest active request expires
            let elapsed = now - oldest_active_timestamp;
            if elapsed >= self.window_milliseconds {
                // The oldest request has already expired (this shouldn't happen with our counting logic)
                self.requests.push((now, weight));
                self.active_weight.fetch_add(weight, Ordering::SeqCst);
                return Ok(());
            }

            // Wait until the oldest request expires
            let wait_time_ms = self.window_milliseconds - elapsed + 1; // Add 1ms buffer

            // Wait for the calculated time
            time::sleep(Duration::from_millis(wait_time_ms)).await;
        }
    }

    /// Tries to acquire permission to make a request without waiting
    ///
    /// # Parameters
    /// * `weight` - The weight of the request
    ///
    /// # Returns
    /// `true` if permission was granted, `false` if the rate limit would be exceeded
    pub fn try_acquire(&self, weight: usize) -> bool {
        let now = self.clock.raw() / 1_000_000; // Convert ns to ms

        // First, try a fast path using the atomic counter
        // This is an approximation and may reject some requests unnecessarily
        let current_weight = self.active_weight.load(Ordering::Relaxed);
        if current_weight + weight <= self.max_weight {
            // Optimistically assume we have capacity
            // We can't use compare_exchange here because we're adding a variable amount (weight)
            // Instead, we'll add the weight and then verify we didn't exceed the limit
            let new_weight = self.active_weight.fetch_add(weight, Ordering::SeqCst) + weight;
            if new_weight <= self.max_weight {
                // We successfully claimed capacity
                self.requests.push((now, weight));
                return true;
            }
            // We exceeded the limit, rollback the weight addition
            self.active_weight.fetch_sub(weight, Ordering::SeqCst);
        }

        // Slow path: count active weight accurately
        let mut active_weight = 0;
        let mut entries = Vec::new();

        // Drain the queue temporarily to scan all entries
        while let Some(entry) = self.requests.pop() {
            entries.push(entry);

            // Check if this entry is still active
            let (timestamp, entry_weight) = entry;
            if now - timestamp <= self.window_milliseconds {
                active_weight += entry_weight;
            }
        }

        // Put all entries back in the queue
        for entry in entries {
            self.requests.push(entry);
        }

        // Update the atomic counter with the accurate weight
        self.active_weight.store(active_weight, Ordering::SeqCst);

        if active_weight + weight <= self.max_weight {
            // We have capacity
            self.requests.push((now, weight));
            self.active_weight.fetch_add(weight, Ordering::SeqCst);
            true
        } else {
            // Rate limit would be exceeded
            false
        }
    }

    /// Marks a request with a specific weight as completed
    pub fn mark_request(&self, weight: usize) {
        let now = self.clock.raw() / 1_000_000; // Convert ns to ms

        // Add the entry to the queue
        self.requests.push((now, weight));

        // Increment the active weight
        self.active_weight.fetch_add(weight, Ordering::SeqCst);

        // Note: We don't need to remove old entries here since they'll be
        // filtered out when counting active weight in acquire/try_acquire
    }

    /// Returns the remaining weight available for immediate use
    pub fn remaining_weight(&self) -> usize {
        // Fast path: use the atomic counter for an approximate count
        let current_weight = self.active_weight.load(Ordering::Relaxed);

        // If we're well below the limit, we can return immediately
        if current_weight < self.max_weight / 2 {
            return self.max_weight - current_weight;
        }

        // Slow path: count active weight accurately
        let now = self.clock.raw() / 1_000_000; // Convert ns to ms
        let mut active_weight = 0;
        let mut entries = Vec::new();

        // Drain the queue temporarily to scan all entries
        while let Some(entry) = self.requests.pop() {
            entries.push(entry);

            // Check if this entry is still active
            let (timestamp, entry_weight) = entry;
            if now - timestamp <= self.window_milliseconds {
                active_weight += entry_weight;
            }
        }

        // Put all entries back in the queue
        for entry in entries {
            self.requests.push(entry);
        }

        // Update the atomic counter with the accurate weight
        self.active_weight.store(active_weight, Ordering::SeqCst);

        self.max_weight.saturating_sub(active_weight)
    }
}