rusty_feeder/limit/

helpers.rs

//! Common helpers for rate limiter implementations
//!
//! This module provides shared utility functions for timestamp cleanup, token refill calculations,
//! and other common rate limiting operations to reduce code duplication across different
//! rate limiter implementations.

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

/// Calculate the number of tokens to refill based on elapsed time
///
/// This function implements the standard proportional token refill algorithm
/// used across all rate limiter implementations. It calculates how many tokens
/// should be available based on the time elapsed since the last refresh.
///
/// # Arguments
/// * `elapsed` - Time elapsed since last refresh
/// * `window_duration` - Duration of the rate limiting window
/// * `max_tokens` - Maximum number of tokens in the bucket
/// * `current_tokens` - Current number of available tokens
///
/// # Returns
/// The new number of available tokens (capped at max_tokens)
#[inline]
pub fn calculate_token_refill(
    elapsed: Duration,
    window_duration: Duration,
    max_tokens: usize,
    current_tokens: usize,
) -> usize {
    if elapsed >= window_duration {
        // Full window has passed, reset to maximum
        max_tokens
    } else {
        // Partial refill based on elapsed time
        let refill_ratio = elapsed.as_secs_f64() / window_duration.as_secs_f64();
        let tokens_to_add = (refill_ratio * max_tokens as f64) as usize;
        (current_tokens + tokens_to_add).min(max_tokens)
    }
}

/// Calculate the proportional time adjustment for partial token refill
///
/// When tokens are partially refilled, the last refresh time needs to be adjusted
/// proportionally to maintain accurate rate limiting.
///
/// # Arguments
/// * `last_refresh` - The last refresh time
/// * `now` - Current time
/// * `elapsed` - Time elapsed since last refresh
/// * `window_duration` - Duration of the rate limiting window
/// * `tokens_added` - Number of tokens that were added
/// * `max_tokens` - Maximum number of tokens in the bucket
///
/// # Returns
/// The adjusted last refresh time
#[inline]
pub fn calculate_proportional_time_adjustment(
    last_refresh: Instant,
    now: Instant,
    elapsed: Duration,
    window_duration: Duration,
    tokens_added: usize,
    max_tokens: usize,
) -> Instant {
    if tokens_added == 0 || max_tokens == 0 {
        return last_refresh;
    }

    // Calculate the time represented by the tokens added
    let token_ratio = tokens_added as f64 / max_tokens as f64;
    let time_for_tokens = Duration::from_secs_f64(window_duration.as_secs_f64() * token_ratio);

    // Adjust the last refresh time forward by the time represented by tokens added
    last_refresh + time_for_tokens
}

/// Clean up timestamps older than the specified cutoff time
///
/// This function removes all timestamps from a collection that are older than
/// the cutoff time, helping to prevent unbounded memory growth in rate limiters
/// that track individual request timestamps.
///
/// # Arguments
/// * `timestamps` - Mutable reference to a vector of timestamps
/// * `cutoff` - The cutoff instant; timestamps older than this will be removed
///
/// # Returns
/// The number of timestamps that were removed
#[inline]
pub fn cleanup_old_timestamps(timestamps: &mut Vec<Instant>, cutoff: Instant) -> usize {
    let initial_len = timestamps.len();

    // Remove timestamps older than cutoff
    timestamps.retain(|&timestamp| timestamp >= cutoff);

    initial_len - timestamps.len()
}

/// Clean up timestamps older than the specified duration from now
///
/// Convenience function that calculates the cutoff time based on the current time
/// and a duration, then calls cleanup_old_timestamps.
///
/// # Arguments
/// * `timestamps` - Mutable reference to a vector of timestamps
/// * `window_duration` - Duration of the rate limiting window
/// * `now` - Current time instant
///
/// # Returns
/// The number of timestamps that were removed
#[inline]
pub fn cleanup_timestamps_by_duration(
    timestamps: &mut Vec<Instant>,
    window_duration: Duration,
    now: Instant,
) -> usize {
    let cutoff = now - window_duration;
    cleanup_old_timestamps(timestamps, cutoff)
}

/// Record a new timestamp and perform cleanup if needed
///
/// This function adds a new timestamp to the collection and optionally performs
/// cleanup of old timestamps to maintain memory efficiency.
///
/// # Arguments
/// * `timestamps` - Mutable reference to a vector of timestamps
/// * `timestamp` - The new timestamp to record
/// * `window_duration` - Duration of the rate limiting window
/// * `max_capacity` - Optional maximum capacity for the timestamp vector
/// * `auto_cleanup` - Whether to automatically cleanup old timestamps
///
/// # Returns
/// The number of timestamps that were cleaned up (if any)
#[inline]
pub fn record_timestamp_with_cleanup(
    timestamps: &mut Vec<Instant>,
    timestamp: Instant,
    window_duration: Duration,
    max_capacity: Option<usize>,
    auto_cleanup: bool,
) -> usize {
    timestamps.push(timestamp);

    let mut cleaned_count = 0;

    // Perform cleanup if requested
    if auto_cleanup {
        cleaned_count += cleanup_timestamps_by_duration(timestamps, window_duration, timestamp);
    }

    // Enforce capacity limit if specified
    if let Some(max_cap) = max_capacity
        && timestamps.len() > max_cap
    {
        let excess = timestamps.len() - max_cap;
        timestamps.drain(0..excess);
        cleaned_count += excess;
    }

    cleaned_count
}

/// Check if enough time has passed for a full window reset
///
/// This is a common check across rate limiters to determine if the window
/// has completely elapsed and tokens should be reset to maximum.
///
/// # Arguments
/// * `elapsed` - Time elapsed since last refresh
/// * `window_duration` - Duration of the rate limiting window
///
/// # Returns
/// True if a full window reset should occur
#[inline]
pub fn should_perform_full_reset(elapsed: Duration, window_duration: Duration) -> bool {
    elapsed >= window_duration
}

/// Calculate token refill for nanosecond-based rate limiters
///
/// Specialized version for rate limiters that work with nanosecond timestamps
/// instead of Duration/Instant types.
///
/// # Arguments
/// * `elapsed_ns` - Nanoseconds elapsed since last refresh
/// * `window_ns` - Duration of the rate limiting window in nanoseconds
/// * `max_tokens` - Maximum number of tokens in the bucket
/// * `current_tokens` - Current number of available tokens
///
/// # Returns
/// The new number of available tokens (capped at max_tokens)
#[inline]
pub fn calculate_token_refill_ns(
    elapsed_ns: u64,
    window_ns: u64,
    max_tokens: usize,
    current_tokens: usize,
) -> usize {
    if elapsed_ns >= window_ns {
        // Full window has passed, reset to maximum
        max_tokens
    } else if elapsed_ns > 0 {
        // Partial refill based on elapsed time
        let refill_ratio = elapsed_ns as f64 / window_ns as f64;
        let tokens_to_add = (refill_ratio * max_tokens as f64) as usize;
        (current_tokens + tokens_to_add).min(max_tokens)
    } else {
        current_tokens
    }
}

/// Calculate proportional time adjustment for nanosecond-based rate limiters
///
/// # Arguments
/// * `last_refill_ns` - The last refresh time in nanoseconds
/// * `elapsed_ns` - Nanoseconds elapsed since last refresh
/// * `window_ns` - Duration of the rate limiting window in nanoseconds
/// * `tokens_added` - Number of tokens that were added
/// * `max_tokens` - Maximum number of tokens in the bucket
///
/// # Returns
/// The adjusted last refresh time in nanoseconds
#[inline]
pub const fn calculate_proportional_time_adjustment_ns(
    last_refill_ns: u64,
    elapsed_ns: u64,
    window_ns: u64,
    tokens_added: usize,
    max_tokens: usize,
) -> u64 {
    if tokens_added == 0 || max_tokens == 0 {
        return last_refill_ns;
    }

    // Calculate the time represented by the tokens added
    let time_for_tokens = (elapsed_ns * tokens_added as u64) / max_tokens as u64;

    // Adjust the last refresh time forward by the time represented by tokens added
    last_refill_ns + time_for_tokens
}

/// Check if enough time has passed for a full window reset (nanosecond version)
///
/// # Arguments
/// * `elapsed_ns` - Nanoseconds elapsed since last refresh
/// * `window_ns` - Duration of the rate limiting window in nanoseconds
///
/// # Returns
/// True if a full window reset should occur
#[inline]
pub const fn should_perform_full_reset_ns(elapsed_ns: u64, window_ns: u64) -> bool {
    elapsed_ns >= window_ns
}

/// Calculate tokens available after time-based refill
///
/// Higher-level function that combines token calculation and proportional time adjustment.
/// This is the main function used by rate limiters to update their state.
///
/// # Arguments
/// * `current_tokens` - Current number of available tokens
/// * `max_tokens` - Maximum number of tokens in the bucket
/// * `last_refresh` - The last refresh time
/// * `now` - Current time
/// * `window_duration` - Duration of the rate limiting window
///
/// # Returns
/// A tuple of (new_token_count, new_last_refresh_time)
#[inline]
pub fn refill_tokens_with_time_adjustment(
    current_tokens: usize,
    max_tokens: usize,
    last_refresh: Instant,
    now: Instant,
    window_duration: Duration,
) -> (usize, Instant) {
    let elapsed = now.saturating_duration_since(last_refresh);

    if should_perform_full_reset(elapsed, window_duration) {
        // Full reset
        (max_tokens, now)
    } else {
        // Partial refill
        let new_tokens =
            calculate_token_refill(elapsed, window_duration, max_tokens, current_tokens);
        let tokens_added = new_tokens.saturating_sub(current_tokens);

        let new_last_refresh = if tokens_added > 0 {
            calculate_proportional_time_adjustment(
                last_refresh,
                now,
                elapsed,
                window_duration,
                tokens_added,
                max_tokens,
            )
        } else {
            last_refresh
        };

        (new_tokens, new_last_refresh)
    }
}

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

    #[test]
    fn test_calculate_token_refill_full_window() {
        let window = Duration::from_secs(1);
        let elapsed = Duration::from_secs(2); // More than full window
        let max_tokens = 100;
        let current_tokens = 50;

        let result = calculate_token_refill(elapsed, window, max_tokens, current_tokens);
        assert_eq!(result, max_tokens);
    }

    #[test]
    fn test_calculate_token_refill_partial() {
        let window = Duration::from_secs(1);
        let elapsed = Duration::from_millis(500); // Half window
        let max_tokens = 100;
        let current_tokens = 0;

        let result = calculate_token_refill(elapsed, window, max_tokens, current_tokens);
        assert_eq!(result, 50); // Half the tokens
    }

    #[test]
    fn test_calculate_token_refill_cap_at_max() {
        let window = Duration::from_secs(1);
        let elapsed = Duration::from_millis(500); // Half window
        let max_tokens = 100;
        let current_tokens = 80; // Already near max

        let result = calculate_token_refill(elapsed, window, max_tokens, current_tokens);
        assert_eq!(result, max_tokens); // Capped at max
    }

    #[test]
    fn test_cleanup_old_timestamps() {
        let now = Instant::now();
        let old_time = now - Duration::from_secs(2);
        let recent_time = now - Duration::from_millis(500);

        let mut timestamps = vec![old_time, recent_time, now];
        let cutoff = now - Duration::from_secs(1);

        let removed_count = cleanup_old_timestamps(&mut timestamps, cutoff);

        assert_eq!(removed_count, 1);
        assert_eq!(timestamps.len(), 2);
        assert!(!timestamps.contains(&old_time));
        assert!(timestamps.contains(&recent_time));
        assert!(timestamps.contains(&now));
    }

    #[test]
    fn test_record_timestamp_with_cleanup() {
        let now = Instant::now();
        let old_time = now - Duration::from_secs(2);

        let mut timestamps = vec![old_time];
        let window = Duration::from_secs(1);

        let cleaned_count = record_timestamp_with_cleanup(&mut timestamps, now, window, None, true);

        assert_eq!(cleaned_count, 1); // Old timestamp was cleaned
        assert_eq!(timestamps.len(), 1);
        assert!(timestamps.contains(&now));
        assert!(!timestamps.contains(&old_time));
    }

    #[test]
    fn test_refill_tokens_with_time_adjustment_full_reset() {
        let now = Instant::now();
        let last_refresh = now - Duration::from_secs(2);
        let window = Duration::from_secs(1);
        let max_tokens = 100;
        let current_tokens = 50;

        let (new_tokens, new_last_refresh) = refill_tokens_with_time_adjustment(
            current_tokens,
            max_tokens,
            last_refresh,
            now,
            window,
        );

        assert_eq!(new_tokens, max_tokens);
        assert_eq!(new_last_refresh, now);
    }

    #[test]
    fn test_refill_tokens_with_time_adjustment_partial() {
        let now = Instant::now();
        let last_refresh = now - Duration::from_millis(500);
        let window = Duration::from_secs(1);
        let max_tokens = 100;
        let current_tokens = 0;

        let (new_tokens, _new_last_refresh) = refill_tokens_with_time_adjustment(
            current_tokens,
            max_tokens,
            last_refresh,
            now,
            window,
        );

        assert_eq!(new_tokens, 50); // Half the tokens for half the time
    }
}