mirror of https://github.com/MingweiSamuel/Riven
Update: method limits to handle all non-application 429s
parent
dafc181bbe
commit
9dcbc626c7
|
@ -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()
|
||||
|
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue