use std::sync::Mutex;
use std::collections::BTreeMap;
use std;
use hyper;
use time::get_time;
use {Result, Error};
#[derive(Default)]
pub struct RateLimits {
global: Mutex<RateLimit>,
endpoints: Mutex<BTreeMap<String, RateLimit>>,
}
impl RateLimits {
pub fn pre_check(&self, url: &str) {
self.global.lock().expect("Rate limits poisoned").pre_check();
if let Some(rl) = self.endpoints.lock().expect("Rate limits poisoned").get_mut(url) {
rl.pre_check();
}
}
pub fn post_update(&self, url: &str, response: &hyper::client::Response) -> bool {
if response.headers.get_raw("X-RateLimit-Global").is_some() {
self.global.lock().expect("Rate limits poisoned").post_update(response)
} else {
self.endpoints.lock().expect("Rate limits poisoned")
.entry(url.to_owned())
.or_insert_with(RateLimit::default)
.post_update(response)
}
}
}
#[derive(Default)]
struct RateLimit {
reset: i64,
limit: i64,
remaining: i64,
}
impl RateLimit {
fn pre_check(&mut self) {
if self.limit == 0 { return }
let difference = self.reset - get_time().sec;
if difference < 0 {
self.reset += 3;
self.remaining = self.limit;
return
}
if self.remaining <= 0 {
let delay = difference as u64 * 1000 + 900;
warn!("pre-ratelimit: sleeping for {}ms", delay);
::sleep_ms(delay);
return
}
self.remaining -= 1;
}
fn post_update(&mut self, response: &hyper::client::Response) -> bool {
match self.try_post_update(response) {
Err(e) => {
error!("rate limit checking error: {}", e);
false
}
Ok(r) => r
}
}
fn try_post_update(&mut self, response: &hyper::client::Response) -> Result<bool> {
if let Some(reset) = try!(read_header(&response.headers, "X-RateLimit-Reset")) {
self.reset = reset;
}
if let Some(limit) = try!(read_header(&response.headers, "X-RateLimit-Limit")) {
self.limit = limit;
}
if let Some(remaining) = try!(read_header(&response.headers, "X-RateLimit-Remaining")) {
self.remaining = remaining;
}
if response.status == hyper::status::StatusCode::TooManyRequests {
if let Some(delay) = try!(read_header(&response.headers, "Retry-After")) {
let delay = delay as u64 + 100; warn!("429: sleeping for {}ms", delay);
::sleep_ms(delay);
return Ok(true); }
}
Ok(false)
}
}
fn read_header(headers: &hyper::header::Headers, name: &str) -> Result<Option<i64>> {
match headers.get_raw(name) {
Some(hdr) => if hdr.len() == 1 {
match std::str::from_utf8(&hdr[0]) {
Ok(text) => match text.parse() {
Ok(val) => Ok(Some(val)),
Err(_) => Err(Error::Other("header is not an i64"))
},
Err(_) => Err(Error::Other("header is not UTF-8"))
}
} else {
Err(Error::Other("header appears multiple times"))
},
None => Ok(None)
}
}