From 1994959fb28c9f9a17128ca27b6f550b2d8df1d3 Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Tue, 24 Sep 2024 19:15:38 -0400
Subject: [PATCH] chore(tree): move persistence_state out of mod.rs (#11180)

---
 crates/engine/tree/src/tree/mod.rs            | 45 +---------------
 .../engine/tree/src/tree/persistence_state.rs | 52 +++++++++++++++++++
 2 files changed, 54 insertions(+), 43 deletions(-)
 create mode 100644 crates/engine/tree/src/tree/persistence_state.rs

diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
index 750e31a866..3b7cf8ae2b 100644
--- a/crates/engine/tree/src/tree/mod.rs
+++ b/crates/engine/tree/src/tree/mod.rs
@@ -65,9 +65,11 @@ use tracing::*;
 pub mod config;
 mod invalid_block_hook;
 mod metrics;
+mod persistence_state;
 use crate::{engine::EngineApiRequest, tree::metrics::EngineApiMetrics};
 pub use config::TreeConfig;
 pub use invalid_block_hook::{InvalidBlockHooks, NoopInvalidBlockHook};
+pub use persistence_state::PersistenceState;
 pub use reth_engine_primitives::InvalidBlockHook;
 
 /// Keeps track of the state of the tree.
@@ -2540,49 +2542,6 @@ pub enum AdvancePersistenceError {
     Provider(#[from] ProviderError),
 }
 
-/// The state of the persistence task.
-#[derive(Default, Debug)]
-pub struct PersistenceState {
-    /// Hash and number of the last block persisted.
-    ///
-    /// This tracks the chain height that is persisted on disk
-    last_persisted_block: BlockNumHash,
-    /// Receiver end of channel where the result of the persistence task will be
-    /// sent when done. A None value means there's no persistence task in progress.
-    rx: Option<(oneshot::Receiver<Option<BlockNumHash>>, Instant)>,
-    /// The block above which blocks should be removed from disk, because there has been an on disk
-    /// reorg.
-    remove_above_state: VecDeque<u64>,
-}
-
-impl PersistenceState {
-    /// Determines if there is a persistence task in progress by checking if the
-    /// receiver is set.
-    const fn in_progress(&self) -> bool {
-        self.rx.is_some()
-    }
-
-    /// Sets state for a started persistence task.
-    fn start(&mut self, rx: oneshot::Receiver<Option<BlockNumHash>>) {
-        self.rx = Some((rx, Instant::now()));
-    }
-
-    /// Sets the `remove_above_state`, to the new tip number specified, only if it is less than the
-    /// current `last_persisted_block_number`.
-    fn schedule_removal(&mut self, new_tip_num: u64) {
-        debug!(target: "engine::tree", ?new_tip_num, prev_remove_state=?self.remove_above_state, last_persisted_block=?self.last_persisted_block, "Scheduling removal");
-        self.remove_above_state.push_back(new_tip_num);
-    }
-
-    /// Sets state for a finished persistence task.
-    fn finish(&mut self, last_persisted_block_hash: B256, last_persisted_block_number: u64) {
-        trace!(target: "engine::tree", block= %last_persisted_block_number, hash=%last_persisted_block_hash, "updating persistence state");
-        self.rx = None;
-        self.last_persisted_block =
-            BlockNumHash::new(last_persisted_block_number, last_persisted_block_hash);
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/crates/engine/tree/src/tree/persistence_state.rs b/crates/engine/tree/src/tree/persistence_state.rs
new file mode 100644
index 0000000000..fca51291bb
--- /dev/null
+++ b/crates/engine/tree/src/tree/persistence_state.rs
@@ -0,0 +1,52 @@
+use alloy_eips::BlockNumHash;
+use reth_primitives::B256;
+use std::{collections::VecDeque, time::Instant};
+use tokio::sync::oneshot;
+use tracing::{debug, trace};
+
+/// The state of the persistence task.
+#[derive(Default, Debug)]
+pub struct PersistenceState {
+    /// Hash and number of the last block persisted.
+    ///
+    /// This tracks the chain height that is persisted on disk
+    pub(crate) last_persisted_block: BlockNumHash,
+    /// Receiver end of channel where the result of the persistence task will be
+    /// sent when done. A None value means there's no persistence task in progress.
+    pub(crate) rx: Option<(oneshot::Receiver<Option<BlockNumHash>>, Instant)>,
+    /// The block above which blocks should be removed from disk, because there has been an on disk
+    /// reorg.
+    pub(crate) remove_above_state: VecDeque<u64>,
+}
+
+impl PersistenceState {
+    /// Determines if there is a persistence task in progress by checking if the
+    /// receiver is set.
+    pub(crate) const fn in_progress(&self) -> bool {
+        self.rx.is_some()
+    }
+
+    /// Sets state for a started persistence task.
+    pub(crate) fn start(&mut self, rx: oneshot::Receiver<Option<BlockNumHash>>) {
+        self.rx = Some((rx, Instant::now()));
+    }
+
+    /// Sets the `remove_above_state`, to the new tip number specified, only if it is less than the
+    /// current `last_persisted_block_number`.
+    pub(crate) fn schedule_removal(&mut self, new_tip_num: u64) {
+        debug!(target: "engine::tree", ?new_tip_num, prev_remove_state=?self.remove_above_state, last_persisted_block=?self.last_persisted_block, "Scheduling removal");
+        self.remove_above_state.push_back(new_tip_num);
+    }
+
+    /// Sets state for a finished persistence task.
+    pub(crate) fn finish(
+        &mut self,
+        last_persisted_block_hash: B256,
+        last_persisted_block_number: u64,
+    ) {
+        trace!(target: "engine::tree", block= %last_persisted_block_number, hash=%last_persisted_block_hash, "updating persistence state");
+        self.rx = None;
+        self.last_persisted_block =
+            BlockNumHash::new(last_persisted_block_number, last_persisted_block_hash);
+    }
+}