chore: add engine terminate (paradigmxyz#20420)

Co-authored-by: joshieDo <93316087+joshieDo@users.noreply.github.com>

Author: mattsse

Cargo.lock

@@ -219,10 +219,19 @@ pub enum HandlerEvent<T> {
 }
 
 /// Internal events issued by the [`ChainOrchestrator`].
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 pub enum FromOrchestrator {
     /// Invoked when backfill sync finished
     BackfillSyncFinished(ControlFlow),
     /// Invoked when backfill sync started
     BackfillSyncStarted,
+    /// Gracefully terminate the engine service.
+    ///
+    /// When this variant is received, the engine will persist all remaining in-memory blocks
+    /// to disk before shutting down. Once persistence is complete, a signal is sent through
+    /// the oneshot channel to notify the caller.
+    Terminate {
+        /// Channel to signal termination completion.
+        tx: tokio::sync::oneshot::Sender<()>,
+    },
 }

crates/engine/tree/src/chain.rs

@@ -39,6 +39,7 @@ use revm::state::EvmState;
 use state::TreeState;
 use std::{
     fmt::Debug,
+    ops,
     sync::{
         mpsc::{Receiver, RecvError, RecvTimeoutError, Sender},
         Arc,
@@ -426,9 +427,13 @@ where
             match self.try_recv_engine_message() {
                 Ok(Some(msg)) => {
                     debug!(target: "engine::tree", %msg, "received new engine message");
-                    if let Err(fatal) = self.on_engine_message(msg) {
-                        error!(target: "engine::tree", %fatal, "insert block fatal error");
-                        return
+                    match self.on_engine_message(msg) {
+                        Ok(ops::ControlFlow::Break(())) => return,
+                        Ok(ops::ControlFlow::Continue(())) => {}
+                        Err(fatal) => {
+                            error!(target: "engine::tree", %fatal, "insert block fatal error");
+                            return
+                        }
                     }
                 }
                 Ok(None) => {
@@ -1315,14 +1320,50 @@ where
             if let Some(new_tip_num) = self.find_disk_reorg()? {
                 self.remove_blocks(new_tip_num)
             } else if self.should_persist() {
-                let blocks_to_persist = self.get_canonical_blocks_to_persist()?;
+                let blocks_to_persist =
+                    self.get_canonical_blocks_to_persist(PersistTarget::Threshold)?;
                 self.persist_blocks(blocks_to_persist);
             }
         }
 
         Ok(())
     }
 
+    /// Finishes termination by persisting all remaining blocks and signaling completion.
+    ///
+    /// This blocks until all persistence is complete. Always signals completion,
+    /// even if an error occurs.
+    fn finish_termination(
+        &mut self,
+        pending_termination: oneshot::Sender<()>,
+    ) -> Result<(), AdvancePersistenceError> {
+        trace!(target: "engine::tree", "finishing termination, persisting remaining blocks");
+        let result = self.persist_until_complete();
+        let _ = pending_termination.send(());
+        result
+    }
+
+    /// Persists all remaining blocks until none are left.
+    fn persist_until_complete(&mut self) -> Result<(), AdvancePersistenceError> {
+        loop {
+            // Wait for any in-progress persistence to complete (blocking)
+            if let Some((rx, start_time, _action)) = self.persistence_state.rx.take() {
+                let result = rx.blocking_recv().map_err(|_| TryRecvError::Closed)?;
+                self.on_persistence_complete(result, start_time)?;
+            }
+
+            let blocks_to_persist = self.get_canonical_blocks_to_persist(PersistTarget::Head)?;
+
+            if blocks_to_persist.is_empty() {
+                debug!(target: "engine::tree", "persistence complete, signaling termination");
+                return Ok(())
+            }
+
+            debug!(target: "engine::tree", count = blocks_to_persist.len(), "persisting remaining blocks before shutdown");
+            self.persist_blocks(blocks_to_persist);
+        }
+    }
+
     /// Handles a completed persistence task.
     fn on_persistence_complete(
         &mut self,
@@ -1348,10 +1389,12 @@ where
     }
 
     /// Handles a message from the engine.
+    ///
+    /// Returns `ControlFlow::Break(())` if the engine should terminate.
     fn on_engine_message(
         &mut self,
         msg: FromEngine<EngineApiRequest<T, N>, N::Block>,
-    ) -> Result<(), InsertBlockFatalError> {
+    ) -> Result<ops::ControlFlow<()>, InsertBlockFatalError> {
         match msg {
             FromEngine::Event(event) => match event {
                 FromOrchestrator::BackfillSyncStarted => {
@@ -1361,14 +1404,21 @@ where
                 FromOrchestrator::BackfillSyncFinished(ctrl) => {
                     self.on_backfill_sync_finished(ctrl)?;
                 }
+                FromOrchestrator::Terminate { tx } => {
+                    debug!(target: "engine::tree", "received terminate request");
+                    if let Err(err) = self.finish_termination(tx) {
+                        error!(target: "engine::tree", %err, "Termination failed");
+                    }
+                    return Ok(ops::ControlFlow::Break(()))
+                }
             },
             FromEngine::Request(request) => {
                 match request {
                     EngineApiRequest::InsertExecutedBlock(block) => {
                         let block_num_hash = block.recovered_block().num_hash();
                         if block_num_hash.number <= self.state.tree_state.canonical_block_number() {
                             // outdated block that can be skipped
-                            return Ok(())
+                            return Ok(ops::ControlFlow::Continue(()))
                         }
 
                         debug!(target: "engine::tree", block=?block_num_hash, "inserting already executed block");
@@ -1476,7 +1526,7 @@ where
                 }
             }
         }
-        Ok(())
+        Ok(ops::ControlFlow::Continue(()))
     }
 
     /// Invoked if the backfill sync has finished to target.
@@ -1710,10 +1760,10 @@ where
     }
 
     /// Returns a batch of consecutive canonical blocks to persist in the range
-    /// `(last_persisted_number .. canonical_head - threshold]`. The expected
-    /// order is oldest -> newest.
+    /// `(last_persisted_number .. target]`. The expected order is oldest -> newest.
     fn get_canonical_blocks_to_persist(
         &self,
+        target: PersistTarget,
     ) -> Result<Vec<ExecutedBlock<N>>, AdvancePersistenceError> {
         // We will calculate the state root using the database, so we need to be sure there are no
         // changes
@@ -1724,9 +1774,12 @@ where
         let last_persisted_number = self.persistence_state.last_persisted_block.number;
         let canonical_head_number = self.state.tree_state.canonical_block_number();
 
-        // Persist only up to block buffer target
-        let target_number =
-            canonical_head_number.saturating_sub(self.config.memory_block_buffer_target());
+        let target_number = match target {
+            PersistTarget::Head => canonical_head_number,
+            PersistTarget::Threshold => {
+                canonical_head_number.saturating_sub(self.config.memory_block_buffer_target())
+            }
+        };
 
         debug!(
             target: "engine::tree",
@@ -2869,3 +2922,12 @@ pub enum InsertPayloadOk {
     /// The payload was valid and inserted into the tree.
     Inserted(BlockStatus),
 }
+
+/// Target for block persistence.
+#[derive(Debug, Clone, Copy)]
+enum PersistTarget {
+    /// Persist up to `canonical_head - memory_block_buffer_target`.
+    Threshold,
+    /// Persist all blocks up to and including the canonical head.
+    Head,
+}

crates/engine/tree/src/tree/mod.rs

@@ -4,7 +4,7 @@ use crate::{
     tree::{
         payload_validator::{BasicEngineValidator, TreeCtx, ValidationOutcome},
         persistence_state::CurrentPersistenceAction,
-        TreeConfig,
+        PersistTarget, TreeConfig,
     },
 };
 
@@ -285,7 +285,8 @@ impl TestHarness {
         let fcu_state = self.fcu_state(block_hash);
 
         let (tx, rx) = oneshot::channel();
-        self.tree
+        let _ = self
+            .tree
             .on_engine_message(FromEngine::Request(
                 BeaconEngineMessage::ForkchoiceUpdated {
                     state: fcu_state,
@@ -498,7 +499,7 @@ fn test_tree_persist_block_batch() {
 
     // process the message
     let msg = test_harness.tree.try_recv_engine_message().unwrap().unwrap();
-    test_harness.tree.on_engine_message(msg).unwrap();
+    let _ = test_harness.tree.on_engine_message(msg).unwrap();
 
     // we now should receive the other batch
     let msg = test_harness.tree.try_recv_engine_message().unwrap().unwrap();
@@ -577,7 +578,7 @@ async fn test_engine_request_during_backfill() {
         .with_backfill_state(BackfillSyncState::Active);
 
     let (tx, rx) = oneshot::channel();
-    test_harness
+    let _ = test_harness
         .tree
         .on_engine_message(FromEngine::Request(
             BeaconEngineMessage::ForkchoiceUpdated {
@@ -658,7 +659,7 @@ async fn test_holesky_payload() {
         TestHarness::new(HOLESKY.clone()).with_backfill_state(BackfillSyncState::Active);
 
     let (tx, rx) = oneshot::channel();
-    test_harness
+    let _ = test_harness
         .tree
         .on_engine_message(FromEngine::Request(
             BeaconEngineMessage::NewPayload {
@@ -883,7 +884,8 @@ async fn test_get_canonical_blocks_to_persist() {
         .with_persistence_threshold(persistence_threshold)
         .with_memory_block_buffer_target(memory_block_buffer_target);
 
-    let blocks_to_persist = test_harness.tree.get_canonical_blocks_to_persist().unwrap();
+    let blocks_to_persist =
+        test_harness.tree.get_canonical_blocks_to_persist(PersistTarget::Threshold).unwrap();
 
     let expected_blocks_to_persist_length: usize =
         (canonical_head_number - memory_block_buffer_target - last_persisted_block_number)
@@ -902,7 +904,8 @@ async fn test_get_canonical_blocks_to_persist() {
 
     assert!(test_harness.tree.state.tree_state.sealed_header_by_hash(&fork_block_hash).is_some());
 
-    let blocks_to_persist = test_harness.tree.get_canonical_blocks_to_persist().unwrap();
+    let blocks_to_persist =
+        test_harness.tree.get_canonical_blocks_to_persist(PersistTarget::Threshold).unwrap();
     assert_eq!(blocks_to_persist.len(), expected_blocks_to_persist_length);
 
     // check that the fork block is not included in the blocks to persist
@@ -981,7 +984,7 @@ async fn test_engine_tree_live_sync_transition_required_blocks_requested() {
     let backfill_tip_block = main_chain[(backfill_finished_block_number - 1) as usize].clone();
     // add block to mock provider to enable persistence clean up.
     test_harness.provider.add_block(backfill_tip_block.hash(), backfill_tip_block.into_block());
-    test_harness.tree.on_engine_message(FromEngine::Event(backfill_finished)).unwrap();
+    let _ = test_harness.tree.on_engine_message(FromEngine::Event(backfill_finished)).unwrap();
 
     let event = test_harness.from_tree_rx.recv().await.unwrap();
     match event {
@@ -991,7 +994,7 @@ async fn test_engine_tree_live_sync_transition_required_blocks_requested() {
         _ => panic!("Unexpected event: {event:#?}"),
     }
 
-    test_harness
+    let _ = test_harness
         .tree
         .on_engine_message(FromEngine::DownloadedBlocks(vec![main_chain
             .last()
@@ -1047,7 +1050,7 @@ async fn test_fcu_with_canonical_ancestor_updates_latest_block() {
 
     // Send FCU to the canonical ancestor
     let (tx, rx) = oneshot::channel();
-    test_harness
+    let _ = test_harness
         .tree
         .on_engine_message(FromEngine::Request(
             BeaconEngineMessage::ForkchoiceUpdated {
@@ -1943,4 +1946,53 @@ mod forkchoice_updated_tests {
             .unwrap();
         assert!(result.is_some(), "OpStack should handle canonical head");
     }
+
+    /// Test that engine termination persists all blocks and signals completion.
+    #[test]
+    fn test_engine_termination_with_everything_persisted() {
+        let chain_spec = MAINNET.clone();
+        let mut test_block_builder = TestBlockBuilder::eth().with_chain_spec((*chain_spec).clone());
+
+        // Create 10 blocks to persist
+        let blocks: Vec<_> = test_block_builder.get_executed_blocks(1..11).collect();
+        let canonical_tip = blocks.last().unwrap().recovered_block().number;
+        let test_harness = TestHarness::new(chain_spec).with_blocks(blocks);
+
+        // Create termination channel
+        let (terminate_tx, mut terminate_rx) = oneshot::channel();
+
+        let to_tree_tx = test_harness.to_tree_tx.clone();
+        let action_rx = test_harness.action_rx;
+
+        // Spawn tree in background thread
+        std::thread::Builder::new()
+            .name("Engine Task".to_string())
+            .spawn(|| test_harness.tree.run())
+            .unwrap();
+
+        // Send terminate request
+        to_tree_tx
+            .send(FromEngine::Event(FromOrchestrator::Terminate { tx: terminate_tx }))
+            .unwrap();
+
+        // Handle persistence actions until termination completes
+        let mut last_persisted_number = 0;
+        loop {
+            if terminate_rx.try_recv().is_ok() {
+                break;
+            }
+
+            if let Ok(PersistenceAction::SaveBlocks(saved_blocks, sender)) =
+                action_rx.recv_timeout(std::time::Duration::from_millis(100))
+            {
+                if let Some(last) = saved_blocks.last() {
+                    last_persisted_number = last.recovered_block().number;
+                }
+                sender.send(saved_blocks.last().map(|b| b.recovered_block().num_hash())).unwrap();
+            }
+        }
+
+        // Ensure we persisted right to the tip
+        assert_eq!(last_persisted_number, canonical_tip);
+    }
 }