From 312defed86fbbbc1dfee489be373af1417ee624a Mon Sep 17 00:00:00 2001
From: Patrick <pab@matterlabs.dev>
Date: Mon, 8 Jul 2024 15:53:41 +0200
Subject: [PATCH] fix(tee): Introduce a 1 second delay in the batch poll
 (#2398)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## What ❔

Introduce a 1 second delay in the batch poll.

## Why ❔

We don't want to poll the batches too often if they are unavailable, as
it wastes CPU resources.

## Checklist

- [x] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [ ] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [x] Code has been formatted via `zk fmt` and `zk lint`.
---
 core/bin/zksync_tee_prover/src/tee_prover.rs | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/core/bin/zksync_tee_prover/src/tee_prover.rs b/core/bin/zksync_tee_prover/src/tee_prover.rs
index 3d0af9cc884c..b14d07b72db6 100644
--- a/core/bin/zksync_tee_prover/src/tee_prover.rs
+++ b/core/bin/zksync_tee_prover/src/tee_prover.rs
@@ -181,7 +181,7 @@ impl Task for TeeProver {
                 return Ok(());
             }
             let result = self.step().await;
-            match result {
+            let need_to_sleep = match result {
                 Ok(batch_number) => {
                     retries = 1;
                     backoff = self.config.initial_retry_backoff;
@@ -191,6 +191,9 @@ impl Task for TeeProver {
                         METRICS
                             .last_batch_number_processed
                             .set(batch_number.0 as u64);
+                        false
+                    } else {
+                        true
                     }
                 }
                 Err(err) => {
@@ -200,14 +203,17 @@ 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,
                     );
+                    true
                 }
+            };
+            if need_to_sleep {
+                tokio::time::timeout(backoff, stop_receiver.0.changed())
+                    .await
+                    .ok();
             }
         }
     }