From 7b07f3044afed6cbcdc923569d8299c3398e6bd9 Mon Sep 17 00:00:00 2001 From: Patrick Beza Date: Wed, 16 Oct 2024 16:07:33 +0200 Subject: [PATCH] feat(data-proof-handler): add first processed batch option Add an option to the proof data handler to allow the first verified batch to be set. --- core/lib/config/src/configs/mod.rs | 2 +- .../config/src/configs/proof_data_handler.rs | 50 +++++++++++++++++-- core/lib/config/src/testonly.rs | 5 +- core/lib/dal/src/tee_proof_generation_dal.rs | 4 +- core/lib/env_config/src/proof_data_handler.rs | 9 +++- .../protobuf_config/src/proof_data_handler.rs | 14 ++++-- .../src/proto/config/prover.proto | 1 + core/node/proof_data_handler/src/lib.rs | 2 +- .../src/tee_request_processor.rs | 4 +- 9 files changed, 76 insertions(+), 15 deletions(-) diff --git a/core/lib/config/src/configs/mod.rs b/core/lib/config/src/configs/mod.rs index a8d136d632ea..b3a7c2913437 100644 --- a/core/lib/config/src/configs/mod.rs +++ b/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}, diff --git a/core/lib/config/src/configs/proof_data_handler.rs b/core/lib/config/src/configs/proof_data_handler.rs index de7f6969b05f..9ce89421b34b 100644 --- a/core/lib/config/src/configs/proof_data_handler.rs +++ b/core/lib/config/src/configs/proof_data_handler.rs @@ -1,12 +1,24 @@ -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 { @@ -14,3 +26,35 @@ impl ProofDataHandlerConfig { 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 +where + D: de::Deserializer<'de>, + T: FromStr, + T::Err: Display, +{ + deserializer.deserialize_any(StringifiedAnyVisitor(PhantomData)) +} + +pub struct StringifiedAnyVisitor(PhantomData); + +impl<'de, T> de::Visitor<'de> for StringifiedAnyVisitor +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(self, v: &str) -> Result + where + E: de::Error, + { + Self::Value::from_str(v).map_err(E::custom) + } +} diff --git a/core/lib/config/src/testonly.rs b/core/lib/config/src/testonly.rs index 0fdd927d19f0..1b52d36cd959 100644 --- a/core/lib/config/src/testonly.rs +++ b/core/lib/config/src/testonly.rs @@ -680,7 +680,10 @@ impl Distribution 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()), + }, } } } diff --git a/core/lib/dal/src/tee_proof_generation_dal.rs b/core/lib/dal/src/tee_proof_generation_dal.rs index d865212f190c..599409758144 100644 --- a/core/lib/dal/src/tee_proof_generation_dal.rs +++ b/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, + min_batch_number: L1BatchNumber, ) -> DalResult> { 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 ( diff --git a/core/lib/env_config/src/proof_data_handler.rs b/core/lib/env_config/src/proof_data_handler.rs index f69aa1d6dc59..fbd5d3850efe 100644 --- a/core/lib/env_config/src/proof_data_handler.rs +++ b/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); diff --git a/core/lib/protobuf_config/src/proof_data_handler.rs b/core/lib/protobuf_config/src/proof_data_handler.rs index 4b7bd2fd7c32..48dcf0c54f69 100644 --- a/core/lib/protobuf_config/src/proof_data_handler.rs +++ b/core/lib/protobuf_config/src/proof_data_handler.rs @@ -14,9 +14,14 @@ 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")?, + }, }) } @@ -24,7 +29,8 @@ impl ProtoRepr for proto::ProofDataHandler { 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), } } } diff --git a/core/lib/protobuf_config/src/proto/config/prover.proto b/core/lib/protobuf_config/src/proto/config/prover.proto index 4fe3861183bf..ae8a636f52af 100644 --- a/core/lib/protobuf_config/src/proto/config/prover.proto +++ b/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 } diff --git a/core/node/proof_data_handler/src/lib.rs b/core/node/proof_data_handler/src/lib.rs index 2c6243fdb51c..05023ce18be9 100644 --- a/core/node/proof_data_handler/src/lib.rs +++ b/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(); diff --git a/core/node/proof_data_handler/src/tee_request_processor.rs b/core/node/proof_data_handler/src/tee_request_processor.rs index 8f10f573cd84..03c8e8bd680a 100644 --- a/core/node/proof_data_handler/src/tee_request_processor.rs +++ b/core/node/proof_data_handler/src/tee_request_processor.rs @@ -47,7 +47,7 @@ impl TeeRequestProcessor { ) -> Result, RequestProcessorError> { tracing::info!("Received request for proof generation data: {:?}", request); - let mut min_batch_number: Option = 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 { @@ -151,7 +151,7 @@ impl TeeRequestProcessor { async fn lock_batch_for_proving( &self, tee_type: TeeType, - min_batch_number: Option, + min_batch_number: L1BatchNumber, ) -> Result { self.pool .connection_tagged("tee_request_processor")