From 5f1661610dae88065f46db991f84357c67ef8ff8 Mon Sep 17 00:00:00 2001
From: Federico Gimenez <federico.gimenez@gmail.com>
Date: Tue, 6 Aug 2024 14:04:35 +0200
Subject: [PATCH 1/2] chore: ensure persistence receiver is dropped

---
 crates/engine/tree/src/tree/mod.rs | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
index 6cfa9365b111..db19abb2e91f 100644
--- a/crates/engine/tree/src/tree/mod.rs
+++ b/crates/engine/tree/src/tree/mod.rs
@@ -535,7 +535,10 @@ where
                 }
             }
 
-            self.advance_persistence();
+            if let Err(err) = self.advance_persistence() {
+                error!(target: "engine", %err, "Advancing persistence failed");
+                break
+            }
         }
     }
 
@@ -568,7 +571,7 @@ where
     ///
     /// If we're currently awaiting a response this will try to receive the response (non-blocking)
     /// or send a new persistence action if necessary.
-    fn advance_persistence(&mut self) {
+    fn advance_persistence(&mut self) -> Result<(), TryRecvError> {
         if self.should_persist() && !self.persistence_state.in_progress() {
             let blocks_to_persist = self.get_canonical_blocks_to_persist();
             if !blocks_to_persist.is_empty() {
@@ -581,10 +584,10 @@ where
         }
 
         if self.persistence_state.in_progress() {
-            let rx = self
+            let mut rx = self
                 .persistence_state
                 .rx
-                .as_mut()
+                .take()
                 .expect("if a persistence task is in progress Receiver must be Some");
 
             // Check if persistence has completed
@@ -594,7 +597,7 @@ where
                         // if this happened, then we persisted no blocks because we sent an empty
                         // vec of blocks
                         warn!(target: "engine", "Persistence task completed but did not persist any blocks");
-                        return
+                        return Ok(())
                     };
                     if let Some(block) =
                         self.state.tree_state.block_by_hash(last_persisted_block_hash)
@@ -605,12 +608,10 @@ where
                         error!("could not find persisted block with hash {last_persisted_block_hash} in memory");
                     }
                 }
-                Err(TryRecvError::Closed) => {
-                    panic!("persistence task has been terminated");
-                }
-                _ => {}
+                Err(err) => return Err(err),
             }
         }
+        Ok(())
     }
 
     /// Handles a message from the engine.

From dde3f1c13570ed2a4fb5f0c6b92581f341853561 Mon Sep 17 00:00:00 2001
From: Federico Gimenez <federico.gimenez@gmail.com>
Date: Tue, 6 Aug 2024 14:28:27 +0200
Subject: [PATCH 2/2] reassing persistence_state.rx on TryRecvError::Empty

---
 crates/engine/tree/src/tree/mod.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
index db19abb2e91f..97ae3658fa2f 100644
--- a/crates/engine/tree/src/tree/mod.rs
+++ b/crates/engine/tree/src/tree/mod.rs
@@ -608,7 +608,8 @@ where
                         error!("could not find persisted block with hash {last_persisted_block_hash} in memory");
                     }
                 }
-                Err(err) => return Err(err),
+                Err(TryRecvError::Closed) => return Err(TryRecvError::Closed),
+                Err(TryRecvError::Empty) => self.persistence_state.rx = Some(rx),
             }
         }
         Ok(())