neighborly 0.0.1

mod tuning;
mod util;

use serde::{Deserialize, Serialize};
use std::{cmp::Ordering, collections::VecDeque, num::NonZeroUsize, sync::Mutex, time::Duration};
use thiserror::Error;
use tokio::time::{sleep_until, Instant};

/// Object-safe subset of the rate limiter's functionality.
pub trait RateLimiterCore {
    /// Schedule a request to run. If it may run immediately, returns `None`.
    /// Otherwise, returns `Some` with a future time at which the request has been
    /// scheduled to run. The caller should sleep until that time.
    ///
    /// Unless you have a specific need for this function, use `wait_until_ready()`
    /// instead.
    fn schedule(&self) -> Option<Instant>;
}

/// If you need to use dynamic dispatch with a `RateLimiter`, see
/// `RateLimiterCore`.
///
/// This is auto-implemented on any type implementing `RateLimiterCore`.
#[allow(async_fn_in_trait)]
pub trait RateLimiter: RateLimiterCore {
    /// Schedule a request & sleep until the designated time.
    async fn wait_until_ready(&self) {
        if let Some(i) = <Self as RateLimiterCore>::schedule(&self) {
            sleep_until(i).await
        }
    }
}

impl<T: RateLimiterCore> RateLimiter for T {}

#[derive(Debug)]
/// A simple rate limiter based on the sliding window algorithm.
pub struct SimpleRateLimiter {
    /// Window of time over which rates are calculated
    window: Duration,
    /// Minimum interval between requests.
    burst_interval: Duration,
    /// Maximum number of requests in a given window
    limit: usize,
    /// Average interval of a rate limited stream of requests
    steady_interval: Duration,
    /// Timestamps of scheduled requests. May be in the future.
    history: Mutex<VecDeque<Instant>>,
}
impl SimpleRateLimiter {
    /// Create a new instance of a rate limiter.
    pub fn new(window: Duration, burst_interval: Duration, limit: NonZeroUsize) -> Self {
        Self::with_capacity(window, burst_interval, limit, 2 * limit.get())
    }
    /// Create a new instance of a rate limiter, using the given capacity for the
    /// underlying VecDeque<Instant>.
    pub fn with_capacity(
        window: Duration,
        burst_interval: Duration,
        limit: NonZeroUsize,
        capacity: usize,
    ) -> Self {
        assert!(window.as_nanos() > 0, "Window must not be empty");
        assert!(
            burst_interval.as_nanos() > 0,
            "Burst interval must not be empty"
        );
        let limit = limit.get();
        let interval = Duration::from_secs_f32(window.as_secs_f32() / (limit as f32));

        Self {
            window,
            steady_interval: interval,
            burst_interval,
            limit,
            history: Mutex::new(VecDeque::with_capacity(capacity)),
        }
    }
    pub fn builder() -> SimpleRateLimiterBuilder {
        Default::default()
    }
    /// Remove any expired timestamps from the history.
    fn clean_history(&self, history: &mut VecDeque<Instant>, now: Instant) {
        while history.front().map_or(false, |oldest| {
            now.saturating_duration_since(*oldest) > self.window
        }) {
            history.pop_front();
        }
    }
}
impl RateLimiterCore for SimpleRateLimiter {
    fn schedule(&self) -> Option<Instant> {
        let mut history = self
            .history
            .lock()
            .expect("rate limiter mutex was poisoned");

        let now = Instant::now();
        self.clean_history(&mut history, now);

        match history.len().cmp(&self.limit) {
            Ordering::Less => {
                // We have not yet reached our rate limit. Schedule incoming requests at the
                // burst rate.
                if let Some(oldest) = history.back() {
                    let scheduled_time =
                        now.max(*oldest + self.burst_interval + Duration::from_nanos(1));
                    history.push_back(scheduled_time);
                    if scheduled_time > now {
                        Some(scheduled_time)
                    } else {
                        None
                    }
                } else {
                    history.push_back(now);
                    None
                }
            }
            Ordering::Equal => {
                // We have reached our rate limit. Schedule when the oldest request expires
                // modulo the burst limit.
                let oldest = history
                    .front()
                    .expect("limit is nonzero, therefore history is not empty");
                let youngest = history
                    .back()
                    .expect("limit is nonzero, therefore history is not empty");
                let oldest_expiry = *oldest + self.window + Duration::from_nanos(1);
                let burst_limit_expiry = *youngest + self.burst_interval + Duration::from_nanos(1);
                let scheduled_time = oldest_expiry.max(burst_limit_expiry);

                history.push_back(scheduled_time);
                Some(scheduled_time)
            }
            Ordering::Greater => {
                // We have saturated our rate limit. Schedule requests at regular intervals.
                let youngest = history
                    .back()
                    .expect("limit is nonzero, therefore history is not empty");
                let scheduled_time = *youngest + self.steady_interval;

                history.push_back(scheduled_time);
                Some(scheduled_time)
            }
        }
    }
}
impl Default for SimpleRateLimiter {
    fn default() -> Self {
        Self::builder()
            .window(Duration::from_secs(1))
            .limit(10)
            .burst_factor(2.)
            .build()
    }
}

#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct SimpleRateLimiterBuilder {
    /// Required.
    pub window: Option<Duration>,
    /// Required.
    pub burst_config: Option<BurstConfig>,
    /// Required.
    pub limit: Option<usize>,
    /// Optional.
    pub capacity: Option<usize>,
}
impl SimpleRateLimiterBuilder {
    pub fn new() -> Self {
        Self {
            window: None,
            burst_config: None,
            limit: None,
            capacity: None,
        }
    }
    pub fn window(self, window: Duration) -> Self {
        Self {
            window: Some(window),
            ..self
        }
    }
    pub fn limit(self, limit: usize) -> Self {
        Self {
            limit: Some(limit),
            ..self
        }
    }
    pub fn burst_interval(self, interval: Duration) -> Self {
        Self {
            burst_config: Some(BurstConfig::Interval(interval)),
            ..self
        }
    }
    pub fn burst_factor(self, factor: f32) -> Self {
        Self {
            burst_config: Some(BurstConfig::Factor(factor)),
            ..self
        }
    }
    pub fn try_build(&self) -> Result<SimpleRateLimiter, BuilderError> {
        // Ensure all required values are supplied
        let (window, burst_config, limit) = match (&self.window, &self.burst_config, &self.limit) {
            (Some(w), Some(b), Some(l)) => (w, b, l),
            (None, None, None) => Err(BuilderError::MissingAll)?,
            (None, _, _) => Err(BuilderError::MissingWindow)?,
            (_, None, _) => Err(BuilderError::MissingBurstConfig)?,
            (_, _, None) => Err(BuilderError::MissingLimit)?,
        };

        // Ensure all supplied values are valid
        let limit_nz: NonZeroUsize;
        if let Some(l) = NonZeroUsize::new(*limit) {
            limit_nz = l;
        } else {
            return Err(BuilderError::LimitIsZero);
        }
        if window.as_nanos() == 0 {
            Err(BuilderError::WindowIsZero)?
        }

        Ok(SimpleRateLimiter::with_capacity(
            *window,
            burst_config.burst_interval(*window, *limit)?,
            limit_nz,
            self.capacity.unwrap_or(2 * limit),
        ))
    }
    pub fn build(&self) -> SimpleRateLimiter {
        self.try_build().expect("Failed to build SimpleRateLimiter")
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum BurstConfig {
    Factor(f32),
    Interval(Duration),
}
impl BurstConfig {
    pub fn burst_interval(&self, window: Duration, limit: usize) -> Result<Duration, BuilderError> {
        match self {
            BurstConfig::Factor(f) => {
                if *f > 1. {
                    Ok(Duration::from_secs_f32(
                        window.as_secs_f32() / (limit as f32) / f,
                    ))
                } else {
                    Err(BuilderError::BurstFactorIsLessThanOrEqualToOne)
                }
            }
            BurstConfig::Interval(d) => {
                if d.as_nanos() > 0 {
                    Ok(*d)
                } else {
                    Err(BuilderError::BurstIntervalIsZero)
                }
            }
        }
    }
}

#[derive(Clone, Copy, Debug, Error)]
pub enum BuilderError {
    #[error("All required values are missing")]
    MissingAll,
    #[error("Missing required attrbiute window")]
    MissingWindow,
    #[error("Missing required attrbiute burst_config")]
    MissingBurstConfig,
    #[error("Missing required attrbiute limit")]
    MissingLimit,
    #[error("Window cannot be zero")]
    WindowIsZero,
    #[error("Burst interval cannot be zero")]
    BurstIntervalIsZero,
    #[error("Limit cannot be zero")]
    LimitIsZero,
    #[error("Burst factor must be > 1")]
    BurstFactorIsLessThanOrEqualToOne,
}

#[cfg(test)]
mod test {
    use super::*;
    use tokio::time::{advance, pause};

    #[test]
    fn core_is_objsafe() {
        // If this function compiles, then we are object safe.
        let _: Option<Box<dyn RateLimiterCore>> = None;
    }

    #[tokio::test]
    async fn first_request_is_scheduled_immediately() {
        let rl = SimpleRateLimiter::default();
        assert_eq!(rl.schedule(), None)
    }

    #[tokio::test]
    async fn second_request_is_scheduled_after_burst_interval() {
        let rl = SimpleRateLimiter::default();
        assert_eq!(rl.schedule(), None)
    }

    #[tokio::test]
    async fn bursting_to_limit() {
        let window = Duration::from_secs(1);
        let limit = 5;
        let burst_factor = 2.;
        let burst_interval =
            Duration::from_secs_f32(window.as_secs_f32() / (limit as f32) / burst_factor);
        let mut results: Vec<Instant> = Vec::with_capacity(limit);

        pause();
        let rl = SimpleRateLimiter::builder()
            .window(window)
            .limit(limit)
            .burst_factor(burst_factor)
            .build();

        // Burst to limit with requests every nanosecond, recording when the request was scheduled
        // to run
        for _ in 0..limit {
            results.push(rl.schedule().unwrap_or(Instant::now()));
            advance(Duration::from_nanos(1)).await;
        }

        assert_eq!(rl.len(), limit);

        // Requests were scheduled at burst rate
        for i in 1..limit {
            let scheduled_time = results[i];
            let prev_scheduled_time = results[i - 1];

            assert_eq!(
                scheduled_time
                    .saturating_duration_since(prev_scheduled_time)
                    .as_millis(),
                burst_interval.as_millis()
            );
        }
    }

    #[tokio::test]
    async fn expiring_single_request() {
        let window = Duration::from_secs(1);
        let limit = 5;
        let burst_factor = 2.;

        pause();
        let rl = SimpleRateLimiter::builder()
            .window(window)
            .limit(limit)
            .burst_factor(burst_factor)
            .build();

        assert!(rl.schedule().is_none());
        assert_eq!(rl.len(), 1);

        advance(window + Duration::from_nanos(1)).await;
        assert_eq!(rl.len(), 0);
    }

    #[tokio::test]
    async fn expiring_single_request_after_burst() {
        let window = Duration::from_secs(1);
        let limit = 5;
        let burst_factor = 2.;

        pause();

        let rl = SimpleRateLimiter::builder()
            .window(window)
            .limit(limit)
            .burst_factor(burst_factor)
            .build();

        assert!(rl.schedule().is_none());
        assert_eq!(rl.len(), 1);

        advance(window / 2).await;
        for i in 1..limit {
            assert!(rl.schedule().is_some() || i == 1);
            advance(Duration::from_nanos(1)).await;
        }

        assert_eq!(rl.len(), limit);

        advance(window / 2).await;
        assert_eq!(rl.len(), limit - 1);
    }

    #[tokio::test]
    async fn expiring_burst() {
        let window = Duration::from_secs(1);
        let limit = 5;
        let burst_factor = 2.;
        let burst_interval =
            Duration::from_secs_f32(window.as_secs_f32() / (limit as f32) / burst_factor);

        pause();
        let start = Instant::now();

        let rl = SimpleRateLimiter::builder()
            .window(window)
            .limit(limit)
            .burst_factor(burst_factor)
            .build();

        for i in 0..limit {
            assert!(rl.schedule().is_some() || i == 0);
            advance(Duration::from_nanos(1)).await;
        }

        assert_eq!(rl.len(), limit);
        let last_scheduled_time = start + (burst_interval * limit as u32);
        let expiry = last_scheduled_time + window + Duration::from_nanos(1);

        advance(last_scheduled_time.saturating_duration_since(Instant::now())).await;
        assert_eq!(rl.len(), limit);

        advance(expiry.saturating_duration_since(Instant::now())).await;
        assert_eq!(rl.len(), 0);
    }

    #[tokio::test]
    async fn bursting_to_oversaturation() {
        let window = Duration::from_secs(1);
        let limit = 5;
        let burst_factor = 2.;
        let interval = Duration::from_secs_f32(window.as_secs_f32() / (limit as f32));
        let mut results: Vec<Instant> = Vec::with_capacity(2 * limit);

        pause();

        let rl = SimpleRateLimiter::builder()
            .window(window)
            .limit(limit)
            .burst_factor(burst_factor)
            .build();

        for _ in 0..(2 * limit) {
            results.push(rl.schedule().unwrap_or(Instant::now()));
            advance(Duration::from_nanos(1)).await;
        }
        assert_eq!(rl.len(), 2 * limit);

        for i in (limit + 1)..(2 * limit) {
            let scheduled_time = results[i];
            let prev_scheduled_time = results[i - 1];
            assert_eq!(
                scheduled_time
                    .saturating_duration_since(prev_scheduled_time)
                    .as_millis(),
                interval.as_millis()
            );
        }

        let last_scheduled_time = *results.last().unwrap();
        let expiry = last_scheduled_time + window + Duration::from_nanos(1);

        advance(last_scheduled_time.saturating_duration_since(Instant::now())).await;
        assert!(rl.len() > 0);
        advance(expiry.saturating_duration_since(Instant::now())).await;
        assert_eq!(rl.len(), 0);
    }
}