diff --git a/crates/stages/stages/src/stages/prune.rs b/crates/stages/stages/src/stages/prune.rs index e61acf54c1..fe107d1405 100644 --- a/crates/stages/stages/src/stages/prune.rs +++ b/crates/stages/stages/src/stages/prune.rs @@ -1,6 +1,8 @@ use reth_db_api::database::Database; -use reth_provider::{DatabaseProviderRW, PruneCheckpointReader}; -use reth_prune::{PruneMode, PruneModes, PruneSegment, PrunerBuilder}; +use reth_provider::{DatabaseProviderRW, PruneCheckpointReader, PruneCheckpointWriter}; +use reth_prune::{ + PruneMode, PruneModes, PruneSegment, PrunerBuilder, SegmentOutput, SegmentOutputCheckpoint, +}; use reth_stages_api::{ ExecInput, ExecOutput, Stage, StageCheckpoint, StageError, StageId, UnwindInput, UnwindOutput, }; @@ -49,7 +51,34 @@ impl Stage for PruneStage { if result.progress.is_finished() { Ok(ExecOutput { checkpoint: StageCheckpoint::new(input.target()), done: true }) } else { - info!(target: "sync::stages::prune::exec", segments = ?result.segments, "Pruner has more data to prune"); + if let Some((last_segment, last_segment_output)) = result.segments.last() { + match last_segment_output { + SegmentOutput { + progress, + pruned, + checkpoint: + checkpoint @ Some(SegmentOutputCheckpoint { block_number: Some(_), .. }), + } => { + info!( + target: "sync::stages::prune::exec", + ?last_segment, + ?progress, + ?pruned, + ?checkpoint, + "Last segment has more data to prune" + ) + } + SegmentOutput { progress, pruned, checkpoint: _ } => { + info!( + target: "sync::stages::prune::exec", + ?last_segment, + ?progress, + ?pruned, + "Last segment has more data to prune" + ) + } + } + } // We cannot set the checkpoint yet, because prune segments may have different highest // pruned block numbers Ok(ExecOutput { checkpoint: input.checkpoint(), done: false }) @@ -58,10 +87,16 @@ impl Stage for PruneStage { fn unwind( &mut self, - _provider: &DatabaseProviderRW, + provider: &DatabaseProviderRW, input: UnwindInput, ) -> Result { - info!(target: "sync::stages::prune::unwind", "Stage is always skipped"); + // We cannot recover the data that was pruned in `execute`, so we just update the + // checkpoints. + let prune_checkpoints = provider.get_prune_checkpoints()?; + for (segment, mut checkpoint) in prune_checkpoints { + checkpoint.block_number = Some(input.unwind_to); + provider.save_prune_checkpoint(segment, checkpoint)?; + } Ok(UnwindOutput { checkpoint: StageCheckpoint::new(input.unwind_to) }) } } diff --git a/crates/storage/provider/src/providers/database/mod.rs b/crates/storage/provider/src/providers/database/mod.rs index 3faaf3f1d2..bfd0a39bcb 100644 --- a/crates/storage/provider/src/providers/database/mod.rs +++ b/crates/storage/provider/src/providers/database/mod.rs @@ -582,6 +582,10 @@ impl PruneCheckpointReader for ProviderFactory { ) -> ProviderResult> { self.provider()?.get_prune_checkpoint(segment) } + + fn get_prune_checkpoints(&self) -> ProviderResult> { + self.provider()?.get_prune_checkpoints() + } } impl Clone for ProviderFactory { diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs index afabaf46e5..3aace4f86f 100644 --- a/crates/storage/provider/src/providers/database/provider.rs +++ b/crates/storage/provider/src/providers/database/provider.rs @@ -3327,6 +3327,14 @@ impl PruneCheckpointReader for DatabaseProvider { ) -> ProviderResult> { Ok(self.tx.get::(segment)?) } + + fn get_prune_checkpoints(&self) -> ProviderResult> { + Ok(self + .tx + .cursor_read::()? + .walk(None)? + .collect::>()?) + } } impl PruneCheckpointWriter for DatabaseProvider { diff --git a/crates/storage/provider/src/providers/mod.rs b/crates/storage/provider/src/providers/mod.rs index 4788db247e..4349a9ecc7 100644 --- a/crates/storage/provider/src/providers/mod.rs +++ b/crates/storage/provider/src/providers/mod.rs @@ -581,6 +581,10 @@ where ) -> ProviderResult> { self.database.provider()?.get_prune_checkpoint(segment) } + + fn get_prune_checkpoints(&self) -> ProviderResult> { + self.database.provider()?.get_prune_checkpoints() + } } impl ChainSpecProvider for BlockchainProvider diff --git a/crates/storage/provider/src/test_utils/noop.rs b/crates/storage/provider/src/test_utils/noop.rs index 0364895543..877c4d7afc 100644 --- a/crates/storage/provider/src/test_utils/noop.rs +++ b/crates/storage/provider/src/test_utils/noop.rs @@ -474,6 +474,10 @@ impl PruneCheckpointReader for NoopProvider { ) -> ProviderResult> { Ok(None) } + + fn get_prune_checkpoints(&self) -> ProviderResult> { + Ok(Vec::new()) + } } impl StaticFileProviderFactory for NoopProvider { diff --git a/crates/storage/storage-api/src/prune_checkpoint.rs b/crates/storage/storage-api/src/prune_checkpoint.rs index 6a6e32bf64..ba83ea5103 100644 --- a/crates/storage/storage-api/src/prune_checkpoint.rs +++ b/crates/storage/storage-api/src/prune_checkpoint.rs @@ -4,11 +4,14 @@ use reth_storage_errors::provider::ProviderResult; /// The trait for fetching prune checkpoint related data. #[auto_impl::auto_impl(&, Arc)] pub trait PruneCheckpointReader: Send + Sync { - /// Fetch the checkpoint for the given prune segment. + /// Fetch the prune checkpoint for the given segment. fn get_prune_checkpoint( &self, segment: PruneSegment, ) -> ProviderResult>; + + /// Fetch all the prune checkpoints. + fn get_prune_checkpoints(&self) -> ProviderResult>; } /// The trait for updating prune checkpoint related data.