Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add initial header-based rate limit handling
- Loading branch information
1 parent
931dae5
commit b71d1e3
Showing
3 changed files
with
187 additions
and
39 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,123 @@ | ||
| 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 { | ||
| /// Check before issuing a request for the given URL. | ||
| 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(); | ||
| } | ||
| } | ||
|
|
||
| /// Update based on rate limit headers in the response for given URL. | ||
| /// Returns `true` if the request was rate limited and should be retried. | ||
| 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) { | ||
| // break out if uninitialized | ||
| if self.limit == 0 { return } | ||
|
|
||
| let difference = self.reset - get_time().sec; | ||
| if difference < 0 { | ||
| // If reset is apparently in the past, optimistically assume that | ||
| // the reset has occurred and we're good for the next three seconds | ||
| // or so. When the response comes back we will know for real. | ||
| self.reset += 3; | ||
| self.remaining = self.limit; | ||
| return | ||
| } | ||
|
|
||
| // if no requests remain, wait a bit | ||
| if self.remaining <= 0 { | ||
| // 900ms in case "difference" is off by 1 | ||
| let delay = difference as u64 * 1000 + 900; | ||
| warn!("pre-ratelimit: sleeping for {}ms", delay); | ||
| ::sleep_ms(delay); | ||
| return | ||
| } | ||
|
|
||
| // Deduct from our remaining requests. If a lot of requests are issued | ||
| // before any responses are received, this will mean we can still limit | ||
| // preemptively. | ||
| 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; // 100ms of leeway | ||
| warn!("429: sleeping for {}ms", delay); | ||
| ::sleep_ms(delay); | ||
| return Ok(true); // retry the request | ||
| } | ||
| } | ||
| 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) | ||
| } | ||
| } |