From faf2629d1fadb15d5c4cc771890f9980ab28ee75 Mon Sep 17 00:00:00 2001 From: Patrick Beza Date: Mon, 21 Oct 2024 12:34:54 +0200 Subject: [PATCH 1/6] feat(data-handler): add first_tee_processed_batch option in staging --- 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 | 6 +-- core/node/proof_data_handler/src/tests.rs | 12 +++-- 10 files changed, 86 insertions(+), 19 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 9b1ec13e2d2e..5aca45d6ee6f 100644 --- a/core/lib/config/src/testonly.rs +++ b/core/lib/config/src/testonly.rs @@ -669,7 +669,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 bde07f732802..755d02769101 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 661c76d20006..e014fca15d77 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 2c2a56300097..3a4d2a0bba3c 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 { @@ -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, + min_batch_number: L1BatchNumber, ) -> Result, RequestProcessorError> { self.pool .connection_tagged("tee_request_processor") diff --git a/core/node/proof_data_handler/src/tests.rs b/core/node/proof_data_handler/src/tests.rs index a10044cacd9c..63ea087a81c4 100644 --- a/core/node/proof_data_handler/src/tests.rs +++ b/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(), From 9e96fd1b3ea4532e9c1ea32acd4cc42c4d45cfec Mon Sep 17 00:00:00 2001 From: Patrick Beza Date: Tue, 22 Oct 2024 13:20:11 +0200 Subject: [PATCH 2/6] Address code review comments --- .../config/src/configs/proof_data_handler.rs | 51 +++++-------------- core/lib/env_config/src/proof_data_handler.rs | 9 +++- .../src/proto/config/prover.proto | 2 +- 3 files changed, 23 insertions(+), 39 deletions(-) diff --git a/core/lib/config/src/configs/proof_data_handler.rs b/core/lib/config/src/configs/proof_data_handler.rs index 9ce89421b34b..bbe9e881d5d6 100644 --- a/core/lib/config/src/configs/proof_data_handler.rs +++ b/core/lib/config/src/configs/proof_data_handler.rs @@ -1,23 +1,32 @@ -use std::{fmt, fmt::Display, marker::PhantomData, str::FromStr, time::Duration}; +use std::time::Duration; -use serde::{de, Deserialize}; +use serde::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, } +impl Default for TeeConfig { + fn default() -> Self { + TeeConfig { + tee_support: false, + first_tee_processed_batch: L1BatchNumber(0), + } + } +} + #[derive(Debug, Deserialize, Clone, PartialEq)] pub struct ProofDataHandlerConfig { pub http_port: u16, pub proof_generation_timeout_in_secs: u16, - #[serde(flatten)] + #[serde(skip)] + // ^ Filled in separately in `Self::from_env()`. We cannot use `serde(flatten)` because it + // doesn't work with `envy`: https://github.com/softprops/envy/issues/26 pub tee_config: TeeConfig, } @@ -26,35 +35,3 @@ 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/env_config/src/proof_data_handler.rs b/core/lib/env_config/src/proof_data_handler.rs index fbd5d3850efe..716e7a890cd4 100644 --- a/core/lib/env_config/src/proof_data_handler.rs +++ b/core/lib/env_config/src/proof_data_handler.rs @@ -4,7 +4,14 @@ use crate::{envy_load, FromEnv}; impl FromEnv for ProofDataHandlerConfig { fn from_env() -> anyhow::Result { - envy_load("proof_data_handler", "PROOF_DATA_HANDLER_") + Ok(Self { + http_port: std::env::var("PROOF_DATA_HANDLER_HTTP_PORT")?.parse()?, + proof_generation_timeout_in_secs: std::env::var( + "PROOF_DATA_HANDLER_PROOF_GENERATION_TIMEOUT_IN_SECS", + )? + .parse()?, + tee_config: envy_load("proof_data_handler", "PROOF_DATA_HANDLER_")?, + }) } } diff --git a/core/lib/protobuf_config/src/proto/config/prover.proto b/core/lib/protobuf_config/src/proto/config/prover.proto index ae8a636f52af..5ce173d91dd9 100644 --- a/core/lib/protobuf_config/src/proto/config/prover.proto +++ b/core/lib/protobuf_config/src/proto/config/prover.proto @@ -108,5 +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 + optional uint64 first_tee_processed_batch = 4; // optional } From a75a0a2c7b8b86fd972e5364619c61f4e482fe1b Mon Sep 17 00:00:00 2001 From: Patrick Beza Date: Wed, 23 Oct 2024 10:21:21 +0200 Subject: [PATCH 3/6] Simplify tee_config loading in ProofDataHandlerConfig --- core/lib/env_config/src/proof_data_handler.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/core/lib/env_config/src/proof_data_handler.rs b/core/lib/env_config/src/proof_data_handler.rs index 716e7a890cd4..b5bfda4544e7 100644 --- a/core/lib/env_config/src/proof_data_handler.rs +++ b/core/lib/env_config/src/proof_data_handler.rs @@ -5,12 +5,8 @@ use crate::{envy_load, FromEnv}; impl FromEnv for ProofDataHandlerConfig { fn from_env() -> anyhow::Result { Ok(Self { - http_port: std::env::var("PROOF_DATA_HANDLER_HTTP_PORT")?.parse()?, - proof_generation_timeout_in_secs: std::env::var( - "PROOF_DATA_HANDLER_PROOF_GENERATION_TIMEOUT_IN_SECS", - )? - .parse()?, - tee_config: envy_load("proof_data_handler", "PROOF_DATA_HANDLER_")?, + tee_config: envy_load("proof_data_handler.tee", "PROOF_DATA_HANDLER_")?, + ..envy_load("proof_data_handler", "PROOF_DATA_HANDLER_")? }) } } From 4c66736c983811d8dfe95de65754c31ff6896ccf Mon Sep 17 00:00:00 2001 From: Patrick Beza Date: Wed, 23 Oct 2024 11:09:39 +0200 Subject: [PATCH 4/6] Set defaults in conversion from the Protobuf type to the domain type --- core/lib/config/src/configs/proof_data_handler.rs | 14 ++++++++++++-- core/lib/protobuf_config/src/proof_data_handler.rs | 8 ++------ 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/core/lib/config/src/configs/proof_data_handler.rs b/core/lib/config/src/configs/proof_data_handler.rs index bbe9e881d5d6..1094b1bb1801 100644 --- a/core/lib/config/src/configs/proof_data_handler.rs +++ b/core/lib/config/src/configs/proof_data_handler.rs @@ -14,12 +14,22 @@ pub struct TeeConfig { impl Default for TeeConfig { fn default() -> Self { TeeConfig { - tee_support: false, - first_tee_processed_batch: L1BatchNumber(0), + tee_support: Self::default_tee_support(), + first_tee_processed_batch: Self::default_first_tee_processed_batch(), } } } +impl TeeConfig { + pub fn default_tee_support() -> bool { + false + } + + pub fn default_first_tee_processed_batch() -> L1BatchNumber { + L1BatchNumber(0) + } +} + #[derive(Debug, Deserialize, Clone, PartialEq)] pub struct ProofDataHandlerConfig { pub http_port: u16, diff --git a/core/lib/protobuf_config/src/proof_data_handler.rs b/core/lib/protobuf_config/src/proof_data_handler.rs index 48dcf0c54f69..0347d71412a3 100644 --- a/core/lib/protobuf_config/src/proof_data_handler.rs +++ b/core/lib/protobuf_config/src/proof_data_handler.rs @@ -15,12 +15,8 @@ impl ProtoRepr for proto::ProofDataHandler { .and_then(|x| Ok((*x).try_into()?)) .context("proof_generation_timeout_in_secs")?, 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")?, + tee_support: configs::TeeConfig::default_tee_support(), + first_tee_processed_batch: configs::TeeConfig::default_first_tee_processed_batch(), }, }) } From 060f9a4f0c75959215326b38a480dcc4425a0325 Mon Sep 17 00:00:00 2001 From: Patrick Beza Date: Wed, 23 Oct 2024 11:42:31 +0200 Subject: [PATCH 5/6] fixup! Set defaults in conversion from the Protobuf type to the domain type --- core/lib/protobuf_config/src/proof_data_handler.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/core/lib/protobuf_config/src/proof_data_handler.rs b/core/lib/protobuf_config/src/proof_data_handler.rs index 0347d71412a3..a587c702633f 100644 --- a/core/lib/protobuf_config/src/proof_data_handler.rs +++ b/core/lib/protobuf_config/src/proof_data_handler.rs @@ -1,6 +1,7 @@ use anyhow::Context as _; use zksync_config::configs; use zksync_protobuf::{repr::ProtoRepr, required}; +use zksync_types::L1BatchNumber; use crate::proto::prover as proto; @@ -15,8 +16,13 @@ impl ProtoRepr for proto::ProofDataHandler { .and_then(|x| Ok((*x).try_into()?)) .context("proof_generation_timeout_in_secs")?, tee_config: configs::TeeConfig { - tee_support: configs::TeeConfig::default_tee_support(), - first_tee_processed_batch: configs::TeeConfig::default_first_tee_processed_batch(), + tee_support: self + .tee_support + .unwrap_or_else(configs::TeeConfig::default_tee_support), + first_tee_processed_batch: self + .first_tee_processed_batch + .map(|x| L1BatchNumber(x as u32)) + .unwrap_or_else(configs::TeeConfig::default_first_tee_processed_batch), }, }) } From d943314829de7e41bbbea6bc0d8cf30fdd83a6f1 Mon Sep 17 00:00:00 2001 From: Patrick Beza Date: Wed, 23 Oct 2024 12:04:54 +0200 Subject: [PATCH 6/6] Make tee_support optional (default: false) --- core/lib/protobuf_config/src/proto/config/prover.proto | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/lib/protobuf_config/src/proto/config/prover.proto b/core/lib/protobuf_config/src/proto/config/prover.proto index 5ce173d91dd9..92ba770a7560 100644 --- a/core/lib/protobuf_config/src/proto/config/prover.proto +++ b/core/lib/protobuf_config/src/proto/config/prover.proto @@ -107,6 +107,6 @@ message WitnessVectorGenerator { 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 bool tee_support = 3; // optional optional uint64 first_tee_processed_batch = 4; // optional }