feat(pruner): respect batch size per run (#4246)

Co-authored-by: joshieDo <ranriver@protonmail.com>

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

@@ -3,7 +3,7 @@
 use futures::FutureExt;
 use reth_db::database::Database;
 use reth_primitives::BlockNumber;
-use reth_prune::{Pruner, PrunerError, PrunerWithResult};
+use reth_prune::{Pruner, PrunerResult, PrunerWithResult};
 use reth_tasks::TaskSpawner;
 use std::task::{ready, Context, Poll};
 use tokio::sync::oneshot;
@@ -116,7 +116,7 @@ pub(crate) enum EnginePruneEvent {
     /// If this is returned, the pruner is idle.
     Finished {
         /// Final result of the pruner run.
-        result: Result<(), PrunerError>,
+        result: PrunerResult,
     },
     /// Pruner task was dropped after it was started, unable to receive it because channel
     /// closed. This would indicate a panicked pruner task

crates/interfaces/src/test_utils/generators.rs

@@ -239,7 +239,7 @@ where
         // deposit in receiving account and update storage
         let (prev_to, storage): &mut (Account, BTreeMap<H256, U256>) = state.get_mut(&to).unwrap();
 
-        let old_entries = new_entries
+        let mut old_entries: Vec<_> = new_entries
             .into_iter()
             .filter_map(|entry| {
                 let old = if entry.value != U256::ZERO {
@@ -254,9 +254,12 @@ where
                 Some(StorageEntry { value: old.unwrap_or(U256::from(0)), ..entry })
             })
             .collect();
+        old_entries.sort_by_key(|entry| entry.key);
 
         changeset.push((to, *prev_to, old_entries));
 
+        changeset.sort_by_key(|(address, _, _)| *address);
+
         prev_to.balance = prev_to.balance.wrapping_add(transfer);
 
         changesets.push(changeset);

crates/primitives/src/prune/checkpoint.rs

@@ -1,4 +1,4 @@
-use crate::{prune::PruneMode, BlockNumber};
+use crate::{prune::PruneMode, BlockNumber, TxNumber};
 use reth_codecs::{main_codec, Compact};
 
 /// Saves the pruning progress of a stage.
@@ -7,7 +7,10 @@ use reth_codecs::{main_codec, Compact};
 #[cfg_attr(test, derive(Default))]
 pub struct PruneCheckpoint {
     /// Highest pruned block number.
-    pub block_number: BlockNumber,
+    /// If it's [None], the pruning for block `0` is not finished yet.
+    pub block_number: Option<BlockNumber>,
+    /// Highest pruned transaction number, if applicable.
+    pub tx_number: Option<TxNumber>,
     /// Prune mode.
     pub prune_mode: PruneMode,
 }

crates/primitives/src/prune/mod.rs

@@ -49,10 +49,14 @@ impl ReceiptsLogPruneConfig {
             // the BTreeMap (block = 0), otherwise it will be excluded.
             // Reminder that this BTreeMap works as an inclusion list that excludes (prunes) all
             // other receipts.
+            //
+            // Reminder, that we increment because the [`BlockNumber`] key of the new map should be
+            // viewed as `PruneMode::Before(block)`
             let block = (pruned_block + 1).max(
                 mode.prune_target_block(tip, MINIMUM_PRUNING_DISTANCE, PrunePart::ContractLogs)?
                     .map(|(block, _)| block)
-                    .unwrap_or_default(),
+                    .unwrap_or_default() +
+                    1,
             );
 
             map.entry(block).or_insert_with(Vec::new).push(address)