diff --git a/Cargo.lock b/Cargo.lock index 9217b871..3af52fd7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3439,6 +3439,7 @@ dependencies = [ "proof_gen", "ruint", "serde", + "serde_json", "tokio", "trace_decoder", "tracing", diff --git a/common/src/fs.rs b/common/src/fs.rs new file mode 100644 index 00000000..7603f510 --- /dev/null +++ b/common/src/fs.rs @@ -0,0 +1,7 @@ +use std::path::PathBuf; + +pub fn generate_block_proof_file_name(directory: &Option<&str>, block_height: u64) -> PathBuf { + let mut path = PathBuf::from(directory.unwrap_or("")); + path.push(format!("b{}.zkproof", block_height)); + path +} diff --git a/common/src/lib.rs b/common/src/lib.rs index 635eba33..1b0b18e9 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -1,4 +1,5 @@ pub mod block_interval; pub mod debug_utils; +pub mod fs; pub mod parsing; pub mod prover_state; diff --git a/leader/src/cli.rs b/leader/src/cli.rs index 06a71b61..96f2c072 100644 --- a/leader/src/cli.rs +++ b/leader/src/cli.rs @@ -54,6 +54,15 @@ pub(crate) enum Command { /// to determine the blockchain node polling interval. #[arg(short, long, env = "ZERO_BIN_BLOCK_TIME", default_value_t = 2000)] block_time: u64, + /// Keep intermediate proofs. Default action is to + /// delete them after the final proof is generated. + #[arg( + short, + long, + env = "ZERO_BIN_KEEP_INTERMEDIATE_PROOFS", + default_value_t = false + )] + keep_intermediate_proofs: bool, }, /// Reads input from HTTP and writes output to a directory. Http { diff --git a/leader/src/jerigon.rs b/leader/src/jerigon.rs index f3cf1e74..ce03f615 100644 --- a/leader/src/jerigon.rs +++ b/leader/src/jerigon.rs @@ -1,61 +1,64 @@ -use std::{ - fs::{create_dir_all, File}, - io::Write, - path::PathBuf, -}; +use std::path::PathBuf; use alloy::providers::RootProvider; use anyhow::Result; use common::block_interval::BlockInterval; +use common::fs::generate_block_proof_file_name; use paladin::runtime::Runtime; use proof_gen::proof_types::GeneratedBlockProof; +use tracing::{error, warn}; + +#[derive(Debug, Default)] +pub struct ProofParams { + pub checkpoint_block_number: u64, + pub previous_proof: Option, + pub proof_output_dir: Option, + pub save_inputs_on_error: bool, + pub keep_intermediate_proofs: bool, +} /// The main function for the jerigon mode. pub(crate) async fn jerigon_main( runtime: Runtime, rpc_url: &str, block_interval: BlockInterval, - checkpoint_block_number: u64, - previous_proof: Option, - proof_output_dir_opt: Option, - save_inputs_on_error: bool, + mut params: ProofParams, ) -> Result<()> { let prover_input = rpc::prover_input( RootProvider::new_http(rpc_url.parse()?), block_interval, - checkpoint_block_number.into(), + params.checkpoint_block_number.into(), ) .await?; - let block_proofs = prover_input - .prove(&runtime, previous_proof, save_inputs_on_error) + // If `keep_intermediate_proofs` is not set we only keep the last block + // proof from the interval. It contains all the necessary information to + // verify the whole sequence. + let proved_blocks = prover_input + .prove( + &runtime, + params.previous_proof.take(), + params.save_inputs_on_error, + params.proof_output_dir.clone(), + ) .await?; runtime.close().await?; - for block_proof in block_proofs { - let block_proof_str = serde_json::to_vec(&block_proof)?; - write_proof( - block_proof_str, - proof_output_dir_opt.clone().map(|mut path| { - path.push(format!("b{}.zkproof", block_proof.b_height)); - path - }), - )?; - } - Ok(()) -} - -fn write_proof(proof: Vec, proof_output_dir_opt: Option) -> Result<()> { - match proof_output_dir_opt { - Some(p) => { - if let Some(parent) = p.parent() { - create_dir_all(parent)?; - } - - let mut f = File::create(p)?; - f.write_all(&proof)?; - } - None => std::io::stdout().write_all(&proof)?, + if params.keep_intermediate_proofs { + warn!("Skipping cleanup, intermediate proofs are kept"); + } else if let Some(proof_output_dir) = params.proof_output_dir.as_ref() { + proved_blocks + .into_iter() + .rev() + .skip(1) + .map(|b| generate_block_proof_file_name(&proof_output_dir.to_str(), b)) + .for_each(|path| { + if let Err(e) = std::fs::remove_file(path) { + error!("Failed to remove intermediate proof file: {e}"); + } + }); + } else { + // Proofs are written to stdio, so no need to clean up } Ok(()) diff --git a/leader/src/main.rs b/leader/src/main.rs index f03822c6..f04cfc57 100644 --- a/leader/src/main.rs +++ b/leader/src/main.rs @@ -11,6 +11,7 @@ use paladin::runtime::Runtime; use proof_gen::proof_types::GeneratedBlockProof; use tracing::info; +use crate::jerigon::{jerigon_main, ProofParams}; use crate::utils::get_package_version; mod cli; @@ -93,6 +94,7 @@ async fn main() -> Result<()> { proof_output_dir, save_inputs_on_error, block_time, + keep_intermediate_proofs, } => { let previous_proof = get_previous_proof(previous_proof)?; let mut block_interval = BlockInterval::new(&block_interval)?; @@ -106,15 +108,17 @@ async fn main() -> Result<()> { } info!("Proving interval {block_interval}"); - - jerigon::jerigon_main( + jerigon_main( runtime, &rpc_url, block_interval, - checkpoint_block_number, - previous_proof, - proof_output_dir, - save_inputs_on_error, + ProofParams { + checkpoint_block_number, + previous_proof, + proof_output_dir, + save_inputs_on_error, + keep_intermediate_proofs, + }, ) .await?; } diff --git a/prover/Cargo.toml b/prover/Cargo.toml index 05dc591d..ff5c1227 100644 --- a/prover/Cargo.toml +++ b/prover/Cargo.toml @@ -18,6 +18,7 @@ anyhow = { workspace = true } futures = { workspace = true } alloy.workspace = true tokio = {workspace = true} +serde_json = {workspace = true} ruint = { version = "1.12.1", features = ["num-traits", "primitive-types"] } ops = { path = "../ops" } common = { path = "../common" } diff --git a/prover/src/lib.rs b/prover/src/lib.rs index 9f0d528d..0758ea1f 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -1,7 +1,9 @@ use std::future::Future; +use std::path::PathBuf; -use alloy::primitives::U256; -use anyhow::Result; +use alloy::primitives::{BlockNumber, U256}; +use anyhow::{Context, Result}; +use common::fs::generate_block_proof_file_name; use futures::{future::BoxFuture, stream::FuturesOrdered, FutureExt, TryFutureExt, TryStreamExt}; use num_traits::ToPrimitive as _; use ops::TxProof; @@ -11,6 +13,7 @@ use paladin::{ }; use proof_gen::proof_types::GeneratedBlockProof; use serde::{Deserialize, Serialize}; +use tokio::io::AsyncWriteExt; use tokio::sync::oneshot; use trace_decoder::{ processed_block_trace::ProcessingMeta, @@ -133,7 +136,8 @@ impl ProverInput { runtime: &Runtime, previous_proof: Option, save_inputs_on_error: bool, - ) -> Result> { + proof_output_dir: Option, + ) -> Result> { let mut prev: Option>> = previous_proof.map(|proof| Box::pin(futures::future::ok(proof)) as BoxFuture<_>); @@ -147,16 +151,21 @@ impl ProverInput { let (tx, rx) = oneshot::channel::(); // Prove the block + let proof_output_dir = proof_output_dir.clone(); let fut = block .prove(runtime, prev.take(), save_inputs_on_error) - .then(|proof| async { + .then(move |proof| async move { let proof = proof?; + let block_number = proof.b_height; - if tx.send(proof.clone()).is_err() { + // Write latest generated proof to disk or stdout + ProverInput::write_proof(proof_output_dir, &proof).await?; + + if tx.send(proof).is_err() { anyhow::bail!("Failed to send proof"); } - Ok(proof) + Ok(block_number) }) .boxed(); @@ -168,4 +177,30 @@ impl ProverInput { results.try_collect().await } + + /// Write the proof to the disk (if `output_dir` is provided) or stdout. + pub(crate) async fn write_proof( + output_dir: Option, + proof: &GeneratedBlockProof, + ) -> Result<()> { + let proof_serialized = serde_json::to_vec(proof)?; + let block_proof_file_path = + output_dir.map(|path| generate_block_proof_file_name(&path.to_str(), proof.b_height)); + match block_proof_file_path { + Some(p) => { + if let Some(parent) = p.parent() { + tokio::fs::create_dir_all(parent).await?; + } + + let mut f = tokio::fs::File::create(p).await?; + f.write_all(&proof_serialized) + .await + .context("Failed to write proof to disk") + } + None => tokio::io::stdout() + .write_all(&proof_serialized) + .await + .context("Failed to write proof to stdout"), + } + } }