diff --git a/CHANGELOG.md b/CHANGELOG.md index b73c52fc..11ddd6cb 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/ic-agent/src/agent/agent_config.rs b/ic-agent/src/agent/agent_config.rs index 11ed2f35..471acadb 100644 --- a/ic-agent/src/agent/agent_config.rs +++ b/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`. + fn clone_box(&self) -> Box; +} + +impl CloneableBackoff for T +where + T: Backoff + Clone + Send + Sync + 'static, +{ + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +impl Clone for Box { + 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>, + /// See [`with_retry_policy`](super::AgentBuilder::with_retry_policy). + pub retry_policy: Option>, /// 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, diff --git a/ic-agent/src/agent/builder.rs b/ic-agent/src/agent/builder.rs index f0cfee48..5de46fa6 100644 --- a/ic-agent/src/agent/builder.rs +++ b/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(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) -> Self { + self.config.retry_policy = Some(retry_policy); + self + } } diff --git a/ic-agent/src/agent/mod.rs b/ic-agent/src/agent/mod.rs index 1263b8ca..a9b20f9c 100644 --- a/ic-agent/src/agent/mod.rs +++ b/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, verify_query_signatures: bool, max_response_body_size: Option, - max_polling_time: Duration, + retry_policy: Box, #[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 { - 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 { + self.retry_policy.clone() } /// Wait for `request_status` to return a Replied response and return the arg.