feat: add ETL to TxLookup stage

bin/reth/src/commands/stage/run.rs

@@ -28,7 +28,7 @@ use reth_stages::{
     },
     ExecInput, Stage, StageExt, UnwindInput,
 };
-use std::{any::Any, net::SocketAddr, path::PathBuf, sync::Arc};
+use std::{any::Any, net::SocketAddr, path::PathBuf, sync::Arc, time::Instant};
 use tracing::*;
 
 /// `reth stage` command
@@ -269,6 +269,8 @@ impl Command {
             checkpoint: Some(checkpoint.with_block_number(self.from)),
         };
 
+        let start = Instant::now();
+        info!(target: "reth::cli", stage= ?self.stage, "Executing stage.", );
         loop {
             exec_stage.execute_ready(input).await?;
             let output = exec_stage.execute(&provider_rw, input)?;
@@ -284,6 +286,7 @@ impl Command {
                 break
             }
         }
+        info!(target: "reth::cli", stage= ?self.stage, time = ?start.elapsed().as_secs(), "Finished stage.");
 
         Ok(())
     }

crates/etl/Cargo.toml

@@ -11,6 +11,7 @@ exclude.workspace = true
 [dependencies]
 tempfile.workspace = true
 reth-db.workspace = true
+rayon.workspace = true
 
 [dev-dependencies]
 reth-primitives.workspace = true

crates/etl/src/lib.rs

@@ -23,6 +23,7 @@ use std::{
     sync::Arc,
 };
 
+use rayon::prelude::*;
 use reth_db::table::{Compress, Encode, Key, Value};
 use tempfile::{NamedTempFile, TempDir};
 
@@ -99,7 +100,7 @@ where
 
     fn flush(&mut self) {
         self.buffer_size_bytes = 0;
-        self.buffer.sort_unstable_by(|a, b| a.0.cmp(&b.0));
+        self.buffer.par_sort_unstable_by(|a, b| a.0.cmp(&b.0));
         let mut buf = Vec::with_capacity(self.buffer.len());
         std::mem::swap(&mut buf, &mut self.buffer);
         self.files.push(EtlFile::new(self.dir.path(), buf).expect("could not flush data to disk"))

crates/node-core/src/init.rs

@@ -209,7 +209,7 @@ pub fn insert_genesis_history<DB: Database>(
 /// Inserts header for the genesis state.
 pub fn insert_genesis_header<DB: Database>(
     tx: &<DB as Database>::TXMut,
-    snapshot_provider: Arc<SnapshotProvider>,
+    snapshot_provider: SnapshotProvider,
     chain: Arc<ChainSpec>,
 ) -> ProviderResult<()> {
     let (header, block_hash) = chain.sealed_genesis_header().split();

crates/snapshot/src/segments/headers.rs

@@ -9,7 +9,7 @@ use reth_provider::{
     providers::{SnapshotProvider, SnapshotWriter},
     DatabaseProviderRO,
 };
-use std::{ops::RangeInclusive, path::Path, sync::Arc};
+use std::{ops::RangeInclusive, path::Path};
 
 /// Snapshot segment responsible for [SnapshotSegment::Headers] part of data.
 #[derive(Debug, Default)]
@@ -23,7 +23,7 @@ impl<DB: Database> Segment<DB> for Headers {
     fn snapshot(
         &self,
         provider: DatabaseProviderRO<DB>,
-        snapshot_provider: Arc<SnapshotProvider>,
+        snapshot_provider: SnapshotProvider,
         block_range: RangeInclusive<BlockNumber>,
     ) -> ProviderResult<()> {
         let mut snapshot_writer =

crates/snapshot/src/segments/mod.rs

@@ -22,7 +22,7 @@ use reth_primitives::{
     BlockNumber, SnapshotSegment,
 };
 use reth_provider::{providers::SnapshotProvider, DatabaseProviderRO, TransactionsProviderExt};
-use std::{ops::RangeInclusive, path::Path, sync::Arc};
+use std::{ops::RangeInclusive, path::Path};
 
 pub(crate) type Rows<const COLUMNS: usize> = [Vec<Vec<u8>>; COLUMNS];
 
@@ -36,7 +36,7 @@ pub trait Segment<DB: Database>: Send + Sync {
     fn snapshot(
         &self,
         provider: DatabaseProviderRO<DB>,
-        snapshot_provider: Arc<SnapshotProvider>,
+        snapshot_provider: SnapshotProvider,
         block_range: RangeInclusive<BlockNumber>,
     ) -> ProviderResult<()>;
 

crates/snapshot/src/segments/receipts.rs

@@ -11,7 +11,7 @@ use reth_provider::{
     providers::{SnapshotProvider, SnapshotWriter},
     BlockReader, DatabaseProviderRO, TransactionsProviderExt,
 };
-use std::{ops::RangeInclusive, path::Path, sync::Arc};
+use std::{ops::RangeInclusive, path::Path};
 
 /// Snapshot segment responsible for [SnapshotSegment::Receipts] part of data.
 #[derive(Debug, Default)]
@@ -25,7 +25,7 @@ impl<DB: Database> Segment<DB> for Receipts {
     fn snapshot(
         &self,
         provider: DatabaseProviderRO<DB>,
-        snapshot_provider: Arc<SnapshotProvider>,
+        snapshot_provider: SnapshotProvider,
         block_range: RangeInclusive<BlockNumber>,
     ) -> ProviderResult<()> {
         let mut snapshot_writer =

crates/snapshot/src/segments/transactions.rs

@@ -11,7 +11,7 @@ use reth_provider::{
     providers::{SnapshotProvider, SnapshotWriter},
     BlockReader, DatabaseProviderRO, TransactionsProviderExt,
 };
-use std::{ops::RangeInclusive, path::Path, sync::Arc};
+use std::{ops::RangeInclusive, path::Path};
 
 /// Snapshot segment responsible for [SnapshotSegment::Transactions] part of data.
 #[derive(Debug, Default)]
@@ -27,7 +27,7 @@ impl<DB: Database> Segment<DB> for Transactions {
     fn snapshot(
         &self,
         provider: DatabaseProviderRO<DB>,
-        snapshot_provider: Arc<SnapshotProvider>,
+        snapshot_provider: SnapshotProvider,
         block_range: RangeInclusive<BlockNumber>,
     ) -> ProviderResult<()> {
         let mut snapshot_writer =

crates/snapshot/src/snapshotter.rs

@@ -10,7 +10,7 @@ use reth_provider::{
     ProviderFactory,
 };
 use reth_tokio_util::EventListeners;
-use std::{ops::RangeInclusive, sync::Arc, time::Instant};
+use std::{ops::RangeInclusive, time::Instant};
 use tokio_stream::wrappers::UnboundedReceiverStream;
 use tracing::{debug, trace};
 
@@ -26,7 +26,7 @@ pub struct Snapshotter<DB> {
     /// Provider factory
     provider_factory: ProviderFactory<DB>,
     /// Snapshot provider
-    snapshot_provider: Arc<SnapshotProvider>,
+    snapshot_provider: SnapshotProvider,
     /// Pruning configuration for every part of the data that can be pruned. Set by user, and
     /// needed in [Snapshotter] to prevent snapshotting the prunable data.
     /// See [Snapshotter::get_snapshot_targets].
@@ -71,7 +71,7 @@ impl<DB: Database> Snapshotter<DB> {
     /// Creates a new [Snapshotter].
     pub fn new(
         provider_factory: ProviderFactory<DB>,
-        snapshot_provider: Arc<SnapshotProvider>,
+        snapshot_provider: SnapshotProvider,
         prune_modes: PruneModes,
     ) -> Self {
         Self { provider_factory, snapshot_provider, prune_modes, listeners: Default::default() }

crates/stages/src/stages/headers.rs

@@ -99,7 +99,7 @@ where
     fn write_headers<DB: Database>(
         &mut self,
         tx: &<DB as Database>::TXMut,
-        snapshot_provider: Arc<SnapshotProvider>,
+        snapshot_provider: SnapshotProvider,
     ) -> Result<BlockNumber, StageError> {
         let total_headers = self.header_collector.len();
 

crates/stages/src/stages/sender_recovery.rs

@@ -1,23 +1,22 @@
 use crate::{BlockErrorKind, ExecInput, ExecOutput, Stage, StageError, UnwindInput, UnwindOutput};
-use itertools::Itertools;
 use reth_db::{
     cursor::DbCursorRW,
     database::Database,
+    snapshot::TransactionMask,
     tables,
     transaction::{DbTx, DbTxMut},
-    RawValue,
 };
-use reth_interfaces::{consensus, RethError};
+use reth_interfaces::consensus;
 use reth_primitives::{
     keccak256,
     stage::{EntitiesCheckpoint, StageCheckpoint, StageId},
-    Address, PruneSegment, TransactionSignedNoHash, TxNumber,
+    Address, PruneSegment, SnapshotSegment, TransactionSignedNoHash, TxNumber,
 };
 use reth_provider::{
     BlockReader, DatabaseProviderRW, HeaderProvider, ProviderError, PruneCheckpointReader,
-    StatsReader, TransactionsProvider,
+    StatsReader,
 };
-use std::{fmt::Debug, sync::mpsc};
+use std::{fmt::Debug, ops::Range, sync::mpsc};
 use thiserror::Error;
 use tracing::*;
 
@@ -83,43 +82,49 @@ impl<DB: Database> Stage<DB> for SenderRecoveryStage {
         // Acquire the cursor for inserting elements
         let mut senders_cursor = tx.cursor_write::<tables::TxSenders>()?;
 
-        // Query the transactions from both database and static files
-        let transactions = provider.raw_transactions_by_tx_range(tx_range.clone())?;
-
         // Iterate over transactions in chunks
         info!(target: "sync::stages::sender_recovery", ?tx_range, "Recovering senders");
 
-        // channels used to return result of sender recovery.
-        let mut channels = Vec::new();
-
         // Spawn recovery jobs onto the default rayon threadpool and send the result through the
         // channel.
         //
-        // We try to evenly divide the transactions to recover across all threads in the threadpool.
-        // Chunks are submitted instead of individual transactions to reduce the overhead of work
-        // stealing in the threadpool workers.
-        let chunk_size = self.commit_threshold as usize / rayon::current_num_threads();
-        // prevents an edge case
-        // where the chunk size is either 0 or too small
-        // to gain anything from using more than 1 thread
-        let chunk_size = chunk_size.max(16);
-
-        for chunk in &tx_range.zip(transactions).chunks(chunk_size) {
+        // Transactions are different size, so chunks will not all take the same processing time. If
+        // chunks are too big, there will be idle threads waiting for work. Choosing an
+        // arbitrary smaller value to make sure it doesn't happen.
+        let chunk_size = 100;
+
+        let chunks = (tx_range.start..tx_range.end)
+            .step_by(chunk_size as usize)
+            .map(|start| start..std::cmp::min(start + chunk_size as u64, tx_range.end))
+            .collect::<Vec<Range<u64>>>();
+
+        let mut channels = Vec::with_capacity(chunks.len());
+        for chunk_range in chunks {
             // An _unordered_ channel to receive results from a rayon job
             let (recovered_senders_tx, recovered_senders_rx) = mpsc::channel();
             channels.push(recovered_senders_rx);
-            // Note: Unfortunate side-effect of how chunk is designed in itertools (it is not Send)
-            let chunk: Vec<_> = chunk.collect();
 
-            // Spawn the sender recovery task onto the global rayon pool
-            // This task will send the results through the channel after it recovered the senders.
+            let manager = provider.snapshot_provider().clone();
+
+            // Spawn the task onto the global rayon pool
+            // This task will send the results through the channel after it has read the transaction
+            // and calculated the sender.
             rayon::spawn(move || {
                 let mut rlp_buf = Vec::with_capacity(128);
-                for entry in chunk {
-                    rlp_buf.clear();
-                    let recovery_result = recover_sender(entry, &mut rlp_buf);
-                    let _ = recovered_senders_tx.send(recovery_result);
-                }
+                let _ = manager.fetch_range_with_predicate(
+                    SnapshotSegment::Transactions,
+                    chunk_range,
+                    |cursor, number| {
+                        Ok(cursor
+                            .get_one::<TransactionMask<TransactionSignedNoHash>>(number.into())?
+                            .map(|tx| {
+                                rlp_buf.clear();
+                                let _ = recovered_senders_tx
+                                    .send(recover_sender((number, tx), &mut rlp_buf));
+                            }))
+                    },
+                    |_| true,
+                );
             });
         }
 
@@ -185,17 +190,11 @@ impl<DB: Database> Stage<DB> for SenderRecoveryStage {
     }
 }
 
+#[inline]
 fn recover_sender(
-    (tx_id, tx): (TxNumber, RawValue<TransactionSignedNoHash>),
+    (tx_id, tx): (TxNumber, TransactionSignedNoHash),
     rlp_buf: &mut Vec<u8>,
 ) -> Result<(u64, Address), Box<SenderRecoveryStageError>> {
-    let tx = tx
-        .value()
-        .map_err(RethError::from)
-        .map_err(StageError::from)
-        .map_err(Into::into)
-        .map_err(Box::new)?;
-
     tx.transaction.encode_without_signature(rlp_buf);
 
     // We call [Signature::recover_signer_unchecked] because transactions run in the pipeline are
@@ -330,7 +329,9 @@ mod tests {
             random_block_range(&mut rng, stage_progress + 1..=previous_stage, B256::ZERO, 0..4); // set tx count range high enough to hit the threshold
         runner.db.insert_blocks(seed.iter(), None).expect("failed to seed execution");
 
-        let total_transactions = runner.db.table::<tables::Transactions>().unwrap().len() as u64;
+        let total_transactions =
+            runner.db.factory.snapshot_provider().count_entries::<tables::Transactions>().unwrap()
+                as u64;
 
         let first_input = ExecInput {
             target: Some(previous_stage),