feat: Split witness generator timeout configs by round (#1505)

## What ❔

Have different values in configs for witness generation timeouts by
round

## Why ❔

To be able to set custom timeouts by round(sometimes times of generation
differs for different rounds)

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [x] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [x] Tests for the changes have been added / updated.
- [x] Documentation comments have been added / updated.
- [x] Code has been formatted via `zk fmt` and `zk lint`.
- [x] Spellcheck has been run via `zk spellcheck`.
- [x] Linkcheck has been run via `zk linkcheck`.

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

@@ -9,6 +9,10 @@ use serde::Deserialize;
 pub struct FriWitnessGeneratorConfig {
     /// Max time for witness to be generated
     pub generation_timeout_in_secs: u16,
+    pub basic_generation_timeout_in_secs: Option<u16>,
+    pub leaf_generation_timeout_in_secs: Option<u16>,
+    pub node_generation_timeout_in_secs: Option<u16>,
+    pub scheduler_generation_timeout_in_secs: Option<u16>,
     /// Max attempts for generating witness
     pub max_attempts: u32,
     // Percentage of the blocks that gets proven in the range [0.0, 1.0]
@@ -25,9 +29,54 @@ pub struct FriWitnessGeneratorConfig {
     // whether to write to public GCS bucket for https://github.com/matter-labs/era-boojum-validator-cli
     pub shall_save_to_public_bucket: bool,
 }
+
+#[derive(Debug)]
+pub struct WitnessGenerationTimeouts {
+    basic: Duration,
+    leaf: Duration,
+    node: Duration,
+    scheduler: Duration,
+}
+
+impl WitnessGenerationTimeouts {
+    pub fn basic(&self) -> Duration {
+        self.basic
+    }
+
+    pub fn leaf(&self) -> Duration {
+        self.leaf
+    }
+
+    pub fn node(&self) -> Duration {
+        self.node
+    }
+
+    pub fn scheduler(&self) -> Duration {
+        self.scheduler
+    }
+
+    pub fn new(basic: u16, leaf: u16, node: u16, scheduler: u16) -> Self {
+        Self {
+            basic: Duration::from_secs(basic as u64),
+            leaf: Duration::from_secs(leaf as u64),
+            node: Duration::from_secs(node as u64),
+            scheduler: Duration::from_secs(scheduler as u64),
+        }
+    }
+}
+
 impl FriWitnessGeneratorConfig {
-    pub fn witness_generation_timeout(&self) -> Duration {
-        Duration::from_secs(self.generation_timeout_in_secs as u64)
+    pub fn witness_generation_timeouts(&self) -> WitnessGenerationTimeouts {
+        WitnessGenerationTimeouts::new(
+            self.basic_generation_timeout_in_secs
+                .unwrap_or(self.generation_timeout_in_secs),
+            self.leaf_generation_timeout_in_secs
+                .unwrap_or(self.generation_timeout_in_secs),
+            self.node_generation_timeout_in_secs
+                .unwrap_or(self.generation_timeout_in_secs),
+            self.scheduler_generation_timeout_in_secs
+                .unwrap_or(self.generation_timeout_in_secs),
+        )
     }
 
     pub fn last_l1_batch_to_process(&self) -> u32 {

core/lib/config/src/testonly.rs

@@ -617,6 +617,10 @@ impl RandomConfig for configs::FriWitnessGeneratorConfig {
     fn sample(g: &mut Gen<impl Rng>) -> Self {
         Self {
             generation_timeout_in_secs: g.gen(),
+            basic_generation_timeout_in_secs: g.gen(),
+            leaf_generation_timeout_in_secs: g.gen(),
+            node_generation_timeout_in_secs: g.gen(),
+            scheduler_generation_timeout_in_secs: g.gen(),
             max_attempts: g.gen(),
             blocks_proving_percentage: g.gen(),
             dump_arguments_for_blocks: g.gen(),

core/lib/env_config/src/fri_witness_generator.rs

@@ -18,6 +18,10 @@ mod tests {
     fn expected_config() -> FriWitnessGeneratorConfig {
         FriWitnessGeneratorConfig {
             generation_timeout_in_secs: 900u16,
+            basic_generation_timeout_in_secs: Some(900u16),
+            leaf_generation_timeout_in_secs: Some(800u16),
+            node_generation_timeout_in_secs: Some(800u16),
+            scheduler_generation_timeout_in_secs: Some(900u16),
             max_attempts: 4,
             blocks_proving_percentage: Some(30),
             dump_arguments_for_blocks: vec![2, 3],
@@ -32,6 +36,10 @@ mod tests {
         let mut lock = MUTEX.lock();
         let config = r#"
             FRI_WITNESS_GENERATION_TIMEOUT_IN_SECS=900
+            FRI_WITNESS_BASIC_GENERATION_TIMEOUT_IN_SECS=900
+            FRI_WITNESS_LEAF_GENERATION_TIMEOUT_IN_SECS=800
+            FRI_WITNESS_NODE_GENERATION_TIMEOUT_IN_SECS=800
+            FRI_WITNESS_SCHEDULER_GENERATION_TIMEOUT_IN_SECS=900
             FRI_WITNESS_MAX_ATTEMPTS=4
             FRI_WITNESS_DUMP_ARGUMENTS_FOR_BLOCKS="2,3"
             FRI_WITNESS_BLOCKS_PROVING_PERCENTAGE="30"
@@ -43,4 +51,32 @@ mod tests {
         let actual = FriWitnessGeneratorConfig::from_env().unwrap();
         assert_eq!(actual, expected_config());
     }
+
+    #[test]
+    fn from_env_default_timeouts() {
+        let mut lock = MUTEX.lock();
+        lock.remove_env(&[
+            "FRI_WITNESS_LEAF_GENERATION_TIMEOUT_IN_SECS",
+            "FRI_WITNESS_NODE_GENERATION_TIMEOUT_IN_SECS",
+        ]);
+        let config = r#"
+            FRI_WITNESS_GENERATION_TIMEOUT_IN_SECS=800
+            FRI_WITNESS_BASIC_GENERATION_TIMEOUT_IN_SECS=100
+            FRI_WITNESS_SCHEDULER_GENERATION_TIMEOUT_IN_SECS=200
+            FRI_WITNESS_MAX_ATTEMPTS=4
+            FRI_WITNESS_DUMP_ARGUMENTS_FOR_BLOCKS="2,3"
+            FRI_WITNESS_BLOCKS_PROVING_PERCENTAGE="30"
+            FRI_WITNESS_FORCE_PROCESS_BLOCK="1"
+            FRI_WITNESS_SHALL_SAVE_TO_PUBLIC_BUCKET=true
+        "#;
+        lock.set_env(config);
+
+        let actual = FriWitnessGeneratorConfig::from_env().unwrap();
+        let timeouts = actual.witness_generation_timeouts();
+
+        assert_eq!(timeouts.basic().as_secs(), 100);
+        assert_eq!(timeouts.leaf().as_secs(), 800);
+        assert_eq!(timeouts.node().as_secs(), 800);
+        assert_eq!(timeouts.scheduler().as_secs(), 200);
+    }
 }

core/lib/protobuf_config/src/fri_witness_generator.rs

@@ -10,7 +10,27 @@ impl ProtoRepr for proto::FriWitnessGenerator {
         Ok(Self::Type {
             generation_timeout_in_secs: required(&self.generation_timeout_in_secs)
                 .and_then(|x| Ok((*x).try_into()?))
-                .context("generation_timeout_in_secs")?,
+                .context("basic_generation_timeout_in_secs")?,
+            basic_generation_timeout_in_secs: self
+                .basic_generation_timeout_in_secs
+                .map(|x| x.try_into())
+                .transpose()
+                .context("basic_generation_timeout_in_secs")?,
+            leaf_generation_timeout_in_secs: self
+                .leaf_generation_timeout_in_secs
+                .map(|x| x.try_into())
+                .transpose()
+                .context("leaf_generation_timeout_in_secs")?,
+            node_generation_timeout_in_secs: self
+                .node_generation_timeout_in_secs
+                .map(|x| x.try_into())
+                .transpose()
+                .context("node_generation_timeout_in_secs")?,
+            scheduler_generation_timeout_in_secs: self
+                .scheduler_generation_timeout_in_secs
+                .map(|x| x.try_into())
+                .transpose()
+                .context("scheduler_generation_timeout_in_secs")?,
             max_attempts: *required(&self.max_attempts).context("max_attempts")?,
             blocks_proving_percentage: self
                 .blocks_proving_percentage
@@ -28,6 +48,14 @@ impl ProtoRepr for proto::FriWitnessGenerator {
     fn build(this: &Self::Type) -> Self {
         Self {
             generation_timeout_in_secs: Some(this.generation_timeout_in_secs.into()),
+            basic_generation_timeout_in_secs: this
+                .basic_generation_timeout_in_secs
+                .map(|x| x.into()),
+            leaf_generation_timeout_in_secs: this.leaf_generation_timeout_in_secs.map(|x| x.into()),
+            node_generation_timeout_in_secs: this.node_generation_timeout_in_secs.map(|x| x.into()),
+            scheduler_generation_timeout_in_secs: this
+                .scheduler_generation_timeout_in_secs
+                .map(|x| x.into()),
             max_attempts: Some(this.max_attempts),
             blocks_proving_percentage: this.blocks_proving_percentage.map(|x| x.into()),
             dump_arguments_for_blocks: this.dump_arguments_for_blocks.clone(),

core/lib/protobuf_config/src/proto/fri_witness_generator.proto

@@ -3,9 +3,13 @@ syntax = "proto3";
 package zksync.config.fri_witness_generator;
 
 message FriWitnessGenerator {
-    optional uint32 generation_timeout_in_secs = 1; // required; s
+    optional uint32 generation_timeout_in_secs = 1; // required;
+    optional uint32 basic_generation_timeout_in_secs = 8; // optional;
+    optional uint32 leaf_generation_timeout_in_secs = 9; // optional;
+    optional uint32 node_generation_timeout_in_secs = 10; // optional;
+    optional uint32 scheduler_generation_timeout_in_secs = 11; // optional;
     optional uint32 max_attempts = 2; // required
-    optional uint32 blocks_proving_percentage = 3; // optional; 0-100 
+    optional uint32 blocks_proving_percentage = 3; // optional; 0-100
     repeated uint32 dump_arguments_for_blocks = 4;
     optional uint32 last_l1_batch_to_process = 5; // optional
     optional uint32 force_process_block = 6; // optional

core/lib/zksync_core/src/house_keeper/fri_witness_generator_jobs_retry_manager.rs

@@ -1,7 +1,6 @@
-use std::time::Duration;
-
 use async_trait::async_trait;
 use prover_dal::{Prover, ProverDal};
+use zksync_config::configs::fri_witness_generator::WitnessGenerationTimeouts;
 use zksync_dal::ConnectionPool;
 
 use crate::house_keeper::periodic_job::PeriodicJob;
@@ -10,20 +9,20 @@ use crate::house_keeper::periodic_job::PeriodicJob;
 pub struct FriWitnessGeneratorJobRetryManager {
     pool: ConnectionPool<Prover>,
     max_attempts: u32,
-    processing_timeout: Duration,
+    processing_timeouts: WitnessGenerationTimeouts,
     retry_interval_ms: u64,
 }
 
 impl FriWitnessGeneratorJobRetryManager {
     pub fn new(
         max_attempts: u32,
-        processing_timeout: Duration,
+        processing_timeouts: WitnessGenerationTimeouts,
         retry_interval_ms: u64,
         pool: ConnectionPool<Prover>,
     ) -> Self {
         Self {
             max_attempts,
-            processing_timeout,
+            processing_timeouts,
             retry_interval_ms,
             pool,
         }
@@ -36,7 +35,7 @@ impl FriWitnessGeneratorJobRetryManager {
             .await
             .unwrap()
             .fri_witness_generator_dal()
-            .requeue_stuck_jobs(self.processing_timeout, self.max_attempts)
+            .requeue_stuck_jobs(self.processing_timeouts.basic(), self.max_attempts)
             .await;
         let job_len = stuck_jobs.len();
         for stuck_job in stuck_jobs {
@@ -52,7 +51,10 @@ impl FriWitnessGeneratorJobRetryManager {
             .await
             .unwrap()
             .fri_witness_generator_dal()
-            .requeue_stuck_leaf_aggregations_jobs(self.processing_timeout, self.max_attempts)
+            .requeue_stuck_leaf_aggregations_jobs(
+                self.processing_timeouts.leaf(),
+                self.max_attempts,
+            )
             .await;
         let job_len = stuck_jobs.len();
         for stuck_job in stuck_jobs {
@@ -71,7 +73,10 @@ impl FriWitnessGeneratorJobRetryManager {
             .await
             .unwrap()
             .fri_witness_generator_dal()
-            .requeue_stuck_node_aggregations_jobs(self.processing_timeout, self.max_attempts)
+            .requeue_stuck_node_aggregations_jobs(
+                self.processing_timeouts.node(),
+                self.max_attempts,
+            )
             .await;
         let job_len = stuck_jobs.len();
         for stuck_job in stuck_jobs {
@@ -90,7 +95,7 @@ impl FriWitnessGeneratorJobRetryManager {
             .await
             .unwrap()
             .fri_witness_generator_dal()
-            .requeue_stuck_scheduler_jobs(self.processing_timeout, self.max_attempts)
+            .requeue_stuck_scheduler_jobs(self.processing_timeouts.scheduler(), self.max_attempts)
             .await;
         let job_len = stuck_jobs.len();
         for stuck_job in stuck_jobs {

core/lib/zksync_core/src/lib.rs

@@ -1007,7 +1007,7 @@ async fn add_house_keeper_to_task_futures(
         .context("fri_witness_generator_config")?;
     let fri_witness_gen_job_retry_manager = FriWitnessGeneratorJobRetryManager::new(
         fri_witness_gen_config.max_attempts,
-        fri_witness_gen_config.witness_generation_timeout(),
+        fri_witness_gen_config.witness_generation_timeouts(),
         house_keeper_config.fri_witness_generator_job_retrying_interval_ms,
         prover_connection_pool.clone(),
     );

core/node/node_framework/src/implementations/layers/house_keeper.rs

@@ -94,7 +94,7 @@ impl WiringLayer for HouseKeeperLayer {
         let fri_witness_gen_job_retry_manager = FriWitnessGeneratorJobRetryManager::new(
             self.fri_witness_generator_config.max_attempts,
             self.fri_witness_generator_config
-                .witness_generation_timeout(),
+                .witness_generation_timeouts(),
             self.house_keeper_config
                 .fri_witness_generator_job_retrying_interval_ms,
             prover_pool.clone(),

etc/env/base/fri_witness_generator.toml

@@ -1,6 +1,10 @@
 [fri_witness]
-generation_timeout_in_secs=900
-max_attempts=10
-dump_arguments_for_blocks="1"
-force_process_block=1
-shall_save_to_public_bucket=true
+generation_timeout_in_secs = 900
+basic_generation_timeout_in_secs = 900
+leaf_generation_timeout_in_secs = 900
+node_generation_timeout_in_secs = 900
+scheduler_generation_timeout_in_secs = 900
+max_attempts = 10
+dump_arguments_for_blocks = "1"
+force_process_block = 1
+shall_save_to_public_bucket = true