From ac0a8d61daf3ee0e0b3acc73e2e4d8be25c48726 Mon Sep 17 00:00:00 2001
From: Adam Spofford <adam.spofford@dfinity.org>
Date: Tue, 24 Sep 2024 08:32:51 -0700
Subject: [PATCH] Use async_trait for the trait

---
 ic-agent/Cargo.toml       |  3 +-
 ic-agent/src/agent/mod.rs | 77 +++++++++++++++++++++------------------
 2 files changed, 42 insertions(+), 38 deletions(-)

diff --git a/ic-agent/Cargo.toml b/ic-agent/Cargo.toml
index 470c25a8..7fa64e3f 100644
--- a/ic-agent/Cargo.toml
+++ b/ic-agent/Cargo.toml
@@ -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 }
@@ -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",
diff --git a/ic-agent/src/agent/mod.rs b/ic-agent/src/agent/mod.rs
index 8e00d873..41f82da6 100644
--- a/ic-agent/src/agent/mod.rs
+++ b/ic-agent/src/agent/mod.rs
@@ -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};
@@ -1842,58 +1843,62 @@ 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?)
     }
 }
 
@@ -1901,27 +1906,27 @@ 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);
             }
-        })
+        }
     }
 }