From 77b7d778195702e2a55bfa425e0833ab3a89bbbb Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <a.shekhirin@gmail.com>
Date: Thu, 3 Aug 2023 17:00:56 +0100
Subject: [PATCH] fix(engine): poll prune first (#4048)

---
 crates/consensus/beacon/src/engine/mod.rs | 44 +++++++++++++++--------
 1 file changed, 29 insertions(+), 15 deletions(-)

diff --git a/crates/consensus/beacon/src/engine/mod.rs b/crates/consensus/beacon/src/engine/mod.rs
index cd64817792..f5025e049e 100644
--- a/crates/consensus/beacon/src/engine/mod.rs
+++ b/crates/consensus/beacon/src/engine/mod.rs
@@ -1635,6 +1635,21 @@ where
     fn is_prune_active(&self) -> bool {
         !self.is_prune_idle()
     }
+
+    /// Polls the prune controller, if it exists, and processes the event [`EnginePruneEvent`]
+    /// emitted by it.
+    ///
+    /// Returns [`Option::Some`] if prune controller emitted an event which resulted in the error
+    /// (see [`Self::on_prune_event`] for error handling)
+    fn poll_prune(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Option<Result<(), BeaconConsensusEngineError>> {
+        match self.prune.as_mut()?.poll(cx, self.blockchain.canonical_tip().number) {
+            Poll::Ready(prune_event) => self.on_prune_event(prune_event),
+            Poll::Pending => None,
+        }
+    }
 }
 
 /// On initialization, the consensus engine will poll the message receiver and return
@@ -1665,6 +1680,14 @@ where
         // Process all incoming messages from the CL, these can affect the state of the
         // SyncController, hence they are polled first, and they're also time sensitive.
         loop {
+            // Poll prune controller first if it's active, as we will not be able to process any
+            // engine messages until it's finished.
+            if this.is_prune_active() {
+                if let Some(res) = this.poll_prune(cx) {
+                    return Poll::Ready(res)
+                }
+            }
+
             let mut engine_messages_pending = false;
             let mut sync_pending = false;
 
@@ -1715,23 +1738,14 @@ where
 
             // Poll prune controller if all conditions are met:
             // 1. Pipeline is idle
-            // 2. Either of two:
-            //  1. Pruning is running and we need to prioritize checking its events
-            //  2. Both engine and sync messages are pending AND latest FCU status is not INVALID,
-            //     so we may start pruning
+            // 2. No engine and sync messages are pending
+            // 3. Latest FCU status is not INVALID
             if this.sync.is_pipeline_idle() &&
-                (this.is_prune_active() ||
-                    is_pending && !this.forkchoice_state_tracker.is_latest_invalid())
+                is_pending &&
+                !this.forkchoice_state_tracker.is_latest_invalid()
             {
-                if let Some(ref mut prune) = this.prune {
-                    match prune.poll(cx, this.blockchain.canonical_tip().number) {
-                        Poll::Ready(prune_event) => {
-                            if let Some(res) = this.on_prune_event(prune_event) {
-                                return Poll::Ready(res)
-                            }
-                        }
-                        Poll::Pending => {}
-                    }
+                if let Some(res) = this.poll_prune(cx) {
+                    return Poll::Ready(res)
                 }
             }