From ba26e21b73ec99c103582e992ca627f96c207f68 Mon Sep 17 00:00:00 2001
From: HAOYUatHZ <37070449+HAOYUatHZ@users.noreply.github.com>
Date: Thu, 26 Sep 2024 13:39:06 +1000
Subject: [PATCH] feat(coordinator client): add timeout and retry (#14)

---
 src/config.rs                                 |  3 +++
 src/coordinator_handler/api.rs                | 20 ++++++++++++++-----
 src/coordinator_handler/coordinator_client.rs |  2 +-
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index eac552d..e25320e 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -15,6 +15,9 @@ pub struct Config {
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct CoordinatorConfig {
     pub base_url: String,
+    pub retry_count: u32,
+    pub retry_wait_time_sec: u64,
+    pub connection_timeout_sec: u64,
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
diff --git a/src/coordinator_handler/api.rs b/src/coordinator_handler/api.rs
index 5abea89..a93b6a3 100644
--- a/src/coordinator_handler/api.rs
+++ b/src/coordinator_handler/api.rs
@@ -2,23 +2,33 @@ use super::{
     ChallengeResponseData, GetTaskRequest, GetTaskResponseData, LoginRequest, LoginResponseData,
     Response, SubmitProofRequest, SubmitProofResponseData,
 };
+use crate::config::CoordinatorConfig;
+use core::time::Duration;
 use reqwest::{header::CONTENT_TYPE, Url};
 use reqwest_middleware::{ClientBuilder, ClientWithMiddleware};
+use reqwest_retry::{policies::ExponentialBackoff, RetryTransientMiddleware};
 use serde::Serialize;
 
 pub struct Api {
     pub base_url: Url,
+    send_timeout: Duration,
     pub client: ClientWithMiddleware,
 }
 
 impl Api {
-    pub fn new(base_url: &str) -> anyhow::Result<Self> {
+    pub fn new(cfg: CoordinatorConfig) -> anyhow::Result<Self> {
+        let retry_wait_duration = Duration::from_secs(cfg.retry_wait_time_sec);
+        let retry_policy = ExponentialBackoff::builder()
+            .retry_bounds(retry_wait_duration / 2, retry_wait_duration)
+            .build_with_max_retries(cfg.retry_count);
+
         let client = ClientBuilder::new(reqwest::Client::new())
-            // .with(RetryTransientMiddleware::new_with_policy(retry_policy)) // TODO: retry policy
+            .with(RetryTransientMiddleware::new_with_policy(retry_policy))
             .build();
 
         Ok(Self {
-            base_url: Url::parse(base_url)?,
+            base_url: Url::parse(&cfg.base_url)?,
+            send_timeout: core::time::Duration::from_secs(cfg.connection_timeout_sec),
             client,
         })
     }
@@ -47,7 +57,7 @@ impl Api {
             .header(CONTENT_TYPE, "application/json")
             .bearer_auth(token)
             .body(request_body)
-            // .timeout(self.send_timeout) // TODO: send_timeout
+            .timeout(self.send_timeout)
             .send()
             .await?;
 
@@ -76,7 +86,7 @@ impl Api {
             .client
             .get(url)
             .header(CONTENT_TYPE, "application/json")
-            // .timeout(self.send_timeout) // TODO: send_timeout
+            .timeout(self.send_timeout)
             .send()
             .await?;
 
diff --git a/src/coordinator_handler/coordinator_client.rs b/src/coordinator_handler/coordinator_client.rs
index 0a14f26..a687430 100644
--- a/src/coordinator_handler/coordinator_client.rs
+++ b/src/coordinator_handler/coordinator_client.rs
@@ -29,7 +29,7 @@ impl CoordinatorClient {
         let rt = tokio::runtime::Builder::new_current_thread()
             .enable_all()
             .build()?;
-        let api = Api::new(&cfg.base_url)?; // TODO: retry policy
+        let api = Api::new(cfg)?;
         let client = Self {
             circuit_type,
             vks,