From 11850aacab82a96cb6d06e0277bd7d3a9d12976c Mon Sep 17 00:00:00 2001 From: zeapoz Date: Thu, 7 Dec 2023 13:59:02 +0100 Subject: [PATCH] tweak: store the last processed block in snapshot Before, we stored the last fetched block and waited until it had processed fully before terminating the main fetching loop. When the fetcher would repeatedly fail to procure a transaction, it would also delay the shutdown process. This change updates the snapshot with the last *processed* block instead. --- Cargo.lock | 1 + src/main.rs | 9 +++- src/processor/tree/mod.rs | 50 ++++++------------ state-reconstruct-fetcher/Cargo.toml | 1 + state-reconstruct-fetcher/src/l1_fetcher.rs | 57 +++++++++++++++------ 5 files changed, 65 insertions(+), 53 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e035f03..1e94ef8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5078,6 +5078,7 @@ dependencies = [ "serde_json_any_key", "thiserror", "tokio", + "tokio-util", "tracing", ] diff --git a/src/main.rs b/src/main.rs index 4be9674..7854edd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -18,7 +18,7 @@ use cli::{Cli, Command, ReconstructSource}; use eyre::Result; use processor::snapshot::{SnapshotBuilder, SnapshotExporter}; use state_reconstruct_fetcher::{ - constants::storage, + constants::storage::{self, STATE_FILE_NAME}, l1_fetcher::{L1Fetcher, L1FetcherOptions}, snapshot::StateSnapshot, types::CommitBlockInfoV1, @@ -78,7 +78,7 @@ async fn main() -> Result<()> { }; let fetcher = L1Fetcher::new(fetcher_options, Some(snapshot.clone()))?; - let processor = TreeProcessor::new(db_path, snapshot.clone()).await?; + let processor = TreeProcessor::new(db_path.clone(), snapshot.clone()).await?; let (tx, rx) = mpsc::channel::(5); let processor_handle = tokio::spawn(async move { @@ -87,6 +87,11 @@ async fn main() -> Result<()> { fetcher.run(tx).await?; processor_handle.await?; + + // Write the current state to a file. + let snapshot = snapshot.lock().await; + let state_file_path = db_path.join(STATE_FILE_NAME); + snapshot.write(&state_file_path)?; } ReconstructSource::File { file } => { let snapshot = Arc::new(Mutex::new(StateSnapshot::default())); diff --git a/src/processor/tree/mod.rs b/src/processor/tree/mod.rs index 31be019..4128af8 100644 --- a/src/processor/tree/mod.rs +++ b/src/processor/tree/mod.rs @@ -1,7 +1,7 @@ pub mod query_tree; mod tree_wrapper; -use std::{io, path::PathBuf, sync::Arc}; +use std::{path::PathBuf, sync::Arc}; use async_trait::async_trait; use ethers::types::H256; @@ -17,8 +17,6 @@ use super::Processor; pub type RootHash = H256; pub struct TreeProcessor { - /// The path to the directory in which database files and state snapshots will be written. - db_path: PathBuf, /// The internal merkle tree. tree: TreeWrapper, /// The stored state snapshot. @@ -44,45 +42,29 @@ impl TreeProcessor { let index_to_key_map = snapshot.lock().await.index_to_key_map.clone(); let tree = TreeWrapper::new(&db_path, index_to_key_map)?; - Ok(Self { - db_path, - tree, - snapshot, - }) - } - - pub async fn write_state(&self) -> Result<(), io::Error> { - let snapshot = self.snapshot.lock().await; - // Write the current state to a file. - let state_file_path = self.db_path.join(STATE_FILE_NAME); - snapshot.write(&state_file_path) + Ok(Self { tree, snapshot }) } } #[async_trait] impl Processor for TreeProcessor { async fn run(mut self, mut rx: mpsc::Receiver) { - loop { - if let Some(block) = rx.recv().await { - let mut snapshot = self.snapshot.lock().await; - // Check if we've already processed this block. - if snapshot.latest_l2_block_number >= block.block_number { - tracing::debug!( - "Block {} has already been processed, skipping.", - block.block_number - ); - continue; - } + while let Some(block) = rx.recv().await { + let mut snapshot = self.snapshot.lock().await; + // Check if we've already processed this block. + if snapshot.latest_l2_block_number >= block.block_number { + tracing::debug!( + "Block {} has already been processed, skipping.", + block.block_number + ); + continue; + } - self.tree.insert_block(&block); + self.tree.insert_block(&block); - // Update snapshot values. - snapshot.latest_l2_block_number = block.block_number; - snapshot.index_to_key_map = self.tree.index_to_key_map.clone(); - } else { - self.write_state().await.unwrap(); - break; - } + // Update snapshot values. + snapshot.latest_l2_block_number = block.block_number; + snapshot.index_to_key_map = self.tree.index_to_key_map.clone(); } } } diff --git a/state-reconstruct-fetcher/Cargo.toml b/state-reconstruct-fetcher/Cargo.toml index 14bc0c1..b9da78a 100644 --- a/state-reconstruct-fetcher/Cargo.toml +++ b/state-reconstruct-fetcher/Cargo.toml @@ -13,4 +13,5 @@ serde_json = { version = "1.0.107", features = ["std"] } serde_json_any_key = "2.0.0" thiserror = "1.0.50" tokio = { version = "1.33.0", features = ["signal"] } +tokio-util = "0.7.10" tracing = "0.1.40" diff --git a/state-reconstruct-fetcher/src/l1_fetcher.rs b/state-reconstruct-fetcher/src/l1_fetcher.rs index 7ca3293..a07b21c 100644 --- a/state-reconstruct-fetcher/src/l1_fetcher.rs +++ b/state-reconstruct-fetcher/src/l1_fetcher.rs @@ -9,9 +9,10 @@ use eyre::Result; use rand::random; use thiserror::Error; use tokio::{ - sync::{mpsc, oneshot, Mutex}, + sync::{mpsc, Mutex}, time::{sleep, Duration}, }; +use tokio_util::sync::CancellationToken; use crate::{ constants::ethereum::{BLOCK_STEP, GENESIS_BLOCK, ZK_SYNC_ADDR}, @@ -174,11 +175,12 @@ impl L1Fetcher { }); // Wait for shutdown signal in background. - let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let token = CancellationToken::new(); + let cloned_token = token.clone(); tokio::spawn(async move { let _ = tokio::signal::ctrl_c().await; tracing::info!("Shutdown signal received, finishing up and shutting down..."); - let _ = shutdown_tx.send(""); + cloned_token.cancel(); }); let (hash_tx, hash_rx) = mpsc::channel(5); @@ -194,21 +196,33 @@ impl L1Fetcher { // - BlockCommit event filter (main). // - Referred L1 block fetch (tx). // - Calldata parsing (parse). - let tx_handle = - self.spawn_tx_handler(hash_rx, calldata_tx, current_l1_block_number.as_u64()); + let tx_handle = self.spawn_tx_handler( + hash_rx, + calldata_tx, + token.clone(), + current_l1_block_number.as_u64(), + ); let parse_handle = self.spawn_parsing_handler(calldata_rx, sink)?; let main_handle = self.spawn_main_handler( hash_tx, - shutdown_rx, + token, current_l1_block_number, end_block_number, disable_polling, )?; tx_handle.await?; - parse_handle.await?; + let last_processed_l1_block_num = parse_handle.await?; main_handle.await?; + // Store our current L1 block number so we can resume from where we left + // off, we also make sure to update the metrics before printing them. + if let Some(block_num) = last_processed_l1_block_num { + self.metrics.lock().await.latest_l1_block_nbr = block_num; + if let Some(snapshot) = &self.snapshot { + snapshot.lock().await.latest_l1_block_number = U64::from(block_num); + } + } self.metrics.lock().await.print(); Ok(()) @@ -217,7 +231,7 @@ impl L1Fetcher { fn spawn_main_handler( &self, hash_tx: mpsc::Sender, - mut shutdown_rx: oneshot::Receiver<&'static str>, + cancellation_token: CancellationToken, mut current_l1_block_number: U64, end_block_number: U64, disable_polling: bool, @@ -225,7 +239,6 @@ impl L1Fetcher { let metrics = self.metrics.clone(); let event = self.contract.events_by_name("BlockCommit")?[0].clone(); let provider_clone = self.provider.clone(); - let snapshot_clone = self.snapshot.clone(); Ok(tokio::spawn({ async move { @@ -234,14 +247,9 @@ impl L1Fetcher { loop { // Break when reaching the `end_block` or on the receivement of a `ctrl_c` signal. if (disable_polling && current_l1_block_number > end_block_number) - || shutdown_rx.try_recv().is_ok() + || cancellation_token.is_cancelled() { - // Store our current L1 block number so we can resume from where we left - // off, we also make sure to update the metrics before leaving the loop. - metrics.lock().await.latest_l1_block_nbr = current_l1_block_number.as_u64(); - if let Some(snapshot) = &snapshot_clone { - snapshot.lock().await.latest_l1_block_number = current_l1_block_number; - } + tracing::debug!("Shutting down main handle..."); break; } @@ -296,6 +304,7 @@ impl L1Fetcher { &self, mut hash_rx: mpsc::Receiver, l1_tx_tx: mpsc::Sender, + cancellation_token: CancellationToken, mut last_block: u64, ) -> tokio::task::JoinHandle<()> { let metrics = self.metrics.clone(); @@ -315,6 +324,12 @@ impl L1Fetcher { break tx; } _ => { + // Task has been cancelled by user, abort loop. + if cancellation_token.is_cancelled() { + tracing::debug!("Shutting down tx handle..."); + return; + } + tracing::error!( "failed to get transaction for hash: {}, retrying in a bit...", hash @@ -346,12 +361,14 @@ impl L1Fetcher { &self, mut l1_tx_rx: mpsc::Receiver, sink: mpsc::Sender, - ) -> Result> { + ) -> Result>> { let metrics = self.metrics.clone(); let function = self.contract.functions_by_name("commitBlocks")?[0].clone(); Ok(tokio::spawn({ async move { + let mut last_block_number_processed = None; + while let Some(tx) = l1_tx_rx.recv().await { let block_number = tx.block_number.map(|v| v.as_u64()); let blocks = match parse_calldata(block_number, &function, &tx.input) { @@ -369,7 +386,13 @@ impl L1Fetcher { metrics.latest_l2_block_nbr = blk.block_number; sink.send(blk).await.unwrap(); } + + last_block_number_processed = block_number; } + + // Return the last processed l1 block number, + // so we can resume from the same point later on. + last_block_number_processed } })) }