fix(pipeline): don't restart run_loop future on tip change (#368)

Currently the `select!` is configured in `Pipeline::run` to completes on the command waiting regardless of what command is received. This means the `self.run_loop()` future will be cancelled even if it's a command for changing the pipeline's tip - the pipeline has to process the same batch again.

Author: kariy

crates/sync/pipeline/src/lib.rs

@@ -81,7 +81,7 @@ use katana_primitives::block::BlockNumber;
 use katana_provider_api::stage::StageCheckpointProvider;
 use katana_provider_api::ProviderError;
 use katana_stage::{Stage, StageExecutionInput, StageExecutionOutput};
-use tokio::sync::watch;
+use tokio::sync::watch::{self};
 use tokio::task::yield_now;
 use tracing::{debug, error, info, info_span, Instrument};
 
@@ -101,6 +101,9 @@ pub enum Error {
 
     #[error(transparent)]
     Provider(#[from] ProviderError),
+
+    #[error("command channel closed")]
+    CommandChannelClosed,
 }
 
 /// Commands that can be sent to control the pipeline.
@@ -217,8 +220,8 @@ pub struct Pipeline<P> {
     chunk_size: u64,
     provider: P,
     stages: Vec<Box<dyn Stage>>,
-    command_rx: watch::Receiver<Option<PipelineCommand>>,
-    command_tx: watch::Sender<Option<PipelineCommand>>,
+    cmd_rx: watch::Receiver<Option<PipelineCommand>>,
+    cmd_tx: watch::Sender<Option<PipelineCommand>>,
     block_tx: watch::Sender<Option<BlockNumber>>,
     tip: Option<BlockNumber>,
 }
@@ -240,8 +243,8 @@ impl<P> Pipeline<P> {
         let handle = PipelineHandle { tx: tx.clone(), block_tx: block_tx.clone() };
         let pipeline = Self {
             stages: Vec::new(),
-            command_rx: rx,
-            command_tx: tx,
+            cmd_rx: rx,
+            cmd_tx: tx,
             block_tx,
             provider,
             chunk_size,
@@ -269,7 +272,7 @@ impl<P> Pipeline<P> {
     /// The handle can be used to set the target tip block for the pipeline to sync to or to
     /// stop the pipeline.
     pub fn handle(&self) -> PipelineHandle {
-        PipelineHandle { tx: self.command_tx.clone(), block_tx: self.block_tx.clone() }
+        PipelineHandle { tx: self.cmd_tx.clone(), block_tx: self.block_tx.clone() }
     }
 }
 
@@ -285,29 +288,19 @@ impl<P: StageCheckpointProvider> Pipeline<P> {
     /// Returns an error if any stage execution fails or it an error occurs while reading the
     /// checkpoint.
     pub async fn run(&mut self) -> PipelineResult<()> {
-        let mut command_rx = self.command_rx.clone();
+        let mut command_rx = self.cmd_rx.clone();
 
         loop {
             tokio::select! {
                 biased;
 
-                changed = command_rx.changed() => {
+                changed = command_rx.wait_for(|c| matches!(c, &Some(PipelineCommand::Stop))) => {
                     if changed.is_err() {
                         break;
                     }
 
-                    // Check if the handle has sent a signal
-                    match *self.command_rx.borrow_and_update() {
-                        Some(PipelineCommand::Stop) => {
-                            debug!(target: "pipeline", "Received stop command.");
-                            break;
-                        }
-                        Some(PipelineCommand::SetTip(new_tip)) => {
-                            info!(target: "pipeline", tip = %new_tip, "A new tip has been set.");
-                            self.tip = Some(new_tip);
-                        }
-                        None => {}
-                    }
+                    debug!(target: "pipeline", "Received stop command.");
+                    break;
                 }
 
                 result = self.run_loop() => {
@@ -425,27 +418,26 @@ impl<P: StageCheckpointProvider> Pipeline<P> {
                 } else {
                     current_chunk_tip = (last_block_processed + self.chunk_size).min(tip);
                 }
-
-                continue;
+            } else {
+                info!(target: "pipeline", "Waiting to receive new tip.");
             }
 
-            info!(target: "pipeline", "Waiting to receive new tip.");
-
-            // block until a new tip is set
-            self.command_rx
+            if let Some(PipelineCommand::SetTip(new_tip)) = *self
+                .cmd_rx
                 .wait_for(|c| matches!(c, &Some(PipelineCommand::SetTip(_))))
                 .await
-                .expect("qed; channel closed");
+                .map_err(|_| Error::CommandChannelClosed)?
+            {
+                info!(target: "pipeline", tip = %new_tip, "A new tip has been set.");
+                self.tip = Some(new_tip);
+            }
 
             yield_now().await;
         }
     }
 }
 
-impl<P> IntoFuture for Pipeline<P>
-where
-    P: StageCheckpointProvider + 'static,
-{
+impl<P: StageCheckpointProvider + 'static> IntoFuture for Pipeline<P> {
     type Output = PipelineResult<()>;
     type IntoFuture = PipelineFut;
 
@@ -458,13 +450,10 @@ where
     }
 }
 
-impl<P> core::fmt::Debug for Pipeline<P>
-where
-    P: core::fmt::Debug,
-{
+impl<P: core::fmt::Debug> core::fmt::Debug for Pipeline<P> {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.debug_struct("Pipeline")
-            .field("command", &self.command_rx)
+            .field("command", &self.cmd_rx)
             .field("provider", &self.provider)
             .field("chunk_size", &self.chunk_size)
             .field("stages", &self.stages.iter().map(|s| s.id()).collect::<Vec<_>>())