From 3ab286b63bd6bb3e933aba496998308d8ab97c35 Mon Sep 17 00:00:00 2001
From: Patrick Beza <pab@matterlabs.dev>
Date: Mon, 8 Jul 2024 11:55:51 +0200
Subject: [PATCH 1/4] fix(tee): Introduce a 1 second delay in the batch poll

Rationale: we don't want to poll the batch too often, as it's a waste
of CPU resources.
---
 core/bin/zksync_tee_prover/src/tee_prover.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/core/bin/zksync_tee_prover/src/tee_prover.rs b/core/bin/zksync_tee_prover/src/tee_prover.rs
index 9d692e84f10e..6a7d7d1a3017 100644
--- a/core/bin/zksync_tee_prover/src/tee_prover.rs
+++ b/core/bin/zksync_tee_prover/src/tee_prover.rs
@@ -188,15 +188,15 @@ impl Task for TeeProver {
                     }
                     retries += 1;
                     tracing::warn!(%err, "Failed TEE prover step function {retries}/{}, retrying in {} milliseconds.", self.config.max_retries, backoff.as_millis());
-                    tokio::time::timeout(backoff, stop_receiver.0.changed())
-                        .await
-                        .ok();
                     backoff = std::cmp::min(
                         backoff.mul_f32(self.config.retry_backoff_multiplier),
                         self.config.max_backoff,
                     );
                 }
             }
+            tokio::time::timeout(backoff, stop_receiver.0.changed())
+                .await
+                .ok();
         }
     }
 }

From 8e4a7e48389d70d5245803c13fe8d4decdde94e9 Mon Sep 17 00:00:00 2001
From: Patrick Beza <pab@matterlabs.dev>
Date: Mon, 8 Jul 2024 12:52:43 +0200
Subject: [PATCH 2/4] Don't sleep as long as there are pending jobs

---
 core/bin/zksync_tee_prover/src/tee_prover.rs | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/core/bin/zksync_tee_prover/src/tee_prover.rs b/core/bin/zksync_tee_prover/src/tee_prover.rs
index 6a7d7d1a3017..f50acf16ef7d 100644
--- a/core/bin/zksync_tee_prover/src/tee_prover.rs
+++ b/core/bin/zksync_tee_prover/src/tee_prover.rs
@@ -1,3 +1,4 @@
+use core::net;
 use std::time::Duration;
 
 use secp256k1::{ecdsa::Signature, Message, PublicKey, Secp256k1, SecretKey};
@@ -169,6 +170,8 @@ impl Task for TeeProver {
                 return Ok(());
             }
             let result = self.step().await;
+            let mut need_to_sleep = true;
+
             match result {
                 Ok(batch_number) => {
                     retries = 1;
@@ -179,6 +182,7 @@ impl Task for TeeProver {
                         METRICS
                             .last_batch_number_processed
                             .set(batch_number.0 as u64);
+                        need_to_sleep = false;
                     }
                 }
                 Err(err) => {
@@ -194,9 +198,12 @@ impl Task for TeeProver {
                     );
                 }
             }
-            tokio::time::timeout(backoff, stop_receiver.0.changed())
-                .await
-                .ok();
+
+            if need_to_sleep {
+                tokio::time::timeout(backoff, stop_receiver.0.changed())
+                    .await
+                    .ok();
+            }
         }
     }
 }

From 5b208a02992cafadf72a249431b0b50be7322d80 Mon Sep 17 00:00:00 2001
From: Patrick Beza <pab@matterlabs.dev>
Date: Mon, 8 Jul 2024 13:02:20 +0200
Subject: [PATCH 3/4] fixup! Don't sleep as long as there are pending jobs

---
 core/bin/zksync_tee_prover/src/tee_prover.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/core/bin/zksync_tee_prover/src/tee_prover.rs b/core/bin/zksync_tee_prover/src/tee_prover.rs
index f50acf16ef7d..80c7830137ff 100644
--- a/core/bin/zksync_tee_prover/src/tee_prover.rs
+++ b/core/bin/zksync_tee_prover/src/tee_prover.rs
@@ -1,4 +1,3 @@
-use core::net;
 use std::time::Duration;
 
 use secp256k1::{ecdsa::Signature, Message, PublicKey, Secp256k1, SecretKey};

From 5bcd55b2445930eada3f4c85004d52d2f3c619e0 Mon Sep 17 00:00:00 2001
From: Patrick Beza <pab@matterlabs.dev>
Date: Mon, 8 Jul 2024 13:14:05 +0200
Subject: [PATCH 4/4] Be more idiomatic

---
 core/bin/zksync_tee_prover/src/tee_prover.rs | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/core/bin/zksync_tee_prover/src/tee_prover.rs b/core/bin/zksync_tee_prover/src/tee_prover.rs
index 80c7830137ff..3a65f257c68d 100644
--- a/core/bin/zksync_tee_prover/src/tee_prover.rs
+++ b/core/bin/zksync_tee_prover/src/tee_prover.rs
@@ -169,9 +169,7 @@ impl Task for TeeProver {
                 return Ok(());
             }
             let result = self.step().await;
-            let mut need_to_sleep = true;
-
-            match result {
+            let need_to_sleep = match result {
                 Ok(batch_number) => {
                     retries = 1;
                     backoff = self.config.initial_retry_backoff;
@@ -181,7 +179,9 @@ impl Task for TeeProver {
                         METRICS
                             .last_batch_number_processed
                             .set(batch_number.0 as u64);
-                        need_to_sleep = false;
+                        false
+                    } else {
+                        true
                     }
                 }
                 Err(err) => {
@@ -195,9 +195,9 @@ impl Task for TeeProver {
                         backoff.mul_f32(self.config.retry_backoff_multiplier),
                         self.config.max_backoff,
                     );
+                    true
                 }
-            }
-
+            };
             if need_to_sleep {
                 tokio::time::timeout(backoff, stop_receiver.0.changed())
                     .await