Skip to content

Commit a014955

Browse files
committed
chore: cherry pick req extentions (#738)
<!-- Please make sure there is an issue that this PR is correlated to. --> ## Changes <!-- If there are frontend changes, please include screenshots. -->
1 parent e91d538 commit a014955

File tree

6 files changed

+10624
-1
lines changed

6 files changed

+10624
-1
lines changed

lib/global-error/Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,11 @@ protobuf-src = ["types/protobuf-src"]
1010
chirp = ["types"]
1111

1212
[dependencies]
13+
async-trait = "0.1"
1314
formatted-error = { path = "../formatted-error" }
1415
types = { path = "../types/core", optional = true }
1516
http = "0.2"
17+
reqwest = "0.11"
1618
serde = { version = "1.0", features = ["derive"] }
1719
serde_json = "1.0"
1820
thiserror = "1.0"

lib/global-error/src/ext.rs

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
use crate::Location;
1+
use crate::{bail, GlobalResult, Location};
2+
use async_trait::async_trait;
23

34
#[derive(Debug, thiserror::Error)]
45
pub enum AssertionError {
@@ -122,3 +123,23 @@ impl<'a, T> UnwrapOrAssertError for &'a &'a Option<T> {
122123
}
123124
}
124125
}
126+
127+
#[async_trait]
128+
pub trait ToGlobalError: Sized {
129+
async fn to_global_error(self) -> GlobalResult<Self>;
130+
}
131+
132+
#[async_trait]
133+
impl ToGlobalError for reqwest::Response {
134+
async fn to_global_error(self) -> GlobalResult<Self> {
135+
if self.status().is_success() {
136+
Ok(self)
137+
} else {
138+
let url = self.url().clone();
139+
let status = self.status();
140+
let body = self.text().await?;
141+
142+
bail!(format!("{url} ({status}):\n{body}"));
143+
}
144+
}
145+
}

lib/util/core/Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ macros = []
1111
serde = []
1212

1313
[dependencies]
14+
async-trait = "0.1"
1415
bcrypt = "0.13.0"
1516
chrono = "0.4"
1617
formatted-error = { path = "../../formatted-error", optional = true }
@@ -20,6 +21,7 @@ ipnet = { version = "2.7", features = ["serde"] }
2021
lazy_static = "1.4"
2122
rand = "0.8"
2223
regex = "1.4"
24+
reqwest = "0.11"
2325
rivet-util-env = { path = "../env" }
2426
rivet-util-macros = { path = "../macros" }
2527
serde = { version = "1.0", features = ["derive"] }

lib/util/core/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ pub mod geo;
1616
pub mod glob;
1717
pub mod math;
1818
pub mod net;
19+
pub mod req;
1920
pub mod route;
2021
pub mod sort;
2122
pub mod timestamp;

lib/util/core/src/req.rs

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
use std::time::Duration;
2+
3+
use async_trait::async_trait;
4+
use global_error::prelude::*;
5+
6+
#[async_trait]
7+
pub trait SendRetry {
8+
/// Retries the request upon receiving a 429 response.
9+
async fn send_retry(self, mut retries: usize) -> GlobalResult<reqwest::Response>;
10+
}
11+
12+
#[async_trait]
13+
impl SendRetry for reqwest::RequestBuilder {
14+
async fn send_retry(self, mut retries: usize) -> GlobalResult<reqwest::Response> {
15+
loop {
16+
let req = unwrap!(self.try_clone());
17+
let res = req.send().await?;
18+
19+
if let reqwest::StatusCode::TOO_MANY_REQUESTS = res.status() {
20+
if retries != 0 {
21+
retries -= 1;
22+
23+
// TODO: Parse all valid Retry-After formats. Currently only parses duration
24+
let retry_time = res
25+
.headers()
26+
.get("Retry-After")
27+
.map(|x| x.to_str())
28+
.transpose()?
29+
.map(|x| x.parse::<u64>())
30+
.transpose()?
31+
.unwrap_or(5);
32+
tokio::time::sleep(Duration::from_secs(retry_time)).await;
33+
34+
continue;
35+
}
36+
}
37+
38+
break Ok(res);
39+
}
40+
}
41+
}

0 commit comments

Comments
 (0)