feat(bin, prune): pass prune segments from CLI & refactor modes (#4964)

crates/stages/src/stages/index_account_history.rs

@@ -2,7 +2,7 @@ use crate::{ExecInput, ExecOutput, Stage, StageError, UnwindInput, UnwindOutput}
 use reth_db::database::Database;
 use reth_primitives::{
     stage::{StageCheckpoint, StageId},
-    PruneCheckpoint, PruneModes, PruneSegment,
+    PruneCheckpoint, PruneMode, PruneSegment,
 };
 use reth_provider::{
     AccountExtReader, DatabaseProviderRW, HistoryWriter, PruneCheckpointReader,
@@ -19,19 +19,19 @@ pub struct IndexAccountHistoryStage {
     /// flow will be returned to the pipeline for commit.
     pub commit_threshold: u64,
     /// Pruning configuration.
-    pub prune_modes: PruneModes,
+    pub prune_mode: Option<PruneMode>,
 }
 
 impl IndexAccountHistoryStage {
     /// Create new instance of [IndexAccountHistoryStage].
-    pub fn new(commit_threshold: u64, prune_modes: PruneModes) -> Self {
-        Self { commit_threshold, prune_modes }
+    pub fn new(commit_threshold: u64, prune_mode: Option<PruneMode>) -> Self {
+        Self { commit_threshold, prune_mode }
     }
 }
 
 impl Default for IndexAccountHistoryStage {
     fn default() -> Self {
-        Self { commit_threshold: 100_000, prune_modes: PruneModes::none() }
+        Self { commit_threshold: 100_000, prune_mode: None }
     }
 }
 
@@ -48,8 +48,11 @@ impl<DB: Database> Stage<DB> for IndexAccountHistoryStage {
         provider: &DatabaseProviderRW<'_, &DB>,
         mut input: ExecInput,
     ) -> Result<ExecOutput, StageError> {
-        if let Some((target_prunable_block, prune_mode)) =
-            self.prune_modes.prune_target_block_account_history(input.target())?
+        if let Some((target_prunable_block, prune_mode)) = self
+            .prune_mode
+            .map(|mode| mode.prune_target_block(input.target(), PruneSegment::AccountHistory))
+            .transpose()?
+            .flatten()
         {
             if target_prunable_block > input.checkpoint().block_number {
                 input.checkpoint = Some(StageCheckpoint::new(target_prunable_block));
@@ -396,7 +399,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn insert_index_with_prune_modes() {
+    async fn insert_index_with_prune_mode() {
         // init
         let tx = TestTransaction::default();
 
@@ -426,10 +429,7 @@ mod tests {
         // run
         let input = ExecInput { target: Some(20000), ..Default::default() };
         let mut stage = IndexAccountHistoryStage {
-            prune_modes: PruneModes {
-                account_history: Some(PruneMode::Before(36)),
-                ..Default::default()
-            },
+            prune_mode: Some(PruneMode::Before(36)),
             ..Default::default()
         };
         let factory = ProviderFactory::new(tx.tx.as_ref(), MAINNET.clone());
@@ -455,16 +455,12 @@ mod tests {
     struct IndexAccountHistoryTestRunner {
         pub(crate) tx: TestTransaction,
         commit_threshold: u64,
-        prune_modes: PruneModes,
+        prune_mode: Option<PruneMode>,
     }
 
     impl Default for IndexAccountHistoryTestRunner {
         fn default() -> Self {
-            Self {
-                tx: TestTransaction::default(),
-                commit_threshold: 1000,
-                prune_modes: PruneModes::none(),
-            }
+            Self { tx: TestTransaction::default(), commit_threshold: 1000, prune_mode: None }
         }
     }
 
@@ -476,10 +472,7 @@ mod tests {
         }
 
         fn stage(&self) -> Self::S {
-            Self::S {
-                commit_threshold: self.commit_threshold,
-                prune_modes: self.prune_modes.clone(),
-            }
+            Self::S { commit_threshold: self.commit_threshold, prune_mode: self.prune_mode }
         }
     }
 

crates/stages/src/stages/index_storage_history.rs

@@ -2,7 +2,7 @@ use crate::{ExecInput, ExecOutput, Stage, StageError, UnwindInput, UnwindOutput}
 use reth_db::{database::Database, models::BlockNumberAddress};
 use reth_primitives::{
     stage::{StageCheckpoint, StageId},
-    PruneCheckpoint, PruneModes, PruneSegment,
+    PruneCheckpoint, PruneMode, PruneSegment,
 };
 use reth_provider::{
     DatabaseProviderRW, HistoryWriter, PruneCheckpointReader, PruneCheckpointWriter, StorageReader,
@@ -18,19 +18,19 @@ pub struct IndexStorageHistoryStage {
     /// flow will be returned to the pipeline for commit.
     pub commit_threshold: u64,
     /// Pruning configuration.
-    pub prune_modes: PruneModes,
+    pub prune_mode: Option<PruneMode>,
 }
 
 impl IndexStorageHistoryStage {
     /// Create new instance of [IndexStorageHistoryStage].
-    pub fn new(commit_threshold: u64, prune_modes: PruneModes) -> Self {
-        Self { commit_threshold, prune_modes }
+    pub fn new(commit_threshold: u64, prune_mode: Option<PruneMode>) -> Self {
+        Self { commit_threshold, prune_mode }
     }
 }
 
 impl Default for IndexStorageHistoryStage {
     fn default() -> Self {
-        Self { commit_threshold: 100_000, prune_modes: PruneModes::none() }
+        Self { commit_threshold: 100_000, prune_mode: None }
     }
 }
 
@@ -47,8 +47,11 @@ impl<DB: Database> Stage<DB> for IndexStorageHistoryStage {
         provider: &DatabaseProviderRW<'_, &DB>,
         mut input: ExecInput,
     ) -> Result<ExecOutput, StageError> {
-        if let Some((target_prunable_block, prune_mode)) =
-            self.prune_modes.prune_target_block_storage_history(input.target())?
+        if let Some((target_prunable_block, prune_mode)) = self
+            .prune_mode
+            .map(|mode| mode.prune_target_block(input.target(), PruneSegment::StorageHistory))
+            .transpose()?
+            .flatten()
         {
             if target_prunable_block > input.checkpoint().block_number {
                 input.checkpoint = Some(StageCheckpoint::new(target_prunable_block));
@@ -409,7 +412,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn insert_index_with_prune_modes() {
+    async fn insert_index_with_prune_mode() {
         // init
         let tx = TestTransaction::default();
 
@@ -439,10 +442,7 @@ mod tests {
         // run
         let input = ExecInput { target: Some(20000), ..Default::default() };
         let mut stage = IndexStorageHistoryStage {
-            prune_modes: PruneModes {
-                storage_history: Some(PruneMode::Before(36)),
-                ..Default::default()
-            },
+            prune_mode: Some(PruneMode::Before(36)),
             ..Default::default()
         };
         let factory = ProviderFactory::new(tx.tx.as_ref(), MAINNET.clone());
@@ -468,16 +468,12 @@ mod tests {
     struct IndexStorageHistoryTestRunner {
         pub(crate) tx: TestTransaction,
         commit_threshold: u64,
-        prune_modes: PruneModes,
+        prune_mode: Option<PruneMode>,
     }
 
     impl Default for IndexStorageHistoryTestRunner {
         fn default() -> Self {
-            Self {
-                tx: TestTransaction::default(),
-                commit_threshold: 1000,
-                prune_modes: PruneModes::none(),
-            }
+            Self { tx: TestTransaction::default(), commit_threshold: 1000, prune_mode: None }
         }
     }
 
@@ -489,10 +485,7 @@ mod tests {
         }
 
         fn stage(&self) -> Self::S {
-            Self::S {
-                commit_threshold: self.commit_threshold,
-                prune_modes: self.prune_modes.clone(),
-            }
+            Self::S { commit_threshold: self.commit_threshold, prune_mode: self.prune_mode }
         }
     }
 

crates/stages/src/stages/mod.rs

@@ -156,8 +156,10 @@ mod tests {
             );
 
             // Check AccountHistory
-            let mut acc_indexing_stage =
-                IndexAccountHistoryStage { prune_modes: prune_modes.clone(), ..Default::default() };
+            let mut acc_indexing_stage = IndexAccountHistoryStage {
+                prune_mode: prune_modes.account_history,
+                ..Default::default()
+            };
 
             if let Some(PruneMode::Full) = prune_modes.account_history {
                 // Full is not supported
@@ -170,8 +172,10 @@ mod tests {
             }
 
             // Check StorageHistory
-            let mut storage_indexing_stage =
-                IndexStorageHistoryStage { prune_modes: prune_modes.clone(), ..Default::default() };
+            let mut storage_indexing_stage = IndexStorageHistoryStage {
+                prune_mode: prune_modes.storage_history,
+                ..Default::default()
+            };
 
             if let Some(PruneMode::Full) = prune_modes.storage_history {
                 // Full is not supported

crates/stages/src/stages/tx_lookup.rs

@@ -12,7 +12,7 @@ use reth_interfaces::provider::ProviderError;
 use reth_primitives::{
     keccak256,
     stage::{EntitiesCheckpoint, StageCheckpoint, StageId},
-    PruneCheckpoint, PruneModes, PruneSegment, TransactionSignedNoHash, TxNumber, B256,
+    PruneCheckpoint, PruneMode, PruneSegment, TransactionSignedNoHash, TxNumber, B256,
 };
 use reth_provider::{
     BlockReader, DatabaseProviderRW, PruneCheckpointReader, PruneCheckpointWriter,
@@ -29,19 +29,19 @@ use tracing::*;
 pub struct TransactionLookupStage {
     /// The number of lookup entries to commit at once
     commit_threshold: u64,
-    prune_modes: PruneModes,
+    prune_mode: Option<PruneMode>,
 }
 
 impl Default for TransactionLookupStage {
     fn default() -> Self {
-        Self { commit_threshold: 5_000_000, prune_modes: PruneModes::none() }
+        Self { commit_threshold: 5_000_000, prune_mode: None }
     }
 }
 
 impl TransactionLookupStage {
     /// Create new instance of [TransactionLookupStage].
-    pub fn new(commit_threshold: u64, prune_modes: PruneModes) -> Self {
-        Self { commit_threshold, prune_modes }
+    pub fn new(commit_threshold: u64, prune_mode: Option<PruneMode>) -> Self {
+        Self { commit_threshold, prune_mode }
     }
 }
 
@@ -58,8 +58,11 @@ impl<DB: Database> Stage<DB> for TransactionLookupStage {
         provider: &DatabaseProviderRW<'_, &DB>,
         mut input: ExecInput,
     ) -> Result<ExecOutput, StageError> {
-        if let Some((target_prunable_block, prune_mode)) =
-            self.prune_modes.prune_target_block_transaction_lookup(input.target())?
+        if let Some((target_prunable_block, prune_mode)) = self
+            .prune_mode
+            .map(|mode| mode.prune_target_block(input.target(), PruneSegment::TransactionLookup))
+            .transpose()?
+            .flatten()
         {
             if target_prunable_block > input.checkpoint().block_number {
                 input.checkpoint = Some(StageCheckpoint::new(target_prunable_block));
@@ -382,10 +385,7 @@ mod tests {
             random_block_range(&mut rng, stage_progress + 1..=previous_stage, B256::ZERO, 0..2);
         runner.tx.insert_blocks(seed.iter(), None).expect("failed to seed execution");
 
-        runner.set_prune_modes(PruneModes {
-            transaction_lookup: Some(PruneMode::Before(prune_target)),
-            ..Default::default()
-        });
+        runner.set_prune_mode(PruneMode::Before(prune_target));
 
         let rx = runner.execute(input);
 
@@ -469,16 +469,12 @@ mod tests {
     struct TransactionLookupTestRunner {
         tx: TestTransaction,
         commit_threshold: u64,
-        prune_modes: PruneModes,
+        prune_mode: Option<PruneMode>,
     }
 
     impl Default for TransactionLookupTestRunner {
         fn default() -> Self {
-            Self {
-                tx: TestTransaction::default(),
-                commit_threshold: 1000,
-                prune_modes: PruneModes::none(),
-            }
+            Self { tx: TestTransaction::default(), commit_threshold: 1000, prune_mode: None }
         }
     }
 
@@ -487,8 +483,8 @@ mod tests {
             self.commit_threshold = threshold;
         }
 
-        fn set_prune_modes(&mut self, prune_modes: PruneModes) {
-            self.prune_modes = prune_modes;
+        fn set_prune_mode(&mut self, prune_mode: PruneMode) {
+            self.prune_mode = Some(prune_mode);
         }
 
         /// # Panics
@@ -528,7 +524,7 @@ mod tests {
         fn stage(&self) -> Self::S {
             TransactionLookupStage {
                 commit_threshold: self.commit_threshold,
-                prune_modes: self.prune_modes.clone(),
+                prune_mode: self.prune_mode,
             }
         }
     }
@@ -556,9 +552,13 @@ mod tests {
                     let provider = self.tx.inner();
 
                     if let Some((target_prunable_block, _)) = self
-                        .prune_modes
-                        .prune_target_block_transaction_lookup(input.target())
+                        .prune_mode
+                        .map(|mode| {
+                            mode.prune_target_block(input.target(), PruneSegment::TransactionLookup)
+                        })
+                        .transpose()
                         .expect("prune target block for transaction lookup")
+                        .flatten()
                     {
                         if target_prunable_block > input.checkpoint().block_number {
                             input.checkpoint = Some(StageCheckpoint::new(target_prunable_block));