Skip to content

Commit

Permalink
Use async_trait for the trait
Browse files Browse the repository at this point in the history
  • Loading branch information
adamspofford-dfinity committed Sep 24, 2024
1 parent aa557a1 commit ac0a8d6
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 38 deletions.
3 changes: 1 addition & 2 deletions ic-agent/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ include = ["src", "Cargo.toml", "../LICENSE", "README.md"]
arc-swap = { version = "1.7", optional = true }
async-channel = { version = "1.9", optional = true }
async-lock = "3.3"
async-trait = { version = "0.1", optional = true }
async-trait = "0.1"
async-watch = { version = "0.3", optional = true }
backoff = "0.4.0"
cached = { version = "0.52", features = ["ahash"], default-features = false }
Expand Down Expand Up @@ -109,7 +109,6 @@ wasm-bindgen = [
_internal_dynamic-routing = [
"dep:arc-swap",
"dep:async-channel",
"dep:async-trait",
"dep:async-watch",
"dep:stop-token",
"tracing",
Expand Down
77 changes: 41 additions & 36 deletions ic-agent/src/agent/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ pub use agent_config::AgentConfig;
pub use agent_error::AgentError;
use agent_error::HttpErrorPayload;
use async_lock::Semaphore;
use async_trait::async_trait;
pub use builder::AgentBuilder;
use cached::{Cached, TimedCache};
use ed25519_consensus::{Error as Ed25519Error, Signature, VerificationKey};
Expand Down Expand Up @@ -1842,86 +1843,90 @@ impl<'agent> IntoFuture for UpdateBuilder<'agent> {
}

/// HTTP client middleware. Implemented automatically for `reqwest`-compatible by-ref `tower::Service`, such as `reqwest_middleware`.
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
#[cfg_attr(not(target_family = "wasm"), async_trait)]
pub trait HttpService: Send + Sync {
/// Perform a HTTP request. Any retry logic should call `req` again, instead of `Request::try_clone`.
fn call<'a>(
async fn call<'a>(
&'a self,
req: &'a (dyn Fn() -> Result<Request, AgentError> + Send + Sync),
max_retries: usize,
) -> AgentFuture<'a, Response>;
) -> Result<Response, AgentError>;
}
#[cfg(not(target_family = "wasm"))]
#[async_trait]
impl<T> HttpService for T
where
for<'a> &'a T: Service<Request, Response = Response, Error = reqwest::Error>,
for<'a> <&'a Self as Service<Request>>::Future: Send,
T: Send + Sync + ?Sized,
{
fn call<'a>(
#[allow(clippy::needless_arbitrary_self_type)]
async fn call<'a>(
mut self: &'a Self,
req: &'a (dyn Fn() -> Result<Request, AgentError> + Send + Sync),
max_retries: usize,
) -> AgentFuture<'a, Response> {
Box::pin(async move {
let mut retry_count = 0;
loop {
match Service::call(&mut self, req()?).await {
Err(err) => {
// Network-related errors can be retried.
if err.is_connect() {
if retry_count >= max_retries {
return Err(AgentError::TransportError(err));
}
retry_count += 1;
) -> Result<Response, AgentError> {
let mut retry_count = 0;
loop {
match Service::call(&mut self, req()?).await {
Err(err) => {
// Network-related errors can be retried.
if err.is_connect() {
if retry_count >= max_retries {
return Err(AgentError::TransportError(err));
}
retry_count += 1;
}
Ok(resp) => return Ok(resp),
}
Ok(resp) => return Ok(resp),
}
})
}
}
}

#[cfg(target_family = "wasm")]
#[async_trait(?Send)]
impl<T> HttpService for T
where
for<'a> &'a T: Service<Request, Response = Response, Error = reqwest::Error>,
T: Send + Sync + ?Sized,
{
fn call<'a>(
#[allow(clippy::needless_arbitrary_self_type)]
async fn call<'a>(
mut self: &'a Self,
req: &'a (dyn Fn() -> Result<Request, AgentError> + Send + Sync),
_: usize,
) -> AgentFuture<'a, Response> {
Box::pin(async move { Ok(Service::call(&mut self, req()?).await?) })
) -> Result<Response, AgentError> {
Ok(Service::call(&mut self, req()?).await?)
}
}

struct Retry429Logic {
client: Client,
}

#[cfg_attr(target_family = "wasm", async_trait(?Send))]
#[cfg_attr(not(target_family = "wasm"), async_trait)]
impl HttpService for Retry429Logic {
fn call<'a>(
async fn call<'a>(
&'a self,
req: &'a (dyn Fn() -> Result<Request, AgentError> + Send + Sync),
_max_retries: usize,
) -> AgentFuture<'a, Response> {
Box::pin(async move {
loop {
#[cfg(not(target_family = "wasm"))]
let resp = self.client.call(req, _max_retries).await?;
// Client inconveniently does not implement Service on wasm
#[cfg(target_family = "wasm")]
let resp = self.client.execute(req()?).await?;
if resp.status() == StatusCode::TOO_MANY_REQUESTS {
crate::util::sleep(Duration::from_millis(250)).await;
continue;
} else {
break Ok(resp);
}
) -> Result<Response, AgentError> {
loop {
#[cfg(not(target_family = "wasm"))]
let resp = self.client.call(req, _max_retries).await?;
// Client inconveniently does not implement Service on wasm
#[cfg(target_family = "wasm")]
let resp = self.client.execute(req()?).await?;
if resp.status() == StatusCode::TOO_MANY_REQUESTS {
crate::util::sleep(Duration::from_millis(250)).await;
continue;
} else {
break Ok(resp);
}
})
}
}
}

Expand Down

0 comments on commit ac0a8d6

Please sign in to comment.