fix: ensure MerkleChangeSets pruner only runs if pipeline stage has finished (#20073)

Author: joshieDo

Cargo.lock

@@ -21,6 +21,7 @@ reth-tokio-util.workspace = true
 reth-config.workspace = true
 reth-prune-types.workspace = true
 reth-primitives-traits.workspace = true
+reth-stages-types.workspace = true
 reth-static-file-types.workspace = true
 
 # ethereum

crates/prune/prune/Cargo.toml

@@ -7,7 +7,7 @@ use reth_primitives_traits::NodePrimitives;
 use reth_provider::{
     providers::StaticFileProvider, BlockReader, ChainStateBlockReader, DBProvider,
     DatabaseProviderFactory, NodePrimitivesProvider, PruneCheckpointReader, PruneCheckpointWriter,
-    StaticFileProviderFactory, StorageSettingsCache,
+    StageCheckpointReader, StaticFileProviderFactory, StorageSettingsCache,
 };
 use reth_prune_types::PruneModes;
 use std::time::Duration;
@@ -81,6 +81,7 @@ impl PrunerBuilder {
                                 + BlockReader<Transaction: Encodable2718>
                                 + ChainStateBlockReader
                                 + StorageSettingsCache
+                                + StageCheckpointReader
                                 + StaticFileProviderFactory<
                     Primitives: NodePrimitives<SignedTx: Value, Receipt: Value, BlockHeader: Value>,
                 >,
@@ -114,7 +115,8 @@ impl PrunerBuilder {
             + ChainStateBlockReader
             + PruneCheckpointWriter
             + PruneCheckpointReader
-            + StorageSettingsCache,
+            + StorageSettingsCache
+            + StageCheckpointReader,
     {
         let segments = SegmentSet::<Provider>::from_components(static_file_provider, self.segments);
 

crates/prune/prune/src/builder.rs

@@ -8,8 +8,10 @@ use alloy_primitives::BlockNumber;
 use reth_exex_types::FinishedExExHeight;
 use reth_provider::{
     DBProvider, DatabaseProviderFactory, PruneCheckpointReader, PruneCheckpointWriter,
+    StageCheckpointReader,
 };
 use reth_prune_types::{PruneProgress, PrunedSegmentInfo, PrunerOutput};
+use reth_stages_types::StageId;
 use reth_tokio_util::{EventSender, EventStream};
 use std::time::{Duration, Instant};
 use tokio::sync::watch;
@@ -100,7 +102,7 @@ where
 
 impl<Provider, S> Pruner<Provider, S>
 where
-    Provider: PruneCheckpointReader + PruneCheckpointWriter,
+    Provider: PruneCheckpointReader + PruneCheckpointWriter + StageCheckpointReader,
 {
     /// Listen for events on the pruner.
     pub fn events(&self) -> EventStream<PrunerEvent> {
@@ -200,6 +202,19 @@ where
                 .transpose()?
                 .flatten()
             {
+                // Check if segment has a required stage that must be finished first
+                if let Some(required_stage) = segment.required_stage() &&
+                    !is_stage_finished(provider, required_stage)?
+                {
+                    debug!(
+                        target: "pruner",
+                        segment = ?segment.segment(),
+                        ?required_stage,
+                        "Segment's required stage not finished, skipping"
+                    );
+                    continue
+                }
+
                 debug!(
                     target: "pruner",
                     segment = ?segment.segment(),
@@ -318,7 +333,9 @@ where
 
 impl<PF> Pruner<PF::ProviderRW, PF>
 where
-    PF: DatabaseProviderFactory<ProviderRW: PruneCheckpointWriter + PruneCheckpointReader>,
+    PF: DatabaseProviderFactory<
+        ProviderRW: PruneCheckpointWriter + PruneCheckpointReader + StageCheckpointReader,
+    >,
 {
     /// Run the pruner. This will only prune data up to the highest finished ExEx height, if there
     /// are no ExExes.
@@ -333,6 +350,19 @@ where
     }
 }
 
+/// Checks if the given stage has caught up with the `Finish` stage.
+///
+/// Returns `true` if the stage checkpoint is >= the Finish stage checkpoint.
+fn is_stage_finished<Provider: StageCheckpointReader>(
+    provider: &Provider,
+    stage_id: StageId,
+) -> Result<bool, PrunerError> {
+    let stage_checkpoint = provider.get_stage_checkpoint(stage_id)?.map(|c| c.block_number);
+    let finish_checkpoint = provider.get_stage_checkpoint(StageId::Finish)?.map(|c| c.block_number);
+
+    Ok(stage_checkpoint >= finish_checkpoint)
+}
+
 #[cfg(test)]
 mod tests {
     use crate::Pruner;

crates/prune/prune/src/pruner.rs

@@ -11,6 +11,7 @@ use reth_prune_types::{
     PruneCheckpoint, PruneMode, PruneProgress, PrunePurpose, PruneSegment, SegmentOutput,
     SegmentOutputCheckpoint,
 };
+use reth_stages_types::StageId;
 use reth_static_file_types::StaticFileSegment;
 pub use set::SegmentSet;
 use std::{fmt::Debug, ops::RangeInclusive};
@@ -84,6 +85,14 @@ pub trait Segment<Provider>: Debug + Send + Sync {
     {
         provider.save_prune_checkpoint(self.segment(), checkpoint)
     }
+
+    /// Returns the stage this segment depends on, if any.
+    ///
+    /// If this returns `Some(stage_id)`, the pruner will skip this segment if the stage
+    /// has not yet caught up with the `Finish` stage checkpoint.
+    fn required_stage(&self) -> Option<StageId> {
+        None
+    }
 }
 
 /// Segment pruning input, see [`Segment::prune`].

crates/prune/prune/src/segments/mod.rs

@@ -13,6 +13,7 @@ use reth_provider::{
 use reth_prune_types::{
     PruneCheckpoint, PruneMode, PrunePurpose, PruneSegment, SegmentOutput, SegmentOutputCheckpoint,
 };
+use reth_stages_types::StageId;
 use tracing::{instrument, trace};
 
 #[derive(Debug)]
@@ -47,6 +48,10 @@ where
         PrunePurpose::User
     }
 
+    fn required_stage(&self) -> Option<StageId> {
+        Some(StageId::MerkleChangeSets)
+    }
+
     #[instrument(level = "trace", target = "pruner", skip(self, provider), ret)]
     fn prune(&self, provider: &Provider, input: PruneInput) -> Result<SegmentOutput, PrunerError> {
         let Some(block_range) = input.get_next_block_range() else {

crates/prune/prune/src/segments/user/merkle_change_sets.rs

@@ -2,7 +2,7 @@ use reth_db_api::{table::Value, transaction::DbTxMut};
 use reth_primitives_traits::NodePrimitives;
 use reth_provider::{
     BlockReader, ChainStateBlockReader, DBProvider, PruneCheckpointReader, PruneCheckpointWriter,
-    StaticFileProviderFactory, StorageSettingsCache,
+    StageCheckpointReader, StaticFileProviderFactory, StorageSettingsCache,
 };
 use reth_prune::{
     PruneMode, PruneModes, PruneSegment, PrunerBuilder, SegmentOutput, SegmentOutputCheckpoint,
@@ -43,6 +43,7 @@ where
         + PruneCheckpointWriter
         + BlockReader
         + ChainStateBlockReader
+        + StageCheckpointReader
         + StaticFileProviderFactory<
             Primitives: NodePrimitives<SignedTx: Value, Receipt: Value, BlockHeader: Value>,
         > + StorageSettingsCache,
@@ -144,6 +145,7 @@ where
         + PruneCheckpointWriter
         + BlockReader
         + ChainStateBlockReader
+        + StageCheckpointReader
         + StaticFileProviderFactory<
             Primitives: NodePrimitives<SignedTx: Value, Receipt: Value, BlockHeader: Value>,
         > + StorageSettingsCache,