feat(prune, stages): prune pipeline stage (paradigmxyz#9419)

crates/cli/commands/src/prune.rs

@@ -30,8 +30,8 @@ impl PruneCommand {
             info!(target: "reth::cli", ?prune_tip, ?prune_config, "Pruning data from database...");
             // Run the pruner according to the configuration, and don't enforce any limits on it
             let mut pruner = PrunerBuilder::new(prune_config)
-                .prune_delete_limit(usize::MAX)
-                .build(provider_factory);
+                .delete_limit_per_block(usize::MAX)
+                .build_with_provider_factory(provider_factory);
 
             pruner.run(prune_tip)?;
             info!(target: "reth::cli", "Pruned data from database");

crates/config/src/config.rs

@@ -71,6 +71,8 @@ pub struct StageConfig {
     pub sender_recovery: SenderRecoveryConfig,
     /// Execution stage configuration.
     pub execution: ExecutionConfig,
+    /// Prune stage configuration.
+    pub prune: PruneStageConfig,
     /// Account Hashing stage configuration.
     pub account_hashing: HashingConfig,
     /// Storage Hashing stage configuration.
@@ -228,6 +230,20 @@ impl From<ExecutionConfig> for ExecutionStageThresholds {
     }
 }
 
+/// Prune stage configuration.
+#[derive(Debug, Clone, Copy, Deserialize, PartialEq, Eq, Serialize)]
+#[serde(default)]
+pub struct PruneStageConfig {
+    /// The maximum number of entries to prune before committing progress to the database.
+    pub commit_threshold: usize,
+}
+
+impl Default for PruneStageConfig {
+    fn default() -> Self {
+        Self { commit_threshold: 1_000_000 }
+    }
+}
+
 /// Hashing stage configuration.
 #[derive(Debug, Clone, Copy, Deserialize, PartialEq, Eq, Serialize)]
 #[serde(default)]

crates/consensus/beacon/src/engine/hooks/prune.rs

@@ -9,6 +9,7 @@ use metrics::Counter;
 use reth_db_api::database::Database;
 use reth_errors::{RethError, RethResult};
 use reth_primitives::BlockNumber;
+use reth_provider::ProviderFactory;
 use reth_prune::{Pruner, PrunerError, PrunerWithResult};
 use reth_tasks::TaskSpawner;
 use std::{
@@ -39,7 +40,10 @@ impl<DB: fmt::Debug> fmt::Debug for PruneHook<DB> {
 
 impl<DB: Database + 'static> PruneHook<DB> {
     /// Create a new instance
-    pub fn new(pruner: Pruner<DB>, pruner_task_spawner: Box<dyn TaskSpawner>) -> Self {
+    pub fn new(
+        pruner: Pruner<DB, ProviderFactory<DB>>,
+        pruner_task_spawner: Box<dyn TaskSpawner>,
+    ) -> Self {
         Self {
             pruner_state: PrunerState::Idle(Some(pruner)),
             pruner_task_spawner,
@@ -151,9 +155,9 @@ impl<DB: Database + 'static> EngineHook for PruneHook<DB> {
 #[derive(Debug)]
 enum PrunerState<DB> {
     /// Pruner is idle.
-    Idle(Option<Pruner<DB>>),
+    Idle(Option<Pruner<DB, ProviderFactory<DB>>>),
     /// Pruner is running and waiting for a response
-    Running(oneshot::Receiver<PrunerWithResult<DB>>),
+    Running(oneshot::Receiver<PrunerWithResult<DB, ProviderFactory<DB>>>),
 }
 
 #[derive(reth_metrics::Metrics)]

crates/consensus/beacon/src/engine/test_utils.rs

@@ -26,7 +26,7 @@ use reth_payload_builder::test_utils::spawn_test_payload_service;
 use reth_primitives::{BlockNumber, B256};
 use reth_provider::{
     providers::BlockchainProvider, test_utils::create_test_provider_factory_with_chain_spec,
-    ExecutionOutcome,
+    ExecutionOutcome, ProviderFactory,
 };
 use reth_prune::Pruner;
 use reth_prune_types::PruneModes;
@@ -400,7 +400,7 @@ where
         let blockchain_provider =
             BlockchainProvider::with_latest(provider_factory.clone(), tree, latest);
 
-        let pruner = Pruner::new(
+        let pruner = Pruner::<_, ProviderFactory<_>>::new(
             provider_factory.clone(),
             vec![],
             5,

crates/engine/tree/src/persistence.rs

@@ -34,7 +34,7 @@ pub struct Persistence<DB> {
     /// Handle for the static file service.
     static_file_handle: StaticFileServiceHandle,
     /// The pruner
-    pruner: Pruner<DB>,
+    pruner: Pruner<DB, ProviderFactory<DB>>,
 }
 
 impl<DB: Database> Persistence<DB> {
@@ -43,7 +43,7 @@ impl<DB: Database> Persistence<DB> {
         provider: ProviderFactory<DB>,
         incoming: Receiver<PersistenceAction>,
         static_file_handle: StaticFileServiceHandle,
-        pruner: Pruner<DB>,
+        pruner: Pruner<DB, ProviderFactory<DB>>,
     ) -> Self {
         Self { provider, incoming, static_file_handle, pruner }
     }
@@ -145,7 +145,7 @@ where
     fn spawn_new(
         provider: ProviderFactory<DB>,
         static_file_handle: StaticFileServiceHandle,
-        pruner: Pruner<DB>,
+        pruner: Pruner<DB, ProviderFactory<DB>>,
     ) -> PersistenceHandle {
         let (tx, rx) = std::sync::mpsc::channel();
         let service = Self::new(provider, rx, static_file_handle, pruner);
@@ -313,7 +313,15 @@ mod tests {
         let (finished_exex_height_tx, finished_exex_height_rx) =
             tokio::sync::watch::channel(FinishedExExHeight::NoExExs);
 
-        let pruner = Pruner::new(provider.clone(), vec![], 5, 0, 5, None, finished_exex_height_rx);
+        let pruner = Pruner::<_, ProviderFactory<_>>::new(
+            provider.clone(),
+            vec![],
+            5,
+            0,
+            5,
+            None,
+            finished_exex_height_rx,
+        );
 
         let (static_file_sender, _static_file_receiver) = channel();
         let static_file_handle = StaticFileServiceHandle::new(static_file_sender);

crates/node/builder/src/launch/common.rs

@@ -329,7 +329,7 @@ impl<R> LaunchContextWith<Attached<WithConfigs, R>> {
     /// Returns an initialized [`PrunerBuilder`] based on the configured [`PruneConfig`]
     pub fn pruner_builder(&self) -> PrunerBuilder {
         PrunerBuilder::new(self.prune_config().unwrap_or_default())
-            .prune_delete_limit(self.chain_spec().prune_delete_limit)
+            .delete_limit_per_block(self.chain_spec().prune_delete_limit)
             .timeout(PrunerBuilder::DEFAULT_TIMEOUT)
     }
 

crates/node/builder/src/launch/mod.rs

@@ -231,14 +231,15 @@ where
 
         let initial_target = ctx.node_config().debug.tip;
 
-        let mut pruner_builder =
-            ctx.pruner_builder().max_reorg_depth(ctx.tree_config().max_reorg_depth() as usize);
+        let mut pruner_builder = ctx
+            .pruner_builder()
+            .prune_max_blocks_per_run(ctx.tree_config().max_reorg_depth() as usize);
         if let Some(exex_manager_handle) = &exex_manager_handle {
             pruner_builder =
                 pruner_builder.finished_exex_height(exex_manager_handle.finished_height());
         }
 
-        let pruner = pruner_builder.build(ctx.provider_factory().clone());
+        let pruner = pruner_builder.build_with_provider_factory(ctx.provider_factory().clone());
 
         let pruner_events = pruner.events();
         info!(target: "reth::cli", prune_config=?ctx.prune_config().unwrap_or_default(), "Pruner initialized");

crates/prune/prune/src/builder.rs

@@ -15,12 +15,12 @@ pub struct PrunerBuilder {
     block_interval: usize,
     /// Pruning configuration for every part of the data that can be pruned.
     segments: PruneModes,
-    /// The number of blocks that can be re-orged.
-    max_reorg_depth: usize,
+    /// The maximum number of blocks that can be pruned per run.
+    prune_max_blocks_per_run: usize,
     /// The delete limit for pruner, per block. In the actual pruner run it will be multiplied by
     /// the amount of blocks between pruner runs to account for the difference in amount of new
     /// data coming in.
-    prune_delete_limit: usize,
+    delete_limit_per_block: usize,
     /// Time a pruner job can run before timing out.
     timeout: Option<Duration>,
     /// The finished height of all `ExEx`'s.
@@ -51,14 +51,14 @@ impl PrunerBuilder {
     }
 
     /// Sets the number of blocks that can be re-orged.
-    pub const fn max_reorg_depth(mut self, max_reorg_depth: usize) -> Self {
-        self.max_reorg_depth = max_reorg_depth;
+    pub const fn prune_max_blocks_per_run(mut self, prune_max_blocks_per_run: usize) -> Self {
+        self.prune_max_blocks_per_run = prune_max_blocks_per_run;
         self
     }
 
     /// Sets the delete limit for pruner, per block.
-    pub const fn prune_delete_limit(mut self, prune_delete_limit: usize) -> Self {
-        self.prune_delete_limit = prune_delete_limit;
+    pub const fn delete_limit_per_block(mut self, delete_limit_per_block: usize) -> Self {
+        self.delete_limit_per_block = delete_limit_per_block;
         self
     }
 
@@ -80,16 +80,33 @@ impl PrunerBuilder {
         self
     }
 
-    /// Builds a [Pruner] from the current configuration.
-    pub fn build<DB: Database>(self, provider_factory: ProviderFactory<DB>) -> Pruner<DB> {
+    /// Builds a [Pruner] from the current configuration with the given provider factory.
+    pub fn build_with_provider_factory<DB: Database>(
+        self,
+        provider_factory: ProviderFactory<DB>,
+    ) -> Pruner<DB, ProviderFactory<DB>> {
         let segments = SegmentSet::<DB>::from_prune_modes(self.segments);
 
-        Pruner::new(
+        Pruner::<_, ProviderFactory<DB>>::new(
             provider_factory,
             segments.into_vec(),
             self.block_interval,
-            self.prune_delete_limit,
-            self.max_reorg_depth,
+            self.delete_limit_per_block,
+            self.prune_max_blocks_per_run,
+            self.timeout,
+            self.finished_exex_height,
+        )
+    }
+
+    /// Builds a [Pruner] from the current configuration.
+    pub fn build<DB: Database>(self) -> Pruner<DB, ()> {
+        let segments = SegmentSet::<DB>::from_prune_modes(self.segments);
+
+        Pruner::<_, ()>::new(
+            segments.into_vec(),
+            self.block_interval,
+            self.delete_limit_per_block,
+            self.prune_max_blocks_per_run,
             self.timeout,
             self.finished_exex_height,
         )
@@ -101,8 +118,8 @@ impl Default for PrunerBuilder {
         Self {
             block_interval: 5,
             segments: PruneModes::none(),
-            max_reorg_depth: 64,
-            prune_delete_limit: MAINNET.prune_delete_limit,
+            prune_max_blocks_per_run: 64,
+            delete_limit_per_block: MAINNET.prune_delete_limit,
             timeout: None,
             finished_exex_height: watch::channel(FinishedExExHeight::NoExExs).1,
         }