Make ETL file size configurable

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

@@ -82,6 +82,10 @@ pub struct Command {
     #[arg(long)]
     batch_size: Option<u64>,
 
+    /// Size for temporary file during ETL stages
+    #[arg(long)]
+    etl_file_size: Option<usize>,
+
     /// Normally, running the stage requires unwinding for stages that already
     /// have been run, in order to not rewrite to the same database slots.
     ///
@@ -151,6 +155,8 @@ impl Command {
 
         let batch_size = self.batch_size.unwrap_or(self.to - self.from + 1);
 
+        let etl_file_size = self.etl_file_size.unwrap_or(500 * 1024 * 1024);
+
         let (mut exec_stage, mut unwind_stage): (Box<dyn Stage<_>>, Option<Box<dyn Stage<_>>>) =
             match self.stage {
                 StageEnum::Bodies => {
@@ -229,7 +235,7 @@ impl Command {
                     )
                 }
                 StageEnum::TxLookup => {
-                    (Box::new(TransactionLookupStage::new(batch_size, None)), None)
+                    (Box::new(TransactionLookupStage::new(batch_size, etl_file_size, None)), None)
                 }
                 StageEnum::AccountHashing => {
                     (Box::new(AccountHashingStage::new(1, batch_size)), None)

book/cli/reth/stage/run.md

@@ -60,6 +60,9 @@ Options:
 
       --batch-size <BATCH_SIZE>
           Batch size for stage execution and unwind
+
+      --etl-file-size <ETL_FILE_SIZE>
+          Batch size for stage execution and unwind
 
   -s, --skip-unwind
           Normally, running the stage requires unwinding for stages that already have been run, in order to not rewrite to the same database slots.

book/run/config.md

@@ -193,6 +193,12 @@ The transaction lookup stage builds an index of transaction hashes to their sequ
 # Lower thresholds correspond to more frequent disk I/O (writes),
 # but lowers memory usage
 chunk_size = 5000000
+
+# The size of temporary file for ETL data collector.
+#
+# Lower threshold corresponds to more frequent flushes,
+# but lowers temporary storage usage
+etl_file-size = 524288000
 ```
 
 ### `index_account_history`

crates/config/src/config.rs

@@ -227,11 +227,13 @@ impl Default for MerkleConfig {
 pub struct TransactionLookupConfig {
     /// The maximum number of transactions to process before writing to disk.
     pub chunk_size: u64,
+    /// Size of temporary file during etl stage.
+    pub etl_file_size: usize,
 }
 
 impl Default for TransactionLookupConfig {
     fn default() -> Self {
-        Self { chunk_size: 5_000_000 }
+        Self { chunk_size: 5_000_000, etl_file_size: 500 * 1024 * 1024 }
     }
 }
 

crates/node-core/src/node_config.rs

@@ -866,6 +866,7 @@ impl NodeConfig {
                 .set(MerkleStage::new_execution(stage_config.merkle.clean_threshold))
                 .set(TransactionLookupStage::new(
                     stage_config.transaction_lookup.chunk_size,
+                    stage_config.transaction_lookup.etl_file_size,
                     prune_modes.transaction_lookup,
                 ))
                 .set(IndexAccountHistoryStage::new(

crates/stages/benches/criterion.rs

@@ -57,7 +57,7 @@ fn transaction_lookup(c: &mut Criterion) {
     let mut group = c.benchmark_group("Stages");
     // don't need to run each stage for that many times
     group.sample_size(10);
-    let stage = TransactionLookupStage::new(DEFAULT_NUM_BLOCKS, None);
+    let stage = TransactionLookupStage::new(DEFAULT_NUM_BLOCKS, 500 * 1024 * 1024, None);
 
     let db = setup::txs_testdata(DEFAULT_NUM_BLOCKS);
 

crates/stages/src/stages/tx_lookup.rs

@@ -34,19 +34,20 @@ pub struct TransactionLookupStage {
     /// The maximum number of lookup entries to hold in memory before pushing them to
     /// [`reth_etl::Collector`].
     chunk_size: u64,
+    etl_file_size: usize,
     prune_mode: Option<PruneMode>,
 }
 
 impl Default for TransactionLookupStage {
     fn default() -> Self {
-        Self { chunk_size: 5_000_000, prune_mode: None }
+        Self { chunk_size: 5_000_000, etl_file_size: 500 * 1024 * 1024, prune_mode: None }
     }
 }
 
 impl TransactionLookupStage {
     /// Create new instance of [TransactionLookupStage].
-    pub fn new(chunk_size: u64, prune_mode: Option<PruneMode>) -> Self {
-        Self { chunk_size, prune_mode }
+    pub fn new(chunk_size: u64, etl_file_size: usize, prune_mode: Option<PruneMode>) -> Self {
+        Self { chunk_size, etl_file_size, prune_mode }
     }
 }
 
@@ -102,7 +103,7 @@ impl<DB: Database> Stage<DB> for TransactionLookupStage {
 
         // 500MB temporary files
         let mut hash_collector: Collector<TxHash, TxNumber> =
-            Collector::new(Arc::new(TempDir::new()?), 500 * (1024 * 1024));
+            Collector::new(Arc::new(TempDir::new()?), self.etl_file_size);
 
         debug!(
             target: "sync::stages::transaction_lookup",
@@ -397,12 +398,13 @@ mod tests {
     struct TransactionLookupTestRunner {
         db: TestStageDB,
         chunk_size: u64,
+        etl_file_size: usize,
         prune_mode: Option<PruneMode>,
     }
 
     impl Default for TransactionLookupTestRunner {
         fn default() -> Self {
-            Self { db: TestStageDB::default(), chunk_size: 1000, prune_mode: None }
+            Self { db: TestStageDB::default(), chunk_size: 1000, etl_file_size: 500 * 1024 * 1024, prune_mode: None }
         }
     }
 
@@ -449,7 +451,7 @@ mod tests {
         }
 
         fn stage(&self) -> Self::S {
-            TransactionLookupStage { chunk_size: self.chunk_size, prune_mode: self.prune_mode }
+            TransactionLookupStage { chunk_size: self.chunk_size, etl_file_size: self.etl_file_size, prune_mode: self.prune_mode }
         }
     }