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

This commit is contained in:
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 {
/// 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";
/// `"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 {
let initial_bucket = VectorTokenBucket::new(
Duration::from_secs(1), 1, Duration::new(0, 0), 1.0, 1.0);
@ -101,33 +111,35 @@ 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()
.get(RateLimit::HEADER_XRATELIMITTYPE)
.get(Self::HEADER_XRATELIMITTYPE)
.or_else(|| {
log::info!("429 response missing {} header.", RateLimit::HEADER_XRATELIMITTYPE);
log::info!("429 response missing {} header.", Self::HEADER_XRATELIMITTYPE);
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))
.map_err(|e| log::info!("429 response, error parsing '{}' header as string: {}. Header value: {:#?}",
Self::HEADER_XRATELIMITTYPE, e, header_value))
.ok());
match header_opt {
None => {
// Take default responsibility, or ignore.
if !self.rate_limit_type.default_responsibility() {
// This match checks the valid possibilities. Otherwise returns none to ignore.
// `Application` handles "application", `Method` handles all other values.
let application_should_handle = match header_opt {
Some(Self::HEADER_XRATELIMITTYPE_APPLICATION) => true,
Some(Self::HEADER_XRATELIMITTYPE_METHOD | Self::HEADER_XRATELIMITTYPE_SERVICE) => false,
other => {
// Method handles unknown values.
log::warn!(
"429 response has None (missing or invalid) or unknown {} header value {:?}, {:?} rate limit obeying retry-after.",
Self::HEADER_XRATELIMITTYPE, other, self.rate_limit_type);
false
},
};
if (self.rate_limit_type == RateLimitType::Application) != application_should_handle {
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.
@ -135,7 +147,7 @@ impl RateLimit {
.get(reqwest::header::RETRY_AFTER)?.to_str()
.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.
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 {
Application,
Method,
}
impl RateLimitType {
pub fn type_name(self) -> &'static str {
match self {
Self::Application => "application",
Self::Method => "method",
}
}
pub fn limit_header(self) -> &'static str {
pub const fn limit_header(self) -> &'static str {
match self {
Self::Application => "X-App-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 {
Self::Application => "X-App-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,
}
}
}