From 238e0f6318a7e290ad34459f1fd191088b70f55b Mon Sep 17 00:00:00 2001 From: Jiahao XU Date: Tue, 11 Jun 2024 22:37:48 +1000 Subject: [PATCH] Fix rate limit checking in `GhApiClient` (#1725) * Fix rate limit checking in `GhApiClient` - Mv logic into `binstalk_downloader` - Check for `RETRY_AFTER` and `x-ratelimit-remaining` on any status code Signed-off-by: Jiahao XU * Fix clippy lint Signed-off-by: Jiahao XU --------- Signed-off-by: Jiahao XU --- crates/binstalk-downloader/src/remote.rs | 47 +++++++++++++------ .../src/gh_api_client.rs | 7 ++- .../src/gh_api_client/common.rs | 29 ++---------- 3 files changed, 40 insertions(+), 43 deletions(-) diff --git a/crates/binstalk-downloader/src/remote.rs b/crates/binstalk-downloader/src/remote.rs index 65a43d53..7bd2c684 100644 --- a/crates/binstalk-downloader/src/remote.rs +++ b/crates/binstalk-downloader/src/remote.rs @@ -9,7 +9,7 @@ use bytes::Bytes; use futures_util::Stream; use httpdate::parse_http_date; use reqwest::{ - header::{HeaderMap, RETRY_AFTER}, + header::{HeaderMap, HeaderValue, RETRY_AFTER}, Request, }; use thiserror::Error as ThisError; @@ -172,6 +172,8 @@ impl Client { url: &Url, ) -> Result>, ReqwestError> { + static HEADER_VALUE_0: HeaderValue = HeaderValue::from_static("0"); + let response = match self.0.service.call(request).await { Err(err) if err.is_timeout() || err.is_connect() => { let duration = RETRY_DURATION_FOR_TIMEOUT; @@ -197,22 +199,37 @@ impl Client { Ok(ControlFlow::Continue(Ok(response))) }; - match status { - // Delay further request on rate limit - StatusCode::SERVICE_UNAVAILABLE | StatusCode::TOO_MANY_REQUESTS => { - let duration = parse_header_retry_after(response.headers()) - .unwrap_or(DEFAULT_RETRY_DURATION_FOR_RATE_LIMIT) - .min(MAX_RETRY_DURATION); + let headers = response.headers(); - add_delay_and_continue(response, duration) + // Some server (looking at you, github GraphQL API) may returns a rate limit + // even when OK is returned or on other status code (e.g. 453 forbidden). + if let Some(duration) = parse_header_retry_after(headers) { + add_delay_and_continue(response, duration.min(MAX_RETRY_DURATION)) + } else if headers.get("x-ratelimit-remaining") == Some(&HEADER_VALUE_0) { + let duration = headers + .get("x-ratelimit-reset") + .and_then(|value| { + let secs = value.to_str().ok()?.parse().ok()?; + Some(Duration::from_secs(secs)) + }) + .unwrap_or(DEFAULT_RETRY_DURATION_FOR_RATE_LIMIT) + .min(MAX_RETRY_DURATION); + + add_delay_and_continue(response, duration) + } else { + match status { + // Delay further request on rate limit + StatusCode::SERVICE_UNAVAILABLE | StatusCode::TOO_MANY_REQUESTS => { + add_delay_and_continue(response, DEFAULT_RETRY_DURATION_FOR_RATE_LIMIT) + } + + // Delay further request on timeout + StatusCode::REQUEST_TIMEOUT | StatusCode::GATEWAY_TIMEOUT => { + add_delay_and_continue(response, RETRY_DURATION_FOR_TIMEOUT) + } + + _ => Ok(ControlFlow::Break(response)), } - - // Delay further request on timeout - StatusCode::REQUEST_TIMEOUT | StatusCode::GATEWAY_TIMEOUT => { - add_delay_and_continue(response, RETRY_DURATION_FOR_TIMEOUT) - } - - _ => Ok(ControlFlow::Break(response)), } } diff --git a/crates/binstalk-git-repo-api/src/gh_api_client.rs b/crates/binstalk-git-repo-api/src/gh_api_client.rs index 693bc3a2..ac7f1565 100644 --- a/crates/binstalk-git-repo-api/src/gh_api_client.rs +++ b/crates/binstalk-git-repo-api/src/gh_api_client.rs @@ -299,14 +299,13 @@ impl GhApiClient { .send(false) .await?; - match check_http_status_and_header(&response) { + match check_http_status_and_header(response) { Err(GhApiError::Unauthorized) => { self.0.is_auth_token_valid.store(false, Relaxed); + Err(GhApiError::Unauthorized) } - res => res?, + res => res.map(Download::from_response), } - - Ok(Download::from_response(response)) } } diff --git a/crates/binstalk-git-repo-api/src/gh_api_client/common.rs b/crates/binstalk-git-repo-api/src/gh_api_client/common.rs index b09823e9..86b50a6d 100644 --- a/crates/binstalk-git-repo-api/src/gh_api_client/common.rs +++ b/crates/binstalk-git-repo-api/src/gh_api_client/common.rs @@ -1,4 +1,4 @@ -use std::{fmt::Debug, future::Future, sync::OnceLock, time::Duration}; +use std::{fmt::Debug, future::Future, sync::OnceLock}; use binstalk_downloader::remote::{self, Response, Url}; use compact_str::CompactString; @@ -18,28 +18,12 @@ pub(super) fn percent_decode_http_url_path(input: &str) -> CompactString { } } -pub(super) fn check_http_status_and_header(response: &Response) -> Result<(), GhApiError> { - let headers = response.headers(); - +pub(super) fn check_http_status_and_header(response: Response) -> Result { match response.status() { - remote::StatusCode::FORBIDDEN - if headers - .get("x-ratelimit-remaining") - .map(|val| val == "0") - .unwrap_or(false) => - { - Err(GhApiError::RateLimit { - retry_after: headers.get("x-ratelimit-reset").and_then(|value| { - let secs = value.to_str().ok()?.parse().ok()?; - Some(Duration::from_secs(secs)) - }), - }) - } - remote::StatusCode::UNAUTHORIZED => Err(GhApiError::Unauthorized), remote::StatusCode::NOT_FOUND => Err(GhApiError::NotFound), - _ => Ok(()), + _ => Ok(response.error_for_status()?), } } @@ -73,9 +57,7 @@ where .send(false); async move { - let response = future.await?; - - check_http_status_and_header(&response)?; + let response = check_http_status_and_header(future.await?)?; Ok(response.json().await?) } @@ -127,8 +109,7 @@ where }); async move { - let response = res?.await?; - check_http_status_and_header(&response)?; + let response = check_http_status_and_header(res?.await?)?; let mut response: GraphQLResponse = response.json().await?;