forked from mirror/Riven
1
0
Fork 0

Increase robustness of 429 handling, fix #4

users/mingwei/unknown-variant-messy
Mingwei Samuel 2021-05-21 19:37:54 -07:00
parent 39af464170
commit 9c8313c604
2 changed files with 38 additions and 12 deletions

View File

@ -23,8 +23,6 @@ pub struct RateLimit {
impl RateLimit { impl RateLimit {
/// Header specifying which RateLimitType caused a 429. /// Header specifying which RateLimitType caused a 429.
const HEADER_XRATELIMITTYPE: &'static str = "X-Rate-Limit-Type"; 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 { pub fn new(rate_limit_type: RateLimitType) -> Self {
let initial_bucket = VectorTokenBucket::new( let initial_bucket = VectorTokenBucket::new(
@ -60,7 +58,7 @@ impl RateLimit {
bucket.get_tokens(1); bucket.get_tokens(1);
} }
log::debug!("Tokens obtained, buckets: APP {:?} METHOD {:?}", app_buckets, method_buckets); log::trace!("Tokens obtained, buckets: APP {:?} METHOD {:?}", app_buckets, method_buckets);
None None
} }
@ -83,22 +81,41 @@ impl RateLimit {
} }
{ {
// Only care if the header that indicates the relevant RateLimit is present. // Check if the header that indicates the relevant RateLimit is present.
let type_name_header = response.headers() let header_opt = response.headers()
.get(RateLimit::HEADER_XRATELIMITTYPE)?.to_str() .get(RateLimit::HEADER_XRATELIMITTYPE)
.expect("Failed to read x-rate-limit-type header as string."); .or_else(|| {
// Only care if that header's value matches us. log::info!("429 response missing {} header.", RateLimit::HEADER_XRATELIMITTYPE);
if self.rate_limit_type.type_name() != type_name_header.to_lowercase() { None
return None; })
.and_then(|header_value| header_value.to_str()
.map_err(|e| log::info!("429 response, failed to parse {} header as string: {}.",
RateLimit::HEADER_XRATELIMITTYPE, e))
.ok());
match header_opt {
None => {
// Take default responsibility, or ignore.
if !self.rate_limit_type.default_responsibility() {
return None;
}
log::warn!("429 response has missing or invalid {} header, {} rate limit taking responsibility.",
RateLimit::HEADER_XRATELIMITTYPE, self.rate_limit_type.type_name());
}
Some(header_value) => {
// Ignore if header's value does not match us.
if self.rate_limit_type.type_name() != header_value.to_lowercase() {
return None;
}
}
} }
} }
// Get retry after header. Only care if it exists. // Get retry after header. Only care if it exists.
let retry_after_header = response.headers() let retry_after_header = response.headers()
.get(RateLimit::HEADER_RETRYAFTER)?.to_str() .get(reqwest::header::RETRY_AFTER)?.to_str()
.expect("Failed to read retry-after header as string."); .expect("Failed to read retry-after header as string.");
log::debug!("Hit 429, retry-after {} secs.", retry_after_header); log::info!("429 response, rate limit {}, retry-after {} secs.", self.rate_limit_type.type_name(), retry_after_header);
// Header currently only returns ints, but float is more general. Can be zero. // Header currently only returns ints, but float is more general. Can be zero.
let retry_after_secs: f32 = retry_after_header.parse() let retry_after_secs: f32 = retry_after_header.parse()

View File

@ -25,4 +25,13 @@ impl RateLimitType {
Self::Method => "X-Method-Rate-Limit-Count", Self::Method => "X-Method-Rate-Limit-Count",
} }
} }
/// Return if this RateLimitType should take responsibility for responses
/// which are lacking a "X-Rate-Limit-Type" header.
pub fn default_responsibility(self) -> bool {
match self {
Self::Application => true,
Self::Method => false,
}
}
} }