forked from mirror/Riven
1
0
Fork 0

Update: method limits to handle all non-application 429s

users/mingwei/unknown-variant-messy
Mingwei Samuel 2021-09-06 18:49:22 -07:00
parent dafc181bbe
commit 9dcbc626c7
2 changed files with 39 additions and 40 deletions

View File

@ -29,8 +29,18 @@ pub struct RateLimit {
impl RateLimit { impl RateLimit {
/// Header specifying which RateLimitType caused a 429. /// Header specifying which RateLimitType caused a 429.
/// This header specifies which rate limit is violated in a 429 (if any).
/// There are three possible values, see [HEADER_XRATELIMITTYPE_APPLICATION],
/// [HEADER_XRATELIMITTYPE_METHOD], and [HEADER_XRATELIMITTYPE_SERVICE].
const HEADER_XRATELIMITTYPE: &'static str = "X-Rate-Limit-Type"; const HEADER_XRATELIMITTYPE: &'static str = "X-Rate-Limit-Type";
/// `"application"` - Entire app/key is rate limited due to violation.
const HEADER_XRATELIMITTYPE_APPLICATION: &'static str = "application";
/// `"method"` - App/key is rate limited on a specific method due to violation.
const HEADER_XRATELIMITTYPE_METHOD: &'static str = "method";
/// `"service"` - Service backend is rate-limiting (no violation).
const HEADER_XRATELIMITTYPE_SERVICE: &'static str = "service";
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(
Duration::from_secs(1), 1, Duration::new(0, 0), 1.0, 1.0); Duration::from_secs(1), 1, Duration::new(0, 0), 1.0, 1.0);
@ -101,32 +111,34 @@ impl RateLimit {
} }
{ {
// Check if the header that indicates the relevant RateLimit is present. // Get the X-Rate-Limit-Type header, `Some("application" | "method" | "service")` or `None`.
let header_opt = response.headers() let header_opt = response.headers()
.get(RateLimit::HEADER_XRATELIMITTYPE) .get(Self::HEADER_XRATELIMITTYPE)
.or_else(|| { .or_else(|| {
log::info!("429 response missing {} header.", RateLimit::HEADER_XRATELIMITTYPE); log::info!("429 response missing {} header.", Self::HEADER_XRATELIMITTYPE);
None None
}) })
.and_then(|header_value| header_value.to_str() .and_then(|header_value| header_value.to_str()
.map_err(|e| log::info!("429 response, failed to parse {} header as string: {}.", .map_err(|e| log::info!("429 response, error parsing '{}' header as string: {}. Header value: {:#?}",
RateLimit::HEADER_XRATELIMITTYPE, e)) Self::HEADER_XRATELIMITTYPE, e, header_value))
.ok()); .ok());
match header_opt {
None => { // This match checks the valid possibilities. Otherwise returns none to ignore.
// Take default responsibility, or ignore. // `Application` handles "application", `Method` handles all other values.
if !self.rate_limit_type.default_responsibility() { let application_should_handle = match header_opt {
return None; Some(Self::HEADER_XRATELIMITTYPE_APPLICATION) => true,
} Some(Self::HEADER_XRATELIMITTYPE_METHOD | Self::HEADER_XRATELIMITTYPE_SERVICE) => false,
log::warn!("429 response has missing or invalid {} header, {} rate limit taking responsibility.", other => {
RateLimit::HEADER_XRATELIMITTYPE, self.rate_limit_type.type_name()); // Method handles unknown values.
} log::warn!(
Some(header_value) => { "429 response has None (missing or invalid) or unknown {} header value {:?}, {:?} rate limit obeying retry-after.",
// Ignore if header's value does not match us. Self::HEADER_XRATELIMITTYPE, other, self.rate_limit_type);
if self.rate_limit_type.type_name() != header_value.to_lowercase() { false
return None; },
} };
}
if (self.rate_limit_type == RateLimitType::Application) != application_should_handle {
return None;
} }
} }
@ -135,7 +147,7 @@ impl RateLimit {
.get(reqwest::header::RETRY_AFTER)?.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::info!("429 response, rate limit {}, retry-after {} secs.", self.rate_limit_type.type_name(), retry_after_header); log::info!("429 response, rate limit {:?}, retry-after {} secs.", self.rate_limit_type, 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

@ -1,37 +1,24 @@
#[derive(Copy, Clone, PartialEq, Eq)] /// The type for a [RateLimit](super::RateLimit). Either a rate limit for the
/// entire app (`Application`) or for a specific method (`Method`).
/// Method rate limit will handle service violations.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum RateLimitType { pub enum RateLimitType {
Application, Application,
Method, Method,
} }
impl RateLimitType { impl RateLimitType {
pub fn type_name(self) -> &'static str { pub const fn limit_header(self) -> &'static str {
match self {
Self::Application => "application",
Self::Method => "method",
}
}
pub fn limit_header(self) -> &'static str {
match self { match self {
Self::Application => "X-App-Rate-Limit", Self::Application => "X-App-Rate-Limit",
Self::Method => "X-Method-Rate-Limit", Self::Method => "X-Method-Rate-Limit",
} }
} }
pub fn count_header(self) -> &'static str { pub const fn count_header(self) -> &'static str {
match self { match self {
Self::Application => "X-App-Rate-Limit-Count", Self::Application => "X-App-Rate-Limit-Count",
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,
}
}
} }