From 471c28e889083ec5e8f5bbdc0b57ede8023471dc Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <a.shekhirin@gmail.com>
Date: Wed, 29 Nov 2023 19:29:39 +0000
Subject: [PATCH] feat(pruner): limit number of blocks to prune per run (#5627)

---
 bin/reth/src/node/mod.rs                      |  6 +++-
 .../consensus/beacon/src/engine/test_utils.rs |  1 +
 crates/prune/src/pruner.rs                    | 31 +++++++++++++------
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/bin/reth/src/node/mod.rs b/bin/reth/src/node/mod.rs
index de078146f..815b77a71 100644
--- a/bin/reth/src/node/mod.rs
+++ b/bin/reth/src/node/mod.rs
@@ -289,9 +289,10 @@ impl<Ext: RethCliExt> NodeCommand<Ext> {
             Arc::clone(&consensus),
             EvmProcessorFactory::new(self.chain.clone()),
         );
+        let tree_config = BlockchainTreeConfig::default();
         let tree = BlockchainTree::new(
             tree_externals,
-            BlockchainTreeConfig::default(),
+            tree_config,
             prune_config.clone().map(|config| config.segments),
         )?
         .with_sync_metrics_tx(sync_metrics_tx.clone());
@@ -473,6 +474,7 @@ impl<Ext: RethCliExt> NodeCommand<Ext> {
             let mut pruner = self.build_pruner(
                 &prune_config,
                 db.clone(),
+                tree_config,
                 snapshotter.highest_snapshot_receiver(),
             );
 
@@ -975,6 +977,7 @@ impl<Ext: RethCliExt> NodeCommand<Ext> {
         &self,
         config: &PruneConfig,
         db: DB,
+        tree_config: BlockchainTreeConfig,
         highest_snapshots_rx: HighestSnapshotsTracker,
     ) -> Pruner<DB> {
         let segments = SegmentSet::default()
@@ -1012,6 +1015,7 @@ impl<Ext: RethCliExt> NodeCommand<Ext> {
             segments.into_vec(),
             config.block_interval,
             self.chain.prune_delete_limit,
+            tree_config.max_reorg_depth() as usize,
             highest_snapshots_rx,
         )
     }
diff --git a/crates/consensus/beacon/src/engine/test_utils.rs b/crates/consensus/beacon/src/engine/test_utils.rs
index 781161bb3..0622c7734 100644
--- a/crates/consensus/beacon/src/engine/test_utils.rs
+++ b/crates/consensus/beacon/src/engine/test_utils.rs
@@ -533,6 +533,7 @@ where
             vec![],
             5,
             self.base_config.chain_spec.prune_delete_limit,
+            config.max_reorg_depth() as usize,
             watch::channel(None).1,
         );
 
diff --git a/crates/prune/src/pruner.rs b/crates/prune/src/pruner.rs
index 12214d6d3..a74825fd4 100644
--- a/crates/prune/src/pruner.rs
+++ b/crates/prune/src/pruner.rs
@@ -34,6 +34,9 @@ pub struct Pruner<DB> {
     previous_tip_block_number: Option<BlockNumber>,
     /// Maximum total entries to prune (delete from database) per block.
     delete_limit: usize,
+    /// Maximum number of blocks to be pruned per run, as an additional restriction to
+    /// `previous_tip_block_number`.
+    prune_max_blocks_per_run: usize,
     #[allow(dead_code)]
     highest_snapshots_tracker: HighestSnapshotsTracker,
     metrics: Metrics,
@@ -48,6 +51,7 @@ impl<DB: Database> Pruner<DB> {
         segments: Vec<Arc<dyn Segment<DB>>>,
         min_block_interval: usize,
         delete_limit: usize,
+        prune_max_blocks_per_run: usize,
         highest_snapshots_tracker: HighestSnapshotsTracker,
     ) -> Self {
         Self {
@@ -56,6 +60,7 @@ impl<DB: Database> Pruner<DB> {
             min_block_interval,
             previous_tip_block_number: None,
             delete_limit,
+            prune_max_blocks_per_run,
             highest_snapshots_tracker,
             metrics: Metrics::default(),
             listeners: Default::default(),
@@ -87,14 +92,22 @@ impl<DB: Database> Pruner<DB> {
         // TODO(alexey): prune snapshotted segments of data (headers, transactions)
         let highest_snapshots = *self.highest_snapshots_tracker.borrow();
 
-        // Multiply `delete_limit` (number of row to delete per block) by number of blocks since
-        // last pruner run. `previous_tip_block_number` is close to `tip_block_number`, usually
-        // within `self.block_interval` blocks, so `delete_limit` will not be too high. Also see
-        // docs for `self.previous_tip_block_number`.
-        let mut delete_limit = self.delete_limit *
-            self.previous_tip_block_number
-                .map_or(1, |previous_tip_block_number| tip_block_number - previous_tip_block_number)
-                as usize;
+        // Multiply `self.delete_limit` (number of rows to delete per block) by number of blocks
+        // since last pruner run. `self.previous_tip_block_number` is close to
+        // `tip_block_number`, usually within `self.block_interval` blocks, so
+        // `delete_limit` will not be too high. If it's too high, we additionally limit it by
+        // `self.prune_max_blocks_per_run`.
+        //
+        // Also see docs for `self.previous_tip_block_number`.
+        let blocks_since_last_run =
+            (self.previous_tip_block_number.map_or(1, |previous_tip_block_number| {
+                // Saturating subtraction is needed for the case when the chain was reverted,
+                // meaning current block number might be less than the previous tip
+                // block number.
+                tip_block_number.saturating_sub(previous_tip_block_number) as usize
+            }))
+            .min(self.prune_max_blocks_per_run);
+        let mut delete_limit = self.delete_limit * blocks_since_last_run;
 
         for segment in &self.segments {
             if delete_limit == 0 {
@@ -259,7 +272,7 @@ mod tests {
     #[test]
     fn is_pruning_needed() {
         let db = create_test_rw_db();
-        let mut pruner = Pruner::new(db, MAINNET.clone(), vec![], 5, 0, watch::channel(None).1);
+        let mut pruner = Pruner::new(db, MAINNET.clone(), vec![], 5, 0, 5, watch::channel(None).1);
 
         // No last pruned block number was set before
         let first_block_number = 1;