perf(l1): pipeline bytecode downloads and background storage healing in snap sync

crates/networking/p2p/sync.rs

@@ -233,6 +233,8 @@ pub enum SyncError {
     RocksDBError(String),
     #[error("Bytecode file error")]
     BytecodeFileError,
+    #[error("Concurrent bytecode download task panicked: {0}")]
+    BytecodeTaskPanicked(String),
     #[error("Error in Peer Table: {0}")]
     PeerTableError(#[from] PeerTableError),
     #[error("Missing fullsync batch")]
@@ -261,6 +263,7 @@ impl SyncError {
             | SyncError::StorageTempDBDirNotFound(_)
             | SyncError::RocksDBError(_)
             | SyncError::BytecodeFileError
+            | SyncError::BytecodeTaskPanicked(_)
             | SyncError::NoLatestCanonical
             | SyncError::PeerTableError(_)
             | SyncError::MissingFullsyncBatch

crates/networking/p2p/sync/code_collector.rs

@@ -6,6 +6,7 @@ use ethrex_common::H256;
 use ethrex_rlp::encode::RLPEncode;
 use std::collections::HashSet;
 use std::path::PathBuf;
+use tokio::sync::mpsc;
 use tokio::task::JoinSet;
 use tracing::error;
 
@@ -19,6 +20,8 @@ pub struct CodeHashCollector {
     file_index: u64,
     // JoinSet to manage async disk writes
     disk_tasks: JoinSet<Result<(), DumpError>>,
+    // Optional channel to stream code hashes for concurrent bytecode downloading
+    code_hash_sender: Option<mpsc::UnboundedSender<Vec<H256>>>,
 }
 
 impl CodeHashCollector {
@@ -29,19 +32,38 @@ impl CodeHashCollector {
             snapshots_dir,
             file_index: 0,
             disk_tasks: JoinSet::new(),
+            code_hash_sender: None,
         }
     }
 
-    /// Adds a code hash to the buffer
+    /// Creates the channel and returns the receiver for concurrent bytecode downloading.
+    /// Must be called before any code hashes are added.
+    pub fn take_receiver(&mut self) -> mpsc::UnboundedReceiver<Vec<H256>> {
+        let (tx, rx) = mpsc::unbounded_channel();
+        self.code_hash_sender = Some(tx);
+        rx
+    }
+
+    /// Adds a code hash to the buffer and sends it through the channel if present
     pub fn add(&mut self, hash: H256) {
-        self.buffer.insert(hash);
+        if self.buffer.insert(hash) && let Some(sender) = &self.code_hash_sender {
+            // Ignore send errors — the receiver may have been dropped if
+            // the bytecode task finished early or errored out
+            let _ = sender.send(vec![hash]);
+        }
     }
 
     // The optimization for rocksdb database doesn't use this method
     #[cfg(not(feature = "rocksdb"))]
     /// Extends the buffer with a list of code hashes
     pub fn extend(&mut self, hashes: impl IntoIterator<Item = H256>) {
-        self.buffer.extend(hashes);
+        let new_hashes: Vec<H256> = hashes
+            .into_iter()
+            .filter(|h| self.buffer.insert(*h))
+            .collect();
+        if !new_hashes.is_empty() && let Some(sender) = &self.code_hash_sender {
+            let _ = sender.send(new_hashes);
+        }
     }
 
     /// Flushes the buffer to a file if the buffer is larger than [`CODE_HASH_WRITE_BUFFER_SIZE`]
@@ -55,14 +77,17 @@ impl CodeHashCollector {
         Ok(())
     }
 
-    /// Finishes the code collector and returns the final index of file
+    /// Finishes the code collector: flushes remaining hashes and closes the channel
     pub async fn finish(mut self) -> Result<(), SyncError> {
         // Final flush if needed
         if !self.buffer.is_empty() {
             let buffer = std::mem::take(&mut self.buffer);
             self.flush_buffer(buffer);
         }
 
+        // Drop the sender to close the channel, signaling the bytecode task to drain
+        self.code_hash_sender.take();
+
         // Wait for all pending writes
         self.disk_tasks
             .join_all()

crates/networking/p2p/sync/healing/storage.rs

@@ -116,6 +116,10 @@ pub struct NodeRequest {
 /// - When a node is downloaded:
 ///    - if it has no missing children, we store it in the db
 ///    - if the node has missing children, we store it in our healing_queue, which is preserved between calls
+///
+/// When `time_limit` is `Some(duration)`, healing will stop early after the duration elapses
+/// and return `Ok(false)`, allowing the caller to proceed with partial healing.
+/// `global_roots_healed` accumulates the number of account storage roots fully healed across calls.
 pub async fn heal_storage_trie(
     state_root: H256,
     storage_accounts: &AccountStorageRoots,
@@ -124,13 +128,17 @@ pub async fn heal_storage_trie(
     healing_queue: StorageHealingQueue,
     staleness_timestamp: u64,
     global_leafs_healed: &mut u64,
+    time_limit: Option<Duration>,
+    global_roots_healed: &mut u64,
 ) -> Result<bool, SyncError> {
     METRICS.current_step.set(CurrentStepValue::HealingStorage);
     let download_queue = get_initial_downloads(&store, state_root, storage_accounts);
     debug!(
         initial_accounts_count = download_queue.len(),
+        ?time_limit,
         "Started Storage Healing",
     );
+    let healing_start_time = Instant::now();
     let mut state = StorageHealer {
         last_update: Instant::now(),
         download_queue,
@@ -203,8 +211,14 @@ pub async fn heal_storage_trie(
 
         let is_done = state.requests.is_empty() && state.download_queue.is_empty();
         let is_stale = current_unix_time() > state.staleness_timestamp;
-
-        if nodes_to_write.values().map(Vec::len).sum::<usize>() > 100_000 || is_done || is_stale {
+        let time_limit_reached =
+            time_limit.is_some_and(|limit| healing_start_time.elapsed() >= limit);
+
+        if nodes_to_write.values().map(Vec::len).sum::<usize>() > 100_000
+            || is_done
+            || is_stale
+            || time_limit_reached
+        {
             let to_write: Vec<_> = nodes_to_write.drain().collect();
             let store = state.store.clone();
             // NOTE: we keep only a single task in the background to avoid out of order deletes
@@ -230,13 +244,24 @@ pub async fn heal_storage_trie(
         }
 
         if is_done {
+            *global_roots_healed += state.roots_healed as u64;
             db_joinset.join_all().await;
             return Ok(true);
         }
 
-        if is_stale {
+        if is_stale || time_limit_reached {
+            *global_roots_healed += state.roots_healed as u64;
             db_joinset.join_all().await;
-            state.healing_queue = HashMap::new();
+            if is_stale {
+                state.healing_queue = HashMap::new();
+            }
+            if time_limit_reached {
+                debug!(
+                    roots_healed = state.roots_healed,
+                    elapsed_ms = healing_start_time.elapsed().as_millis() as u64,
+                    "Storage healing stopped: time limit reached",
+                );
+            }
             return Ok(false);
         }