Improve light client initialization (#1611)

Cargo.toml

@@ -148,7 +148,7 @@ alloy = { version = "0.2.1", default-features = false }
 alloy-eips = { version = "0.2.1", default-features = false }
 alloy-consensus = { version = "0.2.1", default-features = false }
 
-citrea-e2e = { git = "https://github.com/chainwayxyz/citrea-e2e", rev = "efde07d" }
+citrea-e2e = { git = "https://github.com/chainwayxyz/citrea-e2e", rev = "51a4d19" }
 
 [patch.crates-io]
 bitcoincore-rpc = { version = "0.18.0", git = "https://github.com/chainwayxyz/rust-bitcoincore-rpc.git", rev = "ca3cfa2" }

bin/citrea/src/main.rs

@@ -256,13 +256,19 @@ where
             error!("Error: {}", e);
         }
     } else if let Some(light_client_prover_config) = light_client_prover_config {
-        let prover = CitreaRollupBlueprint::create_new_light_client_prover(
+        let (mut prover, rpc_methods) = CitreaRollupBlueprint::create_new_light_client_prover(
             &rollup_blueprint,
             rollup_config,
             light_client_prover_config,
         )
         .await
         .expect("Could not start light client prover");
+
+        prover
+            .start_rpc_server(rpc_methods, None)
+            .await
+            .expect("Failed to start rpc server");
+
         if let Err(e) = prover.run().await {
             error!("Error: {}", e);
         }

bin/citrea/src/rollup/mod.rs

@@ -6,7 +6,7 @@ use citrea_batch_prover::CitreaBatchProver;
 use citrea_common::tasks::manager::TaskManager;
 use citrea_common::{BatchProverConfig, FullNodeConfig, LightClientProverConfig, SequencerConfig};
 use citrea_fullnode::CitreaFullnode;
-use citrea_light_client_prover::runner::{CitreaLightClientProver, LightClientProver};
+use citrea_light_client_prover::runner::CitreaLightClientProver;
 use citrea_primitives::forks::FORKS;
 use citrea_sequencer::CitreaSequencer;
 use jsonrpsee::RpcModule;
@@ -424,7 +424,13 @@ pub trait CitreaRollupBlueprint: RollupBlueprint {
         &self,
         rollup_config: FullNodeConfig<Self::DaConfig>,
         prover_config: LightClientProverConfig,
-    ) -> Result<LightClientProver<Self>, anyhow::Error>
+    ) -> Result<
+        (
+            CitreaLightClientProver<Self::DaService, Self::Vm, Self::ProverService, LedgerDB>,
+            RpcModule<()>,
+        ),
+        anyhow::Error,
+    >
     where
         <Self::NativeContext as Spec>::Storage: NativeStorage,
     {
@@ -501,9 +507,6 @@ pub trait CitreaRollupBlueprint: RollupBlueprint {
             task_manager,
         )?;
 
-        Ok(LightClientProver {
-            runner,
-            rpc_methods,
-        })
+        Ok((runner, rpc_methods))
     }
 }

bin/citrea/tests/bitcoin_e2e/light_client_test.rs

@@ -53,6 +53,7 @@ impl TestCase for LightClientProvingTest {
     fn light_client_prover_config() -> LightClientProverConfig {
         LightClientProverConfig {
             enable_recovery: false,
+            initial_da_height: 171,
             ..Default::default()
         }
     }
@@ -187,6 +188,7 @@ impl TestCase for LightClientProvingTestMultipleProofs {
     fn light_client_prover_config() -> LightClientProverConfig {
         LightClientProverConfig {
             enable_recovery: false,
+            initial_da_height: 171,
             ..Default::default()
         }
     }

bin/citrea/tests/test_helpers/mod.rs

@@ -107,19 +107,22 @@ pub async fn start_rollup(
         rollup.run().instrument(span).await.unwrap();
     } else if let Some(light_client_prover_config) = light_client_prover_config {
         let span = info_span!("LightClientProver");
-        let rollup = CitreaRollupBlueprint::create_new_light_client_prover(
+        let (mut rollup, rpc_methods) = CitreaRollupBlueprint::create_new_light_client_prover(
             &mock_demo_rollup,
             rollup_config.clone(),
             light_client_prover_config,
         )
         .instrument(span.clone())
         .await
         .unwrap();
+
         rollup
-            .run_and_report_rpc_port(Some(rpc_reporting_channel))
-            .instrument(span)
+            .start_rpc_server(rpc_methods, Some(rpc_reporting_channel))
+            .instrument(span.clone())
             .await
             .unwrap();
+
+        rollup.run().instrument(span).await.unwrap();
     } else {
         let span = info_span!("FullNode");
         let (mut rollup, rpc_methods) = CitreaRollupBlueprint::create_new_rollup(

crates/batch-prover/src/runner.rs

@@ -8,7 +8,7 @@ use anyhow::{anyhow, bail, Context as _};
 use backoff::exponential::ExponentialBackoffBuilder;
 use backoff::future::retry as retry_backoff;
 use citrea_common::cache::L1BlockCache;
-use citrea_common::da::get_da_block_at_height;
+use citrea_common::da::{get_da_block_at_height, get_initial_slot_height};
 use citrea_common::tasks::manager::TaskManager;
 use citrea_common::utils::{create_shutdown_signal, soft_confirmation_to_receipt};
 use citrea_common::{BatchProverConfig, RollupPublicKeys, RpcConfig, RunnerConfig};
@@ -583,16 +583,3 @@ async fn sync_l2(
         }
     }
 }
-
-async fn get_initial_slot_height(client: &HttpClient) -> u64 {
-    loop {
-        match client.get_soft_confirmation_by_number(U64::from(1)).await {
-            Ok(Some(batch)) => return batch.da_slot_height,
-            _ => {
-                // sleep 1
-                tokio::time::sleep(std::time::Duration::from_secs(1)).await;
-                continue;
-            }
-        }
-    }
-}

crates/common/Cargo.toml

@@ -20,6 +20,7 @@ hex = { workspace = true }
 hyper = { workspace = true }
 jsonrpsee = { workspace = true, features = ["http-client", "server"] }
 lru = { workspace = true }
+reth-primitives = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 tokio = { workspace = true }
@@ -30,6 +31,7 @@ tracing = { workspace = true }
 
 # Sov SDK deps
 sov-db = { path = "../sovereign-sdk/full-node/db/sov-db" }
+sov-ledger-rpc = { path = "../sovereign-sdk/full-node/sov-ledger-rpc", features = ["client"] }
 sov-mock-da = { path = "../sovereign-sdk/adapters/mock-da" }
 sov-modules-api = { path = "../sovereign-sdk/module-system/sov-modules-api" }
 sov-rollup-interface = { path = "../sovereign-sdk/rollup-interface" }

crates/common/src/config.rs

@@ -253,6 +253,8 @@ pub struct LightClientProverConfig {
     pub proof_sampling_number: usize,
     /// If true prover will try to recover ongoing proving sessions
     pub enable_recovery: bool,
+    /// The starting DA block to sync from
+    pub initial_da_height: u64,
 }
 
 impl Default for BatchProverConfig {
@@ -271,6 +273,7 @@ impl Default for LightClientProverConfig {
             proving_mode: ProverGuestRunConfig::Execute,
             proof_sampling_number: 0,
             enable_recovery: true,
+            initial_da_height: 1,
         }
     }
 }
@@ -291,6 +294,7 @@ impl FromEnv for LightClientProverConfig {
             proving_mode: serde_json::from_str(&format!("\"{}\"", std::env::var("PROVING_MODE")?))?,
             proof_sampling_number: std::env::var("PROOF_SAMPLING_NUMBER")?.parse()?,
             enable_recovery: std::env::var("ENABLE_RECOVERY")?.parse()?,
+            initial_da_height: std::env::var("INITIAL_DA_HEIGHT")?.parse()?,
         })
     }
 }

crates/common/src/da.rs

@@ -4,6 +4,9 @@ use std::time::Duration;
 use anyhow::anyhow;
 use backoff::future::retry as retry_backoff;
 use backoff::ExponentialBackoffBuilder;
+use jsonrpsee::http_client::HttpClient;
+use reth_primitives::U64;
+use sov_ledger_rpc::LedgerRpcClient;
 use sov_rollup_interface::da::{BlockHeaderTrait, SequencerCommitment};
 use sov_rollup_interface::services::da::{DaService, SlotData};
 use sov_rollup_interface::zk::Proof;
@@ -72,3 +75,16 @@ pub async fn extract_zk_proofs<Da: DaService>(
         .extract_relevant_zk_proofs(l1_block, prover_da_pub_key)
         .await
 }
+
+pub async fn get_initial_slot_height(client: &HttpClient) -> u64 {
+    loop {
+        match client.get_soft_confirmation_by_number(U64::from(1)).await {
+            Ok(Some(batch)) => return batch.da_slot_height,
+            _ => {
+                // sleep 1
+                tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+                continue;
+            }
+        }
+    }
+}

crates/light-client-prover/Cargo.toml

@@ -17,7 +17,6 @@ citrea-primitives = { path = "../primitives", optional = true }
 sov-db = { path = "../sovereign-sdk/full-node/db/sov-db", optional = true }
 sov-ledger-rpc = { path = "../sovereign-sdk/full-node/sov-ledger-rpc", features = ["client"], optional = true }
 sov-modules-api = { path = "../sovereign-sdk/module-system/sov-modules-api", default-features = false }
-sov-modules-rollup-blueprint = { path = "../sovereign-sdk/module-system/sov-modules-rollup-blueprint", optional = true }
 sov-rollup-interface = { path = "../sovereign-sdk/rollup-interface" }
 sov-stf-runner = { path = "../sovereign-sdk/full-node/sov-stf-runner", optional = true }
 
@@ -48,7 +47,6 @@ native = [
   "dep:citrea-primitives",
   "dep:citrea-common",
   "dep:sov-db",
-  "dep:sov-modules-rollup-blueprint",
   "dep:sov-stf-runner",
   "dep:sov-ledger-rpc",
   "dep:anyhow",

crates/light-client-prover/src/runner.rs

@@ -7,11 +7,8 @@ use citrea_common::{LightClientProverConfig, RollupPublicKeys, RpcConfig, Runner
 use jsonrpsee::http_client::{HttpClient, HttpClientBuilder};
 use jsonrpsee::server::{BatchRequestConfig, ServerBuilder};
 use jsonrpsee::RpcModule;
-use reth_primitives::U64;
-use sov_db::ledger_db::{LedgerDB, LightClientProverLedgerOps, SharedLedgerOps};
+use sov_db::ledger_db::{LightClientProverLedgerOps, SharedLedgerOps};
 use sov_db::schema::types::SlotNumber;
-use sov_ledger_rpc::LedgerRpcClient;
-use sov_modules_rollup_blueprint::RollupBlueprint;
 use sov_rollup_interface::services::da::DaService;
 use sov_rollup_interface::spec::SpecId;
 use sov_rollup_interface::zk::ZkvmHost;
@@ -23,41 +20,6 @@ use tracing::{error, info, instrument};
 use crate::da_block_handler::L1BlockHandler;
 use crate::rpc::{create_rpc_module, RpcContext};
 
-/// Dependencies needed to run the rollup.
-pub struct LightClientProver<S: RollupBlueprint> {
-    /// The State Transition Runner.
-    #[allow(clippy::type_complexity)]
-    pub runner: CitreaLightClientProver<S::DaService, S::Vm, S::ProverService, LedgerDB>,
-    /// Rpc methods for the rollup.
-    pub rpc_methods: jsonrpsee::RpcModule<()>,
-}
-
-impl<S: RollupBlueprint> LightClientProver<S> {
-    /// Runs the rollup.
-    #[instrument(level = "trace", skip_all, err, ret(level = "error"))]
-    pub async fn run(self) -> Result<(), anyhow::Error> {
-        self.run_and_report_rpc_port(None).await
-    }
-
-    /// Only run the rpc.
-    pub async fn run_rpc(mut self) -> Result<(), anyhow::Error> {
-        self.runner.start_rpc_server(self.rpc_methods, None).await?;
-        Ok(())
-    }
-
-    /// Runs the rollup. Reports rpc port to the caller using the provided channel.
-    pub async fn run_and_report_rpc_port(
-        self,
-        channel: Option<oneshot::Sender<SocketAddr>>,
-    ) -> Result<(), anyhow::Error> {
-        let mut runner = self.runner;
-        runner.start_rpc_server(self.rpc_methods, channel).await?;
-
-        runner.run().await?;
-        Ok(())
-    }
-}
-
 pub struct CitreaLightClientProver<Da, Vm, Ps, DB>
 where
     Da: DaService + Send + Sync,
@@ -186,7 +148,7 @@ where
         let last_l1_height_scanned = match self.ledger_db.get_last_scanned_l1_height()? {
             Some(l1_height) => l1_height,
             // If not found, start from the first L2 block's L1 height
-            None => SlotNumber(get_initial_da_height(&self.sequencer_client).await),
+            None => SlotNumber(self.prover_config.initial_da_height),
         };
 
         let prover_config = self.prover_config.clone();
@@ -250,16 +212,3 @@ where
         Ok(rpc_methods)
     }
 }
-
-async fn get_initial_da_height(client: &HttpClient) -> u64 {
-    loop {
-        match client.get_soft_confirmation_by_number(U64::from(1)).await {
-            Ok(Some(batch)) => return batch.da_slot_height,
-            _ => {
-                // sleep 1
-                tokio::time::sleep(std::time::Duration::from_secs(1)).await;
-                continue;
-            }
-        }
-    }
-}