Riven/src/req/rate_limit.rs

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

use parking_lot::{
    RwLock,
};

use super::token_bucket::{
    TokenBucket,
    VectorTokenBucket,
};
use super::rate_limit_type::RateLimitType;

pub struct RateLimit {
    rate_limit_type: RateLimitType,
    // Buckets for this rate limit (synchronized).
    // Almost always read, written only when rate limit rates are updated
    // from API response.
    // TODO: Question of writer starvation.
    buckets: RwLock<Vec<VectorTokenBucket>>,
    // Set to when we can retry if a retry-after header is received.
    retry_after: Option<Instant>,
}

impl RateLimit {
    /// Header specifying which RateLimitType caused a 429.
    const HEADER_XRATELIMITTYPE: &'static str = "X-Rate-Limit-Type";
    /// Header specifying retry after time in seconds after a 429.
    const HEADER_RETRYAFTER: &'static str = "Retry-After";

    pub fn new(rate_limit_type: RateLimitType) -> Self {
        let initial_bucket = VectorTokenBucket::new(Duration::from_secs(1), 1);
        RateLimit {
            rate_limit_type: rate_limit_type,
            // Rate limit before getting from response: 1/s.
            buckets: RwLock::new(vec![initial_bucket]),
            retry_after: None,
        }
    }

    pub fn get_both_or_delay(app_rate_limit: &Self, method_rate_limit: &Self) -> Option<Duration> {
        // Check retry after.
        let retry_after_delay = app_rate_limit.get_retry_after_delay()
            .and_then(|a| method_rate_limit.get_retry_after_delay().map(|m| cmp::max(a, m)));
        if retry_after_delay.is_some() {
            return retry_after_delay
        }
        // Check buckets.
        let app_buckets = app_rate_limit.buckets.read();
        let method_buckets = method_rate_limit.buckets.read();
        for bucket in app_buckets.iter().chain(method_buckets.iter()) {
            let delay = bucket.get_delay();
            if delay.is_some() {
                return delay;
            }
        }
        // Success.
        for bucket in app_buckets.iter().chain(method_buckets.iter()) {
            bucket.get_tokens(1);
        }
        None
    }

    pub fn get_retry_after_delay(&self) -> Option<Duration> {
        self.retry_after.and_then(|i| Instant::now().checked_duration_since(i))
    }

    pub fn on_response(&self, _response: &reqwest::Response) {
        unimplemented!();
    }
}

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

    fn send_sync() {
        fn is_send_sync<T: Send + Sync>() {}
        is_send_sync::<RateLimit>();
    }
}