perf(l1): optimize snap sync insertion and healing write paths

crates/common/trie/trie_sorted.rs

@@ -149,8 +149,15 @@ fn flush_nodes_to_write(
     db: &dyn TrieDB,
     sender: Sender<Vec<(Nibbles, Node)>>,
 ) -> Result<(), TrieGenerationError> {
+    let start = std::time::Instant::now();
+    let node_count = nodes_to_write.len();
     db.put_batch_no_alloc(&nodes_to_write)
         .map_err(TrieGenerationError::FlushToDbError)?;
+    tracing::debug!(
+        node_count,
+        elapsed_ms = start.elapsed().as_millis() as u64,
+        "flush_nodes_to_write"
+    );
     nodes_to_write.clear();
     let _ = sender.send(nodes_to_write);
     Ok(())
@@ -351,7 +358,10 @@ where
         let _ = buffer_sender.send(Vec::with_capacity(SIZE_TO_WRITE_DB as usize));
     }
     scope(|s| {
-        let pool = ThreadPool::new(12, s);
+        let thread_count = std::thread::available_parallelism()
+            .map(|n| n.get())
+            .unwrap_or(8);
+        let pool = ThreadPool::new(thread_count, s);
         trie_from_sorted_accounts(
             db,
             accounts_iter,

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

@@ -283,28 +283,37 @@ async fn heal_state_trie(
             // PERF: reuse buffers?
             let to_write = std::mem::take(&mut nodes_to_write);
             let store = store.clone();
-            // NOTE: we keep only a single task in the background to avoid out of order deletes
-            if !db_joinset.is_empty() {
+            // Allow up to 2 tasks in flight to overlap encoding with DB commit
+            while db_joinset.len() >= 2 {
                 db_joinset
                     .join_next()
                     .await
-                    .expect("we just checked joinset is not empty")?;
+                    .expect("joinset is not empty")?;
             }
             db_joinset.spawn_blocking(move || {
+                let encode_start = std::time::Instant::now();
+                let node_count = to_write.len();
                 let mut encoded_to_write = BTreeMap::new();
                 for (path, node) in to_write {
                     for i in 0..path.len() {
-                        encoded_to_write.insert(path.slice(0, i), vec![]);
+                        encoded_to_write.entry(path.slice(0, i)).or_insert(vec![]);
                     }
                     encoded_to_write.insert(path, node.encode_to_vec());
                 }
+                let encode_ms = encode_start.elapsed().as_millis() as u64;
+                let db_start = std::time::Instant::now();
                 let trie_db = store
                     .open_direct_state_trie(*EMPTY_TRIE_HASH)
                     .expect("Store should open");
                 let db = trie_db.db();
-                // PERF: use put_batch_no_alloc (note that it needs to remove nodes too)
                 db.put_batch(encoded_to_write.into_iter().collect())
                     .expect("The put batch on the store failed");
+                debug!(
+                    node_count,
+                    encode_ms,
+                    db_write_ms = db_start.elapsed().as_millis() as u64,
+                    "state healing batch write"
+                );
             });
         }
 

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

@@ -212,25 +212,34 @@ pub async fn heal_storage_trie(
         if nodes_to_write.values().map(Vec::len).sum::<usize>() > 100_000 || is_done || is_stale {
             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
-            if !db_joinset.is_empty() {
+            // Allow up to 2 tasks in flight to overlap encoding with DB commit
+            while db_joinset.len() >= 2 {
                 db_joinset.join_next().await;
             }
             db_joinset.spawn_blocking(move || {
+                let encode_start = std::time::Instant::now();
+                let account_count = to_write.len();
                 let mut encoded_to_write = vec![];
                 for (hashed_account, nodes) in to_write {
-                    let mut account_nodes = vec![];
+                    let mut account_nodes = std::collections::BTreeMap::new();
                     for (path, node) in nodes {
                         for i in 0..path.len() {
-                            account_nodes.push((path.slice(0, i), vec![]));
+                            account_nodes.entry(path.slice(0, i)).or_insert(vec![]);
                         }
-                        account_nodes.push((path, node.encode_to_vec()));
+                        account_nodes.insert(path, node.encode_to_vec());
                     }
-                    encoded_to_write.push((hashed_account, account_nodes));
+                    encoded_to_write.push((hashed_account, account_nodes.into_iter().collect()));
                 }
-                // PERF: use put_batch_no_alloc? (it needs to remove parent nodes too)
-                spawned_rt::tasks::block_on(store.write_storage_trie_nodes_batch(encoded_to_write))
+                let encode_ms = encode_start.elapsed().as_millis() as u64;
+                let db_start = std::time::Instant::now();
+                store.write_storage_trie_nodes_batch_sync(encoded_to_write)
                     .expect("db write failed");
+                debug!(
+                    account_count,
+                    encode_ms,
+                    db_write_ms = db_start.elapsed().as_millis() as u64,
+                    "storage healing batch write"
+                );
             });
         }
 

crates/networking/p2p/sync/snap_sync.rs

@@ -956,16 +956,9 @@ async fn insert_accounts(
         .collect();
     db.ingest_external_file(file_paths)
         .map_err(|err| SyncError::RocksDBError(err.into_string()))?;
-    let iter = db.full_iterator(rocksdb::IteratorMode::Start);
-    for account in iter {
-        let account = account.map_err(|err| SyncError::RocksDBError(err.into_string()))?;
-        let account_state = AccountState::decode(&account.1).map_err(SyncError::Rlp)?;
-        if account_state.code_hash != *EMPTY_KECCACK_HASH {
-            code_hash_collector.add(account_state.code_hash);
-            code_hash_collector.flush_if_needed().await?;
-        }
-    }
 
+    let start = std::time::Instant::now();
+    let mut code_hashes: Vec<H256> = Vec::new();
     let iter = db.full_iterator(rocksdb::IteratorMode::Start);
     let compute_state_root = trie_from_sorted_accounts_wrap(
         trie.db(),
@@ -976,6 +969,9 @@ async fn insert_accounts(
                     .account_tries_inserted
                     .fetch_add(1, Ordering::Relaxed);
                 let account_state = AccountState::decode(v).expect("We should have accounts here");
+                if account_state.code_hash != *EMPTY_KECCACK_HASH {
+                    code_hashes.push(account_state.code_hash);
+                }
                 if account_state.storage_root != *EMPTY_TRIE_HASH {
                     storage_accounts.accounts_with_storage_root.insert(
                         H256::from_slice(k),
@@ -986,6 +982,15 @@ async fn insert_accounts(
             .map(|(k, v)| (H256::from_slice(&k), v.to_vec())),
     )
     .map_err(SyncError::TrieGenerationError)?;
+    debug!(
+        elapsed_ms = start.elapsed().as_millis() as u64,
+        code_hashes_count = code_hashes.len(),
+        "insert_accounts trie build"
+    );
+    for hash in code_hashes {
+        code_hash_collector.add(hash);
+        code_hash_collector.flush_if_needed().await?;
+    }
 
     drop(db); // close db before removing directory
 

crates/storage/store.rs

@@ -1139,22 +1139,39 @@ impl Store {
     ) -> Result<(), StoreError> {
         let mut txn = self.backend.begin_write()?;
         tokio::task::spawn_blocking(move || {
-            for (address_hash, nodes) in storage_trie_nodes {
-                for (node_path, node_data) in nodes {
-                    let key = apply_prefix(Some(address_hash), node_path);
-                    if node_data.is_empty() {
-                        txn.delete(STORAGE_TRIE_NODES, key.as_ref())?;
-                    } else {
-                        txn.put(STORAGE_TRIE_NODES, key.as_ref(), &node_data)?;
-                    }
-                }
-            }
-            txn.commit()
+            Self::write_storage_trie_nodes_inner(&mut *txn, storage_trie_nodes)
         })
         .await
         .map_err(|e| StoreError::Custom(format!("Task panicked: {}", e)))?
     }
 
+    /// Synchronous version of [`Self::write_storage_trie_nodes_batch`] for use inside
+    /// `spawn_blocking` closures, avoiding a double thread hop through `block_on`.
+    pub fn write_storage_trie_nodes_batch_sync(
+        &self,
+        storage_trie_nodes: StorageUpdates,
+    ) -> Result<(), StoreError> {
+        let mut txn = self.backend.begin_write()?;
+        Self::write_storage_trie_nodes_inner(&mut *txn, storage_trie_nodes)
+    }
+
+    fn write_storage_trie_nodes_inner(
+        txn: &mut dyn crate::api::StorageWriteBatch,
+        storage_trie_nodes: StorageUpdates,
+    ) -> Result<(), StoreError> {
+        for (address_hash, nodes) in storage_trie_nodes {
+            for (node_path, node_data) in nodes {
+                let key = apply_prefix(Some(address_hash), node_path);
+                if node_data.is_empty() {
+                    txn.delete(STORAGE_TRIE_NODES, key.as_ref())?;
+                } else {
+                    txn.put(STORAGE_TRIE_NODES, key.as_ref(), &node_data)?;
+                }
+            }
+        }
+        txn.commit()
+    }
+
     /// CAUTION: This method writes directly to the underlying database, bypassing any caching layer.
     /// For updating the state after block execution, use [`Self::store_block_updates`].
     pub async fn write_account_code_batch(

crates/storage/trie.rs

@@ -5,8 +5,10 @@ use crate::api::{StorageBackend, StorageLockedView};
 use crate::error::StoreError;
 use crate::layering::apply_prefix;
 use ethrex_common::H256;
-use ethrex_trie::{Nibbles, TrieDB, error::TrieError};
+use ethrex_rlp::encode::RLPEncode;
+use ethrex_trie::{Nibbles, Node, TrieDB, error::TrieError};
 use std::sync::Arc;
+use std::time::Instant;
 
 /// StorageWriteBatch implementation for the TrieDB trait
 /// Wraps a transaction to allow multiple trie operations on the same transaction
@@ -101,17 +103,52 @@ impl TrieDB for BackendTrieDB {
     }
 
     fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError> {
+        let start = Instant::now();
+        let item_count = key_values.len();
         let mut tx = self.db.begin_write().map_err(|e| {
             TrieError::DbError(anyhow::anyhow!("Failed to begin write transaction: {}", e))
         })?;
         for (key, value) in key_values {
             let prefixed_key = self.make_key(key);
             let table = self.table_for_key(&prefixed_key);
-            tx.put_batch(table, vec![(prefixed_key, value)])
+            tx.put(table, &prefixed_key, &value)
                 .map_err(|e| TrieError::DbError(anyhow::anyhow!("Failed to write batch: {}", e)))?;
         }
+        let commit_start = Instant::now();
         tx.commit()
-            .map_err(|e| TrieError::DbError(anyhow::anyhow!("Failed to write batch: {}", e)))
+            .map_err(|e| TrieError::DbError(anyhow::anyhow!("Failed to write batch: {}", e)))?;
+        tracing::debug!(
+            items = item_count,
+            total_ms = start.elapsed().as_millis() as u64,
+            commit_ms = commit_start.elapsed().as_millis() as u64,
+            "BackendTrieDB::put_batch"
+        );
+        Ok(())
+    }
+
+    fn put_batch_no_alloc(&self, key_values: &[(Nibbles, Node)]) -> Result<(), TrieError> {
+        let start = Instant::now();
+        let item_count = key_values.len();
+        let mut tx = self.db.begin_write().map_err(|e| {
+            TrieError::DbError(anyhow::anyhow!("Failed to begin write transaction: {}", e))
+        })?;
+        for (path, node) in key_values {
+            let prefixed_key = self.make_key(path.clone());
+            let table = self.table_for_key(&prefixed_key);
+            let encoded = node.encode_to_vec();
+            tx.put(table, &prefixed_key, &encoded)
+                .map_err(|e| TrieError::DbError(anyhow::anyhow!("Failed to write batch: {}", e)))?;
+        }
+        let commit_start = Instant::now();
+        tx.commit()
+            .map_err(|e| TrieError::DbError(anyhow::anyhow!("Failed to write batch: {}", e)))?;
+        tracing::debug!(
+            items = item_count,
+            total_ms = start.elapsed().as_millis() as u64,
+            commit_ms = commit_start.elapsed().as_millis() as u64,
+            "BackendTrieDB::put_batch_no_alloc"
+        );
+        Ok(())
     }
 }