feat: Allow to choose a custom retry policy

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+* Added the possibility to set a custom `retry_policy`.
 * Add `read_state_canister_controllers` and `read_state_canister_module_hash` functions.
 
 ## [0.40.0] - 2025-03-17

ic-agent/src/agent/agent_config.rs

@@ -1,3 +1,4 @@
+use backoff::backoff::Backoff;
 use reqwest::Client;
 use url::Url;
 
@@ -9,6 +10,27 @@ use std::{sync::Arc, time::Duration};
 
 use super::{route_provider::RouteProvider, HttpService};
 
+/// A helper trait for cloning backoff policies.
+pub trait CloneableBackoff: Backoff + Send + Sync {
+    /// Clone the backoff policy into a `Box<dyn CloneableBackoff>`.
+    fn clone_box(&self) -> Box<dyn CloneableBackoff>;
+}
+
+impl<T> CloneableBackoff for T
+where
+    T: Backoff + Clone + Send + Sync + 'static,
+{
+    fn clone_box(&self) -> Box<dyn CloneableBackoff> {
+        Box::new(self.clone())
+    }
+}
+
+impl Clone for Box<dyn CloneableBackoff> {
+    fn clone(&self) -> Self {
+        self.clone_box()
+    }
+}
+
 /// A configuration for an agent.
 #[non_exhaustive]
 pub struct AgentConfig {
@@ -32,6 +54,8 @@ pub struct AgentConfig {
     pub max_tcp_error_retries: usize,
     /// See [`with_arc_http_middleware`](super::AgentBuilder::with_arc_http_middleware).
     pub http_service: Option<Arc<dyn HttpService>>,
+    /// See [`with_retry_policy`](super::AgentBuilder::with_retry_policy).
+    pub retry_policy: Option<Box<dyn CloneableBackoff>>,
     /// See [`with_max_polling_time`](super::AgentBuilder::with_max_polling_time).
     pub max_polling_time: Duration,
     /// See [`with_background_dynamic_routing`](super::AgentBuilder::with_background_dynamic_routing).
@@ -53,6 +77,7 @@ impl Default for AgentConfig {
             route_provider: None,
             max_response_body_size: None,
             max_tcp_error_retries: 0,
+            retry_policy: None,
             max_polling_time: Duration::from_secs(60 * 5),
             background_dynamic_routing: false,
             url: None,

ic-agent/src/agent/builder.rs

@@ -4,7 +4,7 @@ use crate::{
 };
 use std::sync::Arc;
 
-use super::{route_provider::RouteProvider, HttpService};
+use super::{route_provider::RouteProvider, CloneableBackoff, HttpService};
 
 /// A builder for an [`Agent`].
 #[derive(Default)]
@@ -168,4 +168,18 @@ impl AgentBuilder {
         self.config.max_polling_time = max_polling_time;
         self
     }
+
+    /// Set the retry policy for the agent to use when retrying requests.
+    pub fn with_retry_policy<B>(self, retry_policy: B) -> Self
+    where
+        B: 'static + CloneableBackoff,
+    {
+        self.with_box_retry_policy(Box::new(retry_policy))
+    }
+
+    /// Same as [`Self::with_retry_policy`], but reuses an existing `Box`
+    pub fn with_box_retry_policy(mut self, retry_policy: Box<dyn CloneableBackoff>) -> Self {
+        self.config.retry_policy = Some(retry_policy);
+        self
+    }
 }

ic-agent/src/agent/mod.rs

@@ -11,7 +11,7 @@ pub(crate) mod response_authentication;
 pub mod route_provider;
 pub mod status;
 
-pub use agent_config::AgentConfig;
+pub use agent_config::{AgentConfig, CloneableBackoff};
 pub use agent_error::AgentError;
 use agent_error::{HttpErrorPayload, Operation};
 use async_lock::Semaphore;
@@ -52,7 +52,6 @@ use crate::{
     to_request_id, RequestId,
 };
 use backoff::{backoff::Backoff, ExponentialBackoffBuilder};
-use backoff::{exponential::ExponentialBackoff, SystemClock};
 use ic_certification::{Certificate, Delegation, Label};
 use ic_transport_types::{
     signed::{SignedQuery, SignedRequestStatus, SignedUpdate},
@@ -161,7 +160,7 @@ pub struct Agent {
     concurrent_requests_semaphore: Arc<Semaphore>,
     verify_query_signatures: bool,
     max_response_body_size: Option<usize>,
-    max_polling_time: Duration,
+    retry_policy: Box<dyn CloneableBackoff>,
     #[allow(dead_code)]
     max_tcp_error_retries: usize,
 }
@@ -236,7 +235,17 @@ impl Agent {
             concurrent_requests_semaphore: Arc::new(Semaphore::new(config.max_concurrent_requests)),
             max_response_body_size: config.max_response_body_size,
             max_tcp_error_retries: config.max_tcp_error_retries,
-            max_polling_time: config.max_polling_time,
+            retry_policy: match config.retry_policy {
+                Some(retry_policy) => retry_policy,
+                None => Box::new(
+                    ExponentialBackoffBuilder::new()
+                        .with_initial_interval(Duration::from_millis(500))
+                        .with_max_interval(Duration::from_secs(1))
+                        .with_multiplier(1.4)
+                        .with_max_elapsed_time(Some(config.max_polling_time))
+                        .build(),
+                ),
+            },
         })
     }
 
@@ -713,13 +722,8 @@ impl Agent {
         })
     }
 
-    fn get_retry_policy(&self) -> ExponentialBackoff<SystemClock> {
-        ExponentialBackoffBuilder::new()
-            .with_initial_interval(Duration::from_millis(500))
-            .with_max_interval(Duration::from_secs(1))
-            .with_multiplier(1.4)
-            .with_max_elapsed_time(Some(self.max_polling_time))
-            .build()
+    fn get_retry_policy(&self) -> Box<dyn CloneableBackoff> {
+        self.retry_policy.clone()
     }
 
     /// Wait for `request_status` to return a Replied response and return the arg.