forked from mirror/Riven
Increase robustness of 429 handling, fix #4
parent
39af464170
commit
9c8313c604
|
@ -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()
|
||||||
|
|
|
@ -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,
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue