Make ETL file size configurable (#6927)

Co-authored-by: Alexey Shekhirin <a.shekhirin@gmail.com>
Co-authored-by: joshieDo <ranriver@protonmail.com>

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/lib.rs

@@ -59,6 +59,7 @@
 //!             headers_downloader,
 //!             bodies_downloader,
 //!             executor_factory,
+//!             500*1024*1024,
 //!         )
 //!     )
 //!     .build(provider_factory, static_file_producer);

crates/stages/src/sets.rs

@@ -100,6 +100,7 @@ impl<Provider, H, B, EF> DefaultStages<Provider, H, B, EF> {
         header_downloader: H,
         body_downloader: B,
         executor_factory: EF,
+        etl_file_size: usize,
     ) -> Self
     where
         EF: ExecutorFactory,
@@ -111,6 +112,7 @@ impl<Provider, H, B, EF> DefaultStages<Provider, H, B, EF> {
                 consensus,
                 header_downloader,
                 body_downloader,
+                etl_file_size,
             ),
             executor_factory,
         }
@@ -162,6 +164,8 @@ pub struct OnlineStages<Provider, H, B> {
     header_downloader: H,
     /// The block body downloader
     body_downloader: B,
+    /// The size of temporary files in bytes for ETL data collector.
+    etl_file_size: usize,
 }
 
 impl<Provider, H, B> OnlineStages<Provider, H, B> {
@@ -172,8 +176,9 @@ impl<Provider, H, B> OnlineStages<Provider, H, B> {
         consensus: Arc<dyn Consensus>,
         header_downloader: H,
         body_downloader: B,
+        etl_file_size: usize,
     ) -> Self {
-        Self { provider, header_mode, consensus, header_downloader, body_downloader }
+        Self { provider, header_mode, consensus, header_downloader, body_downloader, etl_file_size }
     }
 }
 
@@ -198,9 +203,16 @@ where
         mode: HeaderSyncMode,
         header_downloader: H,
         consensus: Arc<dyn Consensus>,
+        etl_file_size: usize,
     ) -> StageSetBuilder<DB> {
         StageSetBuilder::default()
-            .add_stage(HeaderStage::new(provider, header_downloader, mode, consensus.clone()))
+            .add_stage(HeaderStage::new(
+                provider,
+                header_downloader,
+                mode,
+                consensus.clone(),
+                etl_file_size,
+            ))
             .add_stage(bodies)
     }
 }
@@ -219,6 +231,7 @@ where
                 self.header_downloader,
                 self.header_mode,
                 self.consensus.clone(),
+                self.etl_file_size,
             ))
             .add_stage(BodyStage::new(self.body_downloader))
     }

crates/stages/src/stages/headers.rs

@@ -74,6 +74,7 @@ where
         downloader: Downloader,
         mode: HeaderSyncMode,
         consensus: Arc<dyn Consensus>,
+        etl_file_size: usize,
     ) -> Self {
         Self {
             provider: database,
@@ -81,8 +82,8 @@ where
             mode,
             consensus,
             sync_gap: None,
-            hash_collector: Collector::new(100 * (1024 * 1024)),
-            header_collector: Collector::new(100 * (1024 * 1024)),
+            hash_collector: Collector::new(etl_file_size / 2),
+            header_collector: Collector::new(etl_file_size / 2),
             is_etl_ready: false,
         }
     }
@@ -419,6 +420,7 @@ mod tests {
                     (*self.downloader_factory)(),
                     HeaderSyncMode::Tip(self.channel.1.clone()),
                     self.consensus.clone(),
+                    500 * (1024 * 1024),
                 )
             }
         }

crates/stages/src/stages/tx_lookup.rs

@@ -32,19 +32,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 }
     }
 }
 
@@ -99,7 +100,7 @@ impl<DB: Database> Stage<DB> for TransactionLookupStage {
         }
 
         // 500MB temporary files
-        let mut hash_collector: Collector<TxHash, TxNumber> = Collector::new(500 * (1024 * 1024));
+        let mut hash_collector: Collector<TxHash, TxNumber> = Collector::new(self.etl_file_size);
 
         debug!(
             target: "sync::stages::transaction_lookup",
@@ -397,12 +398,18 @@ 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 +456,11 @@ 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,
+            }
         }
     }