From f13b6cfd3eb64d96815dae0be216fd889c84611c Mon Sep 17 00:00:00 2001
From: Arpit Temani <temaniarpit27@gmail.com>
Date: Thu, 19 Sep 2024 18:45:07 +0530
Subject: [PATCH] make pool size configurable (#644)

* make pool size configurable

* review comments
---
 zero/src/bin/leader.rs | 3 +++
 zero/src/prover.rs     | 8 +++++---
 zero/src/prover/cli.rs | 5 +++++
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/zero/src/bin/leader.rs b/zero/src/bin/leader.rs
index abf19bde3..728f6bb0e 100644
--- a/zero/src/bin/leader.rs
+++ b/zero/src/bin/leader.rs
@@ -38,6 +38,9 @@ async fn main() -> Result<()> {
 
     let runtime = Arc::new(Runtime::from_config(&args.paladin, register()).await?);
     let prover_config: ProverConfig = args.prover_config.into();
+    if prover_config.block_pool_size == 0 {
+        panic!("block-pool-size must be greater than 0");
+    }
 
     // If not in test_only mode and running in emulation mode, we'll need to
     // initialize the prover state here.
diff --git a/zero/src/prover.rs b/zero/src/prover.rs
index 3916f98d2..02360e0f7 100644
--- a/zero/src/prover.rs
+++ b/zero/src/prover.rs
@@ -31,9 +31,7 @@ use crate::ops;
 //
 // While proving a block interval, we will output proofs corresponding to block
 // batches as soon as they are generated.
-const PARALLEL_BLOCK_PROVING_PERMIT_POOL_SIZE: usize = 16;
-static PARALLEL_BLOCK_PROVING_PERMIT_POOL: Semaphore =
-    Semaphore::const_new(PARALLEL_BLOCK_PROVING_PERMIT_POOL_SIZE);
+static PARALLEL_BLOCK_PROVING_PERMIT_POOL: Semaphore = Semaphore::const_new(0);
 
 #[derive(Debug, Clone)]
 pub struct ProverConfig {
@@ -44,6 +42,7 @@ pub struct ProverConfig {
     pub proof_output_dir: PathBuf,
     pub keep_intermediate_proofs: bool,
     pub block_batch_size: usize,
+    pub block_pool_size: usize,
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
@@ -242,6 +241,9 @@ pub async fn prove(
     let mut task_set: JoinSet<
         std::result::Result<std::result::Result<u64, anyhow::Error>, anyhow::Error>,
     > = JoinSet::new();
+
+    PARALLEL_BLOCK_PROVING_PERMIT_POOL.add_permits(prover_config.block_pool_size);
+
     while let Some((block_prover_input, is_last_block)) = block_receiver.recv().await {
         block_counter += 1;
         let (tx, rx) = oneshot::channel::<GeneratedBlockProof>();
diff --git a/zero/src/prover/cli.rs b/zero/src/prover/cli.rs
index e55141b7a..a6cdaebf9 100644
--- a/zero/src/prover/cli.rs
+++ b/zero/src/prover/cli.rs
@@ -43,6 +43,10 @@ pub struct CliProverConfig {
     /// generate one proof file.
     #[arg(long, default_value_t = 8)]
     block_batch_size: usize,
+    /// The maximum number of block proving tasks that can run in parallel. Must
+    /// be greater than zero.
+    #[arg(long, default_value_t = 16)]
+    block_pool_size: usize,
 }
 
 impl From<CliProverConfig> for super::ProverConfig {
@@ -55,6 +59,7 @@ impl From<CliProverConfig> for super::ProverConfig {
             proof_output_dir: cli.proof_output_dir,
             keep_intermediate_proofs: cli.keep_intermediate_proofs,
             block_batch_size: cli.block_batch_size,
+            block_pool_size: cli.block_pool_size,
         }
     }
 }