feat: add with_max_polling_time for ic-agent (#604)

Author: chenyan-dfinity

CHANGELOG.md

@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+* Added `AgentBuilder::with_max_polling_time` to config the maximum time to wait for a response from the replica.
+
 ## [0.38.2] - 2024-09-30
 
 * Limited the number of HTTP 429 retries. Users receiving this error should configure `with_max_concurrent_requests`.

ic-agent/src/agent/agent_config.rs

@@ -31,6 +31,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_max_polling_time`](super::AgentBuilder::with_max_polling_time).
+    pub max_polling_time: Duration,
 }
 
 impl Default for AgentConfig {
@@ -46,6 +48,7 @@ impl Default for AgentConfig {
             route_provider: None,
             max_response_body_size: None,
             max_tcp_error_retries: 0,
+            max_polling_time: Duration::from_secs(60 * 5),
         }
     }
 }

ic-agent/src/agent/builder.rs

@@ -164,4 +164,9 @@ impl AgentBuilder {
         self.config.max_response_body_size = Some(max_size);
         self
     }
+    /// Set the maximum time to wait for a response from the replica.
+    pub fn with_max_polling_time(mut self, max_polling_time: std::time::Duration) -> Self {
+        self.config.max_polling_time = max_polling_time;
+        self
+    }
 }

ic-agent/src/agent/mod.rs

@@ -156,6 +156,7 @@ pub struct Agent {
     concurrent_requests_semaphore: Arc<Semaphore>,
     verify_query_signatures: bool,
     max_response_body_size: Option<usize>,
+    max_polling_time: Duration,
     #[allow(dead_code)]
     max_tcp_error_retries: usize,
 }
@@ -208,6 +209,7 @@ 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,
         })
     }
 
@@ -615,12 +617,12 @@ impl Agent {
         })
     }
 
-    fn get_retry_policy() -> ExponentialBackoff<SystemClock> {
+    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(Duration::from_secs(60 * 5)))
+            .with_max_elapsed_time(Some(self.max_polling_time))
             .build()
     }
 
@@ -631,7 +633,7 @@ impl Agent {
         effective_canister_id: Principal,
         signed_request_status: Vec<u8>,
     ) -> Result<Vec<u8>, AgentError> {
-        let mut retry_policy = Self::get_retry_policy();
+        let mut retry_policy = self.get_retry_policy();
 
         let mut request_accepted = false;
         loop {
@@ -679,7 +681,7 @@ impl Agent {
         request_id: &RequestId,
         effective_canister_id: Principal,
     ) -> Result<Vec<u8>, AgentError> {
-        let mut retry_policy = Self::get_retry_policy();
+        let mut retry_policy = self.get_retry_policy();
 
         let mut request_accepted = false;
         loop {