feat(proof-data-handler): add tee_proof_generation_timeout_in_secs param

Add `tee_proof_generation_timeout_in_secs` parameter to the
`proof-data-handler` configuration to not share the same
`proof_generation_timeout_in_secs` timeout with the prover.

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

@@ -22,7 +22,7 @@ pub use self::{
     genesis::GenesisConfig,
     object_store::ObjectStoreConfig,
     observability::{ObservabilityConfig, OpentelemetryConfig},
-    proof_data_handler::ProofDataHandlerConfig,
+    proof_data_handler::{ProofDataHandlerConfig, TeeConfig},
     prover_job_monitor::ProverJobMonitorConfig,
     pruning::PruningConfig,
     secrets::{DatabaseSecrets, L1Secrets, Secrets},

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

@@ -1,16 +1,60 @@
-use std::time::Duration;
+use std::{fmt, fmt::Display, marker::PhantomData, str::FromStr, time::Duration};
 
-use serde::Deserialize;
+use serde::{de, Deserialize};
+use zksync_basic_types::L1BatchNumber;
+
+#[derive(Debug, Deserialize, Clone, PartialEq)]
+pub struct TeeConfig {
+    /// If true, the TEE support is enabled.
+    #[serde(deserialize_with = "deserialize_stringified_any")]
+    pub tee_support: bool,
+    /// All batches before this one are considered to be processed.
+    #[serde(deserialize_with = "deserialize_stringified_any", default)]
+    pub first_tee_processed_batch: L1BatchNumber,
+}
 
 #[derive(Debug, Deserialize, Clone, PartialEq)]
 pub struct ProofDataHandlerConfig {
     pub http_port: u16,
     pub proof_generation_timeout_in_secs: u16,
-    pub tee_support: bool,
+    #[serde(flatten)]
+    pub tee_config: TeeConfig,
 }
 
 impl ProofDataHandlerConfig {
     pub fn proof_generation_timeout(&self) -> Duration {
         Duration::from_secs(self.proof_generation_timeout_in_secs as u64)
     }
 }
+
+// Boilerplate to workaround https://github.com/softprops/envy/issues/26
+
+pub fn deserialize_stringified_any<'de, D, T>(deserializer: D) -> Result<T, D::Error>
+where
+    D: de::Deserializer<'de>,
+    T: FromStr,
+    T::Err: Display,
+{
+    deserializer.deserialize_any(StringifiedAnyVisitor(PhantomData))
+}
+
+pub struct StringifiedAnyVisitor<T>(PhantomData<T>);
+
+impl<'de, T> de::Visitor<'de> for StringifiedAnyVisitor<T>
+where
+    T: FromStr,
+    T::Err: Display,
+{
+    type Value = T;
+
+    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+        formatter.write_str("a string containing json data")
+    }
+
+    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
+    where
+        E: de::Error,
+    {
+        Self::Value::from_str(v).map_err(E::custom)
+    }
+}

core/lib/config/src/testonly.rs

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

core/lib/dal/src/tee_proof_generation_dal.rs

@@ -32,10 +32,10 @@ impl TeeProofGenerationDal<'_, '_> {
         &mut self,
         tee_type: TeeType,
         processing_timeout: Duration,
-        min_batch_number: Option<L1BatchNumber>,
+        min_batch_number: L1BatchNumber,
     ) -> DalResult<Option<L1BatchNumber>> {
         let processing_timeout = pg_interval_from_duration(processing_timeout);
-        let min_batch_number = min_batch_number.map_or(0, |num| i64::from(num.0));
+        let min_batch_number = i64::from(min_batch_number.0);
         sqlx::query!(
             r#"
             WITH upsert AS (

core/lib/env_config/src/proof_data_handler.rs

@@ -10,6 +10,9 @@ impl FromEnv for ProofDataHandlerConfig {
 
 #[cfg(test)]
 mod tests {
+    use zksync_basic_types::L1BatchNumber;
+    use zksync_config::configs::TeeConfig;
+
     use super::*;
     use crate::test_utils::EnvMutex;
 
@@ -19,7 +22,10 @@ mod tests {
         ProofDataHandlerConfig {
             http_port: 3320,
             proof_generation_timeout_in_secs: 18000,
-            tee_support: true,
+            tee_config: TeeConfig {
+                tee_support: true,
+                first_tee_processed_batch: L1BatchNumber(1337),
+            },
         }
     }
 
@@ -29,6 +35,7 @@ mod tests {
             PROOF_DATA_HANDLER_PROOF_GENERATION_TIMEOUT_IN_SECS="18000"
             PROOF_DATA_HANDLER_HTTP_PORT="3320"
             PROOF_DATA_HANDLER_TEE_SUPPORT="true"
+            PROOF_DATA_HANDLER_FIRST_TEE_PROCESSED_BATCH="1337"
         "#;
         let mut lock = MUTEX.lock();
         lock.set_env(config);

core/lib/protobuf_config/src/proof_data_handler.rs

@@ -14,17 +14,23 @@ impl ProtoRepr for proto::ProofDataHandler {
             proof_generation_timeout_in_secs: required(&self.proof_generation_timeout_in_secs)
                 .and_then(|x| Ok((*x).try_into()?))
                 .context("proof_generation_timeout_in_secs")?,
-            tee_support: required(&self.tee_support)
-                .copied()
-                .context("tee_support")?,
+            tee_config: configs::TeeConfig {
+                tee_support: required(&self.tee_support)
+                    .copied()
+                    .context("tee_support")?,
+                first_tee_processed_batch: required(&self.first_tee_processed_batch)
+                    .map(|x| (*x as u32).into())
+                    .context("first_tee_processed_batch")?,
+            },
         })
     }
 
     fn build(this: &Self::Type) -> Self {
         Self {
             http_port: Some(this.http_port.into()),
             proof_generation_timeout_in_secs: Some(this.proof_generation_timeout_in_secs.into()),
-            tee_support: Some(this.tee_support),
+            tee_support: Some(this.tee_config.tee_support),
+            first_tee_processed_batch: Some(this.tee_config.first_tee_processed_batch.0 as u64),
         }
     }
 }

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

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

core/node/proof_data_handler/src/lib.rs

@@ -94,7 +94,7 @@ fn create_proof_processing_router(
             ),
         );
 
-    if config.tee_support {
+    if config.tee_config.tee_support {
         let get_tee_proof_gen_processor =
             TeeRequestProcessor::new(blob_store, connection_pool, config.clone(), l2_chain_id);
         let submit_tee_proof_processor = get_tee_proof_gen_processor.clone();

core/node/proof_data_handler/src/tee_request_processor.rs

@@ -47,7 +47,7 @@ impl TeeRequestProcessor {
     ) -> Result<Option<Json<TeeProofGenerationDataResponse>>, RequestProcessorError> {
         tracing::info!("Received request for proof generation data: {:?}", request);
 
-        let mut min_batch_number: Option<L1BatchNumber> = None;
+        let mut min_batch_number = self.config.tee_config.first_tee_processed_batch;
         let mut missing_range: Option<(L1BatchNumber, L1BatchNumber)> = None;
 
         let result = loop {
@@ -72,7 +72,7 @@ impl TeeRequestProcessor {
                         None => Some((l1_batch_number, l1_batch_number)),
                     };
                     self.unlock_batch(l1_batch_number, request.tee_type).await?;
-                    min_batch_number = Some(min_batch_number.unwrap_or(l1_batch_number) + 1);
+                    min_batch_number = l1_batch_number + 1;
                 }
                 Err(err) => {
                     self.unlock_batch(l1_batch_number, request.tee_type).await?;
@@ -155,7 +155,7 @@ impl TeeRequestProcessor {
     async fn lock_batch_for_proving(
         &self,
         tee_type: TeeType,
-        min_batch_number: Option<L1BatchNumber>,
+        min_batch_number: L1BatchNumber,
     ) -> Result<Option<L1BatchNumber>, RequestProcessorError> {
         self.pool
             .connection_tagged("tee_request_processor")

core/node/proof_data_handler/src/tests.rs

@@ -7,7 +7,7 @@ use axum::{
 use serde_json::json;
 use tower::ServiceExt;
 use zksync_basic_types::L2ChainId;
-use zksync_config::configs::ProofDataHandlerConfig;
+use zksync_config::configs::{ProofDataHandlerConfig, TeeConfig};
 use zksync_dal::{ConnectionPool, CoreDal};
 use zksync_object_store::MockObjectStore;
 use zksync_prover_interface::api::SubmitTeeProofRequest;
@@ -25,7 +25,10 @@ async fn request_tee_proof_inputs() {
         ProofDataHandlerConfig {
             http_port: 1337,
             proof_generation_timeout_in_secs: 10,
-            tee_support: true,
+            tee_config: TeeConfig {
+                tee_support: true,
+                first_tee_processed_batch: L1BatchNumber(0),
+            },
         },
         L1BatchCommitmentMode::Rollup,
         L2ChainId::default(),
@@ -80,7 +83,10 @@ async fn submit_tee_proof() {
         ProofDataHandlerConfig {
             http_port: 1337,
             proof_generation_timeout_in_secs: 10,
-            tee_support: true,
+            tee_config: TeeConfig {
+                tee_support: true,
+                first_tee_processed_batch: L1BatchNumber(0),
+            },
         },
         L1BatchCommitmentMode::Rollup,
         L2ChainId::default(),