#![allow(zero_ptr)]
pub use super::routing::Route;
use chrono::{DateTime, Utc};
use hyper::client::Response;
use hyper::header::Headers;
use hyper::status::StatusCode;
use internal::prelude::*;
use parking_lot::Mutex;
use std::{
collections::HashMap,
sync::Arc,
time::Duration,
str,
thread,
i64
};
use super::{HttpError, Request};
static mut OFFSET: Option<i64> = None;
lazy_static! {
pub static ref GLOBAL: Arc<Mutex<()>> = Arc::new(Mutex::new(()));
pub static ref ROUTES: Arc<Mutex<HashMap<Route, Arc<Mutex<RateLimit>>>>> = {
Arc::new(Mutex::new(HashMap::default()))
};
}
pub(super) fn perform(req: Request) -> Result<Response> {
loop {
let _ = GLOBAL.lock();
let (_, route, _) = req.route.deconstruct();
let bucket = Arc::clone(ROUTES
.lock()
.entry(route)
.or_insert_with(|| {
Arc::new(Mutex::new(RateLimit {
limit: i64::MAX,
remaining: i64::MAX,
reset: i64::MAX,
}))
}));
let mut lock = bucket.lock();
lock.pre_hook(&route);
let response = super::retry(&req)?;
if unsafe { OFFSET }.is_none() {
calculate_offset(response.headers.get_raw("date"));
}
if route == Route::None {
return Ok(response);
} else {
let redo = if response.headers.get_raw("x-ratelimit-global").is_some() {
let _ = GLOBAL.lock();
Ok(
if let Some(retry_after) = parse_header(&response.headers, "retry-after")? {
debug!("Ratelimited on route {:?} for {:?}ms", route, retry_after);
thread::sleep(Duration::from_millis(retry_after as u64));
true
} else {
false
},
)
} else {
lock.post_hook(&response, &route)
};
if !redo.unwrap_or(true) {
return Ok(response);
}
}
}
}
#[derive(Clone, Debug, Default)]
pub struct RateLimit {
pub limit: i64,
pub remaining: i64,
pub reset: i64,
}
impl RateLimit {
pub(crate) fn pre_hook(&mut self, route: &Route) {
if self.limit == 0 {
return;
}
let offset = unsafe { OFFSET }.unwrap_or(0);
let now = Utc::now().timestamp();
let current_time = now - offset;
if current_time > self.reset {
self.remaining = self.limit;
return;
}
let diff = (self.reset - current_time) as u64;
if self.remaining == 0 {
let delay = (diff * 1000) + 500;
debug!(
"Pre-emptive ratelimit on route {:?} for {:?}ms",
route,
delay
);
thread::sleep(Duration::from_millis(delay));
return;
}
self.remaining -= 1;
}
pub(crate) fn post_hook(&mut self, response: &Response, route: &Route) -> Result<bool> {
if let Some(limit) = parse_header(&response.headers, "x-ratelimit-limit")? {
self.limit = limit;
}
if let Some(remaining) = parse_header(&response.headers, "x-ratelimit-remaining")? {
self.remaining = remaining;
}
if let Some(reset) = parse_header(&response.headers, "x-ratelimit-reset")? {
self.reset = reset;
}
Ok(if response.status != StatusCode::TooManyRequests {
false
} else if let Some(retry_after) = parse_header(&response.headers, "retry-after")? {
debug!("Ratelimited on route {:?} for {:?}ms", route, retry_after);
thread::sleep(Duration::from_millis(retry_after as u64));
true
} else {
false
})
}
}
pub fn offset() -> Option<i64> {
unsafe { OFFSET }
}
fn calculate_offset(header: Option<&[Vec<u8>]>) {
let now = Utc::now().timestamp();
let header = header
.and_then(|h| h.get(0))
.and_then(|x| str::from_utf8(x).ok());
if let Some(date) = header {
let s = date.replace("GMT", "+0000");
let parsed = DateTime::parse_from_str(&s, "%a, %d %b %Y %T %z");
if let Ok(parsed) = parsed {
let offset = parsed.timestamp();
let diff = offset - now;
unsafe {
OFFSET = Some(diff);
debug!("[ratelimiting] Set the ratelimit offset to {}", diff);
}
}
}
}
fn parse_header(headers: &Headers, header: &str) -> Result<Option<i64>> {
headers.get_raw(header).map_or(Ok(None), |header| {
str::from_utf8(&header[0])
.map_err(|_| Error::Http(HttpError::RateLimitUtf8))
.and_then(|v| {
v.parse::<i64>()
.map(Some)
.map_err(|_| Error::Http(HttpError::RateLimitI64))
})
})
}