From a34c60d7020e7d7a4ad4b70f04a266e625c78eb0 Mon Sep 17 00:00:00 2001
From: Toan Pham <toan@betanet.co.il>
Date: Tue, 15 Oct 2024 21:43:12 +0700
Subject: [PATCH] add retry logic

---
 clients/cli/src/prover.rs | 47 +++++++++++++++++++++++++++++++--------
 1 file changed, 38 insertions(+), 9 deletions(-)

diff --git a/clients/cli/src/prover.rs b/clients/cli/src/prover.rs
index 8ccdfd0..6f35c13 100644
--- a/clients/cli/src/prover.rs
+++ b/clients/cli/src/prover.rs
@@ -125,11 +125,24 @@ async fn main() {
         )),
     };
 
-    if let Err(e) = client
-        .send(Message::Binary(registration.encode_to_vec()))
-        .await
-    {
-        eprintln!("Failed to send message: {:?}", e);
+    let mut retries = 0;
+    let max_retries = 5;
+
+    loop {
+        if let Err(e) = client.send(Message::Binary(registration.encode_to_vec())).await {
+            eprintln!("Failed to send message: {:?}, attempt {}/{}", e, retries + 1, max_retries);
+
+            retries += 1;
+            if retries >= max_retries {
+                eprintln!("Max retries reached, exiting...");
+                break;
+            }
+
+            // Optionally, add a delay before retrying
+            tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
+        } else {
+            break;
+        }
     }
 
     track(
@@ -263,10 +276,26 @@ async fn main() {
                 }),
             );
             progress_time = SystemTime::now();
-            client
-                .send(Message::Binary(progress.encode_to_vec()))
-                .await
-                .unwrap();
+
+            let mut retries = 0;
+            let max_retries = 5;
+            loop {
+                if let Err(e) = client.send(Message::Binary(progress.encode_to_vec())).await {
+                    eprintln!("Failed to send message: {:?}, attempt {}/{}", e, retries + 1, max_retries);
+        
+                    retries += 1;
+                    if retries >= max_retries {
+                        eprintln!("Max retries reached, exiting...");
+                        break;
+                    }
+        
+                    // Optionally, add a delay before retrying
+                    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
+                } else {
+                    break;
+                }
+            }
+
             if step == end - 1 {
                 let mut buf = Vec::new();
                 let mut writer = Box::new(&mut buf);