update

Author: kariy

crates/sync/pipeline/src/lib.rs

@@ -393,7 +393,7 @@ impl<P: StageCheckpointProvider> Pipeline<P> {
     ///
     /// Returns an error if any stage execution fails or if the pipeline fails to read the
     /// checkpoint.
-    pub async fn run_once(&mut self, to: BlockNumber) -> PipelineResult<BlockNumber> {
+    pub async fn execute(&mut self, to: BlockNumber) -> PipelineResult<BlockNumber> {
         if self.stages.is_empty() {
             return Ok(to);
         }
@@ -455,6 +455,39 @@ impl<P: StageCheckpointProvider> Pipeline<P> {
         Ok(last_block_processed_list.into_iter().min().unwrap_or(to))
     }
 
+    /// Runs pruning on all stages.
+    pub async fn prune(&mut self, tip: BlockNumber) -> PipelineResult<()> {
+        if self.stages.is_empty() {
+            return Ok(());
+        }
+
+        for stage in self.stages.iter_mut() {
+            let id = stage.id();
+
+            let span = info_span!(target: "pipeline", "stage.prune", stage = %id, %tip);
+            let enter = span.entered();
+
+            if let Some(checkpoint) = self.provider.checkpoint(id)? {
+                let prune_input = PruneInput::new(checkpoint, self.pruning_config.mode);
+                info!(target: "pipeline", mode = ?self.pruning_config.mode, "Pruning stage.");
+
+                let span_inner = enter.exit();
+                let PruneOutput { pruned_count } = stage
+                    .prune(&prune_input)
+                    .instrument(span_inner.clone())
+                    .await
+                    .map_err(|error| Error::StagePruning { id, error })?;
+
+                let _enter = span_inner.enter();
+                info!(target: "pipeline", %pruned_count, "Stage pruning completed.");
+            } else {
+                info!(target: "pipeline", "Skipping stage - no data to prune (no checkpoint).");
+            }
+        }
+
+        Ok(())
+    }
+
     /// Run the pipeline loop.
     async fn run_loop(&mut self) -> PipelineResult<()> {
         let mut current_chunk_tip = self.chunk_size;
@@ -463,7 +496,7 @@ impl<P: StageCheckpointProvider> Pipeline<P> {
             // Process blocks if we have a tip
             if let Some(tip) = self.tip {
                 let to = current_chunk_tip.min(tip);
-                let last_block_processed = self.run_once(to).await?;
+                let last_block_processed = self.execute(to).await?;
 
                 // Notify subscribers about the newly processed block
                 let _ = self.block_tx.send(Some(last_block_processed));
@@ -514,35 +547,6 @@ impl<P: StageCheckpointProvider> Pipeline<P> {
             None => current_block >= interval,
         }
     }
-
-    /// Runs pruning on all stages.
-    async fn prune(&mut self, tip: BlockNumber) -> PipelineResult<()> {
-        if self.stages.is_empty() {
-            return Ok(());
-        }
-
-        let prune_input = PruneInput::new(tip, self.pruning_config.mode);
-
-        for stage in self.stages.iter_mut() {
-            let id = stage.id();
-            let span = info_span!(target: "pipeline", "stage.prune", stage = %id, %tip);
-            let enter = span.entered();
-
-            info!(target: "pipeline", mode = ?self.pruning_config.mode, "Pruning stage.");
-
-            let span_inner = enter.exit();
-            let PruneOutput { pruned_count } = stage
-                .prune(&prune_input)
-                .instrument(span_inner.clone())
-                .await
-                .map_err(|error| Error::StagePruning { id, error })?;
-
-            let _enter = span_inner.enter();
-            info!(target: "pipeline", %pruned_count, "Stage pruning completed.");
-        }
-
-        Ok(())
-    }
 }
 
 impl<P> IntoFuture for Pipeline<P>

crates/sync/pipeline/tests/pipeline.rs

@@ -8,7 +8,10 @@ use katana_pipeline::Pipeline;
 use katana_primitives::block::BlockNumber;
 use katana_provider::api::stage::StageCheckpointProvider;
 use katana_provider::test_utils::test_provider;
-use katana_stage::{Stage, StageExecutionInput, StageExecutionOutput, StageResult};
+use katana_stage::{
+    PruneInput, PruneOutput, PruneResult, Stage, StageExecutionInput, StageExecutionOutput,
+    StageResult,
+};
 
 /// Simple mock stage that does nothing
 struct MockStage;
@@ -21,6 +24,11 @@ impl Stage for MockStage {
     fn execute<'a>(&'a mut self, input: &'a StageExecutionInput) -> BoxFuture<'a, StageResult> {
         Box::pin(async move { Ok(StageExecutionOutput { last_block_processed: input.to() }) })
     }
+
+    fn prune<'a>(&'a mut self, input: &'a PruneInput) -> BoxFuture<'a, PruneResult> {
+        let _ = input;
+        Box::pin(async move { Ok(PruneOutput::default()) })
+    }
 }
 
 /// Tracks execution calls with their inputs
@@ -146,7 +154,7 @@ async fn run_to_executes_stage_to_target() {
 
     pipeline.add_stage(stage);
     handle.set_tip(5);
-    let result = pipeline.run_once(5).await.unwrap();
+    let result = pipeline.execute(5).await.unwrap();
 
     assert_eq!(result, 5);
     assert_eq!(provider.checkpoint(stage_clone.id()).unwrap(), Some(5));
@@ -170,7 +178,7 @@ async fn run_to_skips_stage_when_checkpoint_equals_target() {
     pipeline.add_stage(stage);
 
     handle.set_tip(5);
-    let result = pipeline.run_once(5).await.unwrap();
+    let result = pipeline.execute(5).await.unwrap();
 
     assert_eq!(result, 5);
     assert_eq!(stage_clone.executions().len(), 0); // Not executed
@@ -189,7 +197,7 @@ async fn run_to_skips_stage_when_checkpoint_exceeds_target() {
     pipeline.add_stage(stage);
 
     handle.set_tip(10);
-    let result = pipeline.run_once(5).await.unwrap();
+    let result = pipeline.execute(5).await.unwrap();
 
     assert_eq!(result, 10); // Returns the checkpoint
     assert_eq!(stage_clone.executions().len(), 0); // Not executed
@@ -207,7 +215,7 @@ async fn run_to_uses_checkpoint_plus_one_as_from() {
     provider.set_checkpoint(stage.id(), 3).unwrap();
     pipeline.add_stage(stage);
     handle.set_tip(10);
-    pipeline.run_once(10).await.unwrap();
+    pipeline.execute(10).await.unwrap();
 
     let execs = stage_clone.executions();
     assert_eq!(execs.len(), 1);
@@ -241,7 +249,7 @@ async fn run_to_executes_all_stages_in_order() {
     ]);
 
     handle.set_tip(5);
-    pipeline.run_once(5).await.unwrap();
+    pipeline.execute(5).await.unwrap();
 
     // All stages should be executed once because the tip is 5 and the chunk size is 10
     assert_eq!(stage1_clone.execution_count(), 1);
@@ -279,7 +287,7 @@ async fn run_to_with_mixed_checkpoints() {
     provider.set_checkpoint(stage2_clone.id(), 3).unwrap();
 
     handle.set_tip(10);
-    pipeline.run_once(10).await.unwrap();
+    pipeline.execute(10).await.unwrap();
 
     // Stage1 should be skipped because its checkpoint (10) >= than the tip (10)
     assert_eq!(stage1_clone.execution_count(), 0);
@@ -317,7 +325,7 @@ async fn run_to_returns_minimum_last_block_processed() {
     ]);
 
     handle.set_tip(20);
-    let result = pipeline.run_once(20).await.unwrap();
+    let result = pipeline.execute(20).await.unwrap();
 
     // make sure that all the stages were executed once
     assert_eq!(stage1_clone.execution_count(), 1);
@@ -353,7 +361,7 @@ async fn run_to_middle_stage_skip_continues() {
     provider.set_checkpoint(stage2_clone.id(), 10).unwrap();
 
     handle.set_tip(10);
-    pipeline.run_once(10).await.unwrap();
+    pipeline.execute(10).await.unwrap();
 
     // Stage1 and Stage3 should execute
     assert_eq!(stage1_clone.execution_count(), 1);
@@ -526,7 +534,7 @@ async fn stage_execution_error_stops_pipeline() {
     pipeline.add_stage(stage);
 
     handle.set_tip(10);
-    let result = pipeline.run_once(10).await;
+    let result = pipeline.execute(10).await;
     assert!(result.is_err());
 
     // Checkpoint should not be set after failure
@@ -549,7 +557,7 @@ async fn stage_error_doesnt_affect_subsequent_runs() {
     pipeline.add_stage(stage2);
 
     handle.set_tip(10);
-    let error = pipeline.run_once(10).await.unwrap_err();
+    let error = pipeline.execute(10).await.unwrap_err();
 
     let katana_pipeline::Error::StageExecution { id, error } = error else {
         panic!("Unexpected error type");
@@ -573,7 +581,7 @@ async fn empty_pipeline_returns_target() {
 
     // No stages added
     handle.set_tip(10);
-    let result = pipeline.run_once(10).await.unwrap();
+    let result = pipeline.execute(10).await.unwrap();
 
     assert_eq!(result, 10);
 }
@@ -591,7 +599,7 @@ async fn tip_equals_checkpoint_no_execution() {
     pipeline.add_stage(stage);
 
     handle.set_tip(10);
-    pipeline.run_once(10).await.unwrap();
+    pipeline.execute(10).await.unwrap();
 
     assert_eq!(executions.lock().unwrap().len(), 0, "Stage1 should not be executed");
 }
@@ -612,7 +620,7 @@ async fn tip_less_than_checkpoint_skip_all() {
     pipeline.add_stage(stage);
 
     handle.set_tip(20);
-    let result = pipeline.run_once(10).await.unwrap();
+    let result = pipeline.execute(10).await.unwrap();
 
     assert_eq!(result, checkpoint);
     assert_eq!(executions.lock().unwrap().len(), 0, "Stage1 should not be executed");
@@ -660,20 +668,20 @@ async fn stage_checkpoint() {
     assert_eq!(initial_checkpoint, None);
 
     handle.set_tip(5);
-    pipeline.run_once(5).await.expect("failed to run the pipeline once");
+    pipeline.execute(5).await.expect("failed to run the pipeline once");
 
     // check that the checkpoint was set
     let actual_checkpoint = provider.checkpoint("Mock").unwrap();
     assert_eq!(actual_checkpoint, Some(5));
 
     handle.set_tip(10);
-    pipeline.run_once(10).await.expect("failed to run the pipeline once");
+    pipeline.execute(10).await.expect("failed to run the pipeline once");
 
     // check that the checkpoint was set
     let actual_checkpoint = provider.checkpoint("Mock").unwrap();
     assert_eq!(actual_checkpoint, Some(10));
 
-    pipeline.run_once(10).await.expect("failed to run the pipeline once");
+    pipeline.execute(10).await.expect("failed to run the pipeline once");
 
     // check that the checkpoint doesn't change
     let actual_checkpoint = provider.checkpoint("Mock").unwrap();

crates/sync/stage/src/blocks/mod.rs

@@ -137,7 +137,8 @@ where
         })
     }
 
-    fn prune<'a>(&'a mut self, _input: &'a PruneInput) -> BoxFuture<'a, PruneResult> {
+    fn prune<'a>(&'a mut self, input: &'a PruneInput) -> BoxFuture<'a, PruneResult> {
+        let _ = input;
         Box::pin(async move {
             // TODO: Implement block pruning once the BlockWriter provider API supports it.
             // For now, this is a no-op. Block pruning would involve: