From 319ab208d4ea4c1180a04d52a9f14231536efcd7 Mon Sep 17 00:00:00 2001 From: Pierugo Pace Date: Thu, 15 May 2025 13:19:54 +0000 Subject: [PATCH 1/3] feat: allow a custom retry policy --- ic-agent/src/agent/agent_config.rs | 25 +++++++++++++++++++++++++ ic-agent/src/agent/builder.rs | 16 +++++++++++++++- ic-agent/src/agent/mod.rs | 26 +++++++++++++++----------- 3 files changed, 55 insertions(+), 12 deletions(-) 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..9abd194b 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_arc_retry_policy(Box::new(retry_policy)) + } + + /// Same as [`Self::with_retry_policy`], but reuses an existing `Box` + pub fn with_arc_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. From c7bd08d6590d056bef7b2a07f5633f07811c41eb Mon Sep 17 00:00:00 2001 From: Pierugo Pace Date: Thu, 15 May 2025 13:36:15 +0000 Subject: [PATCH 2/3] docs: changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) 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 From cd467f0d7c838e5554dba6804e5cdd48a6249345 Mon Sep 17 00:00:00 2001 From: Pierugo Pace Date: Thu, 15 May 2025 13:43:20 +0000 Subject: [PATCH 3/3] typo --- ic-agent/src/agent/builder.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ic-agent/src/agent/builder.rs b/ic-agent/src/agent/builder.rs index 9abd194b..5de46fa6 100644 --- a/ic-agent/src/agent/builder.rs +++ b/ic-agent/src/agent/builder.rs @@ -174,11 +174,11 @@ impl AgentBuilder { where B: 'static + CloneableBackoff, { - self.with_arc_retry_policy(Box::new(retry_policy)) + self.with_box_retry_policy(Box::new(retry_policy)) } /// Same as [`Self::with_retry_policy`], but reuses an existing `Box` - pub fn with_arc_retry_policy(mut self, retry_policy: Box) -> Self { + pub fn with_box_retry_policy(mut self, retry_policy: Box) -> Self { self.config.retry_policy = Some(retry_policy); self }