From dfcca695df21c8e3833cef7cb4fc48bef3592e45 Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Wed, 24 May 2023 15:52:01 -0400
Subject: [PATCH] fix: update rpc on continuous sync (#2707)

---
 crates/consensus/beacon/src/engine/mod.rs  | 21 +++++++++++++++++++++
 crates/consensus/beacon/src/engine/sync.rs |  5 +++++
 crates/stages/src/pipeline/ctrl.rs         |  9 +++++++++
 3 files changed, 35 insertions(+)

diff --git a/crates/consensus/beacon/src/engine/mod.rs b/crates/consensus/beacon/src/engine/mod.rs
index 075648ad38..f39df89766 100644
--- a/crates/consensus/beacon/src/engine/mod.rs
+++ b/crates/consensus/beacon/src/engine/mod.rs
@@ -868,6 +868,27 @@ where
                             return None
                         }
 
+                        // update the canon chain if continuous is enabled
+                        if self.sync.run_pipeline_continuously() {
+                            let max_block = ctrl.progress().unwrap_or_default();
+                            let max_header = match self.blockchain.sealed_header(max_block) {
+                                Ok(header) => match header {
+                                    Some(header) => header,
+                                    None => {
+                                        return Some(Err(Error::Provider(
+                                            ProviderError::HeaderNotFound(max_block.into()),
+                                        )
+                                        .into()))
+                                    }
+                                },
+                                Err(error) => {
+                                    error!(target: "consensus::engine", ?error, "Error getting canonical header for continuous sync");
+                                    return Some(Err(error.into()))
+                                }
+                            };
+                            self.blockchain.set_canonical_head(max_header);
+                        }
+
                         // Update the state and hashes of the blockchain tree if possible.
                         match self.restore_tree_if_possible(current_state) {
                             Ok(_) => self.sync_state_updater.update_sync_state(SyncState::Idle),
diff --git a/crates/consensus/beacon/src/engine/sync.rs b/crates/consensus/beacon/src/engine/sync.rs
index 86eb233100..2e22f05d68 100644
--- a/crates/consensus/beacon/src/engine/sync.rs
+++ b/crates/consensus/beacon/src/engine/sync.rs
@@ -90,6 +90,11 @@ where
         self.inflight_full_block_requests.retain(|req| *req.hash() != hash);
     }
 
+    /// Returns whether or not the sync controller is set to run the pipeline continuously.
+    pub(crate) fn run_pipeline_continuously(&self) -> bool {
+        self.run_pipeline_continuously
+    }
+
     /// Returns `true` if the pipeline is idle.
     pub(crate) fn is_pipeline_idle(&self) -> bool {
         self.pipeline_state.is_idle()
diff --git a/crates/stages/src/pipeline/ctrl.rs b/crates/stages/src/pipeline/ctrl.rs
index 05f84c2b07..9814d10991 100644
--- a/crates/stages/src/pipeline/ctrl.rs
+++ b/crates/stages/src/pipeline/ctrl.rs
@@ -32,4 +32,13 @@ impl ControlFlow {
     pub fn is_unwind(&self) -> bool {
         matches!(self, ControlFlow::Unwind { .. })
     }
+
+    /// Returns the pipeline progress, if the state is not `Unwind`.
+    pub fn progress(&self) -> Option<BlockNumber> {
+        match self {
+            ControlFlow::Unwind { .. } => None,
+            ControlFlow::Continue { progress } => Some(*progress),
+            ControlFlow::NoProgress { stage_progress } => *stage_progress,
+        }
+    }
 }