feat(proof-data-handler): add tee_proof_generation_timeout_in_secs pa…

…ram (#3128)

## What ❔

Add `tee_proof_generation_timeout_in_secs` parameter to the
`proof-data-handler` configuration to avoid sharing the same
`proof_generation_timeout_in_secs` timeout with the ZK prover.

This timeout is for retrying TEE proof generation if it fails. Retries
continue indefinitely until successful.

## Why ❔

The TEE prover is much faster than the ZK prover, so some of the ZK
timeouts are too long to be shared with the TEE-specific code.

## 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 `zkstack dev fmt` and `zkstack dev
lint`.

core/lib/config/src/configs/proof_data_handler.rs

@@ -9,13 +9,18 @@ pub struct TeeConfig {
     pub tee_support: bool,
     /// All batches before this one are considered to be processed.
     pub first_tee_processed_batch: L1BatchNumber,
+    /// Timeout in seconds for retrying TEE proof generation if it fails. Retries continue
+    /// indefinitely until successful.
+    pub tee_proof_generation_timeout_in_secs: u16,
 }
 
 impl Default for TeeConfig {
     fn default() -> Self {
         TeeConfig {
             tee_support: Self::default_tee_support(),
             first_tee_processed_batch: Self::default_first_tee_processed_batch(),
+            tee_proof_generation_timeout_in_secs:
+                Self::default_tee_proof_generation_timeout_in_secs(),
         }
     }
 }
@@ -28,6 +33,14 @@ impl TeeConfig {
     pub fn default_first_tee_processed_batch() -> L1BatchNumber {
         L1BatchNumber(0)
     }
+
+    pub fn default_tee_proof_generation_timeout_in_secs() -> u16 {
+        600
+    }
+
+    pub fn tee_proof_generation_timeout(&self) -> Duration {
+        Duration::from_secs(self.tee_proof_generation_timeout_in_secs.into())
+    }
 }
 
 #[derive(Debug, Deserialize, Clone, PartialEq)]

core/lib/config/src/testonly.rs

@@ -681,6 +681,7 @@ impl Distribution<configs::ProofDataHandlerConfig> for EncodeDist {
             tee_config: configs::TeeConfig {
                 tee_support: self.sample(rng),
                 first_tee_processed_batch: L1BatchNumber(rng.gen()),
+                tee_proof_generation_timeout_in_secs: self.sample(rng),
             },
         }
     }

core/lib/env_config/src/proof_data_handler.rs

@@ -28,6 +28,7 @@ mod tests {
             tee_config: TeeConfig {
                 tee_support: true,
                 first_tee_processed_batch: L1BatchNumber(1337),
+                tee_proof_generation_timeout_in_secs: 600,
             },
         }
     }
@@ -39,6 +40,7 @@ mod tests {
             PROOF_DATA_HANDLER_HTTP_PORT="3320"
             PROOF_DATA_HANDLER_TEE_SUPPORT="true"
             PROOF_DATA_HANDLER_FIRST_TEE_PROCESSED_BATCH="1337"
+            PROOF_DATA_HANDLER_TEE_PROOF_GENERATION_TIMEOUT_IN_SECS="600"
         "#;
         let mut lock = MUTEX.lock();
         lock.set_env(config);

core/lib/protobuf_config/src/proof_data_handler.rs

@@ -23,6 +23,12 @@ impl ProtoRepr for proto::ProofDataHandler {
                     .first_tee_processed_batch
                     .map(|x| L1BatchNumber(x as u32))
                     .unwrap_or_else(configs::TeeConfig::default_first_tee_processed_batch),
+                tee_proof_generation_timeout_in_secs: self
+                    .tee_proof_generation_timeout_in_secs
+                    .map(|x| x as u16)
+                    .unwrap_or_else(
+                        configs::TeeConfig::default_tee_proof_generation_timeout_in_secs,
+                    ),
             },
         })
     }
@@ -33,6 +39,9 @@ impl ProtoRepr for proto::ProofDataHandler {
             proof_generation_timeout_in_secs: Some(this.proof_generation_timeout_in_secs.into()),
             tee_support: Some(this.tee_config.tee_support),
             first_tee_processed_batch: Some(this.tee_config.first_tee_processed_batch.0 as u64),
+            tee_proof_generation_timeout_in_secs: Some(
+                this.tee_config.tee_proof_generation_timeout_in_secs.into(),
+            ),
         }
     }
 }

core/lib/protobuf_config/src/proto/config/prover.proto

@@ -109,4 +109,5 @@ message ProofDataHandler {
   optional uint32 proof_generation_timeout_in_secs = 2; // required; s
   optional bool tee_support = 3; // optional
   optional uint64 first_tee_processed_batch = 4; // optional
+  optional uint32 tee_proof_generation_timeout_in_secs = 5; // optional
 }

core/node/proof_data_handler/src/tee_request_processor.rs

@@ -164,7 +164,7 @@ impl TeeRequestProcessor {
             .tee_proof_generation_dal()
             .lock_batch_for_proving(
                 tee_type,
-                self.config.proof_generation_timeout(),
+                self.config.tee_config.tee_proof_generation_timeout(),
                 min_batch_number,
             )
             .await

core/node/proof_data_handler/src/tests.rs

@@ -28,6 +28,7 @@ async fn request_tee_proof_inputs() {
             tee_config: TeeConfig {
                 tee_support: true,
                 first_tee_processed_batch: L1BatchNumber(0),
+                tee_proof_generation_timeout_in_secs: 600,
             },
         },
         L1BatchCommitmentMode::Rollup,
@@ -86,6 +87,7 @@ async fn submit_tee_proof() {
             tee_config: TeeConfig {
                 tee_support: true,
                 first_tee_processed_batch: L1BatchNumber(0),
+                tee_proof_generation_timeout_in_secs: 600,
             },
         },
         L1BatchCommitmentMode::Rollup,

etc/env/base/proof_data_handler.toml

@@ -1,4 +1,5 @@
 [proof_data_handler]
 http_port = 3320
 proof_generation_timeout_in_secs = 18000
+tee_proof_generation_timeout_in_secs = 600
 tee_support = true

etc/env/file_based/general.yaml

@@ -169,6 +169,7 @@ witness_vector_generator:
 data_handler:
   http_port: 3320
   proof_generation_timeout_in_secs: 18000
+  tee_proof_generation_timeout_in_secs: 600
   tee_support: true
 prover_gateway:
   api_url: http://127.0.0.1:3320