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.

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::<CommitBlockInfoV1>(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()));

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<CommitBlockInfoV1>) {
-        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();
         }
     }
 }

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"

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,15 +231,14 @@ impl L1Fetcher {
     fn spawn_main_handler(
         &self,
         hash_tx: mpsc::Sender<H256>,
-        mut shutdown_rx: oneshot::Receiver<&'static str>,
+        cancellation_token: CancellationToken,
         mut current_l1_block_number: U64,
         end_block_number: U64,
         disable_polling: bool,
     ) -> Result<tokio::task::JoinHandle<()>> {
         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<H256>,
         l1_tx_tx: mpsc::Sender<Transaction>,
+        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<Transaction>,
         sink: mpsc::Sender<CommitBlockInfoV1>,
-    ) -> Result<tokio::task::JoinHandle<()>> {
+    ) -> Result<tokio::task::JoinHandle<Option<u64>>> {
         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
             }
         }))
     }