From de153516646e7e82dab09a0f7dbff65c509477da Mon Sep 17 00:00:00 2001
From: Federico Gimenez <federico.gimenez@gmail.com>
Date: Fri, 9 Aug 2024 14:16:08 +0000
Subject: [PATCH] feat: add setting for how many blocks can be persisted

---
 crates/engine/tree/src/tree/config.rs | 24 +++++++++++++++++++++++-
 crates/engine/tree/src/tree/mod.rs    | 13 ++++++++-----
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/crates/engine/tree/src/tree/config.rs b/crates/engine/tree/src/tree/config.rs
index 702bef0fade0f..16e93be674a8b 100644
--- a/crates/engine/tree/src/tree/config.rs
+++ b/crates/engine/tree/src/tree/config.rs
@@ -1,6 +1,7 @@
 //! Engine tree configuration.
 
 const DEFAULT_PERSISTENCE_THRESHOLD: u64 = 3;
+const DEFAUL_MEMORY_BLOCK_BUFFER_TARGET: u64 = 2;
 const DEFAULT_BLOCK_BUFFER_LIMIT: u32 = 256;
 const DEFAULT_MAX_INVALID_HEADER_CACHE_LENGTH: u32 = 256;
 
@@ -9,8 +10,12 @@ const DEFAULT_MAX_EXECUTE_BLOCK_BATCH_SIZE: usize = 4;
 /// The configuration of the engine tree.
 #[derive(Debug)]
 pub struct TreeConfig {
-    /// Maximum number of blocks to be kept only in memory without triggering persistence.
+    /// Maximum number of blocks to be kept only in memory without triggering
+    /// persistence.
     persistence_threshold: u64,
+    /// How close to the canonical head we persist blocks. Represents the ideal
+    /// number of most recent blocks to keep in memory for quick access and reorgs.
+    memory_block_buffer_target: u64,
     /// Number of pending blocks that cannot be executed due to missing parent and
     /// are kept in cache.
     block_buffer_limit: u32,
@@ -24,6 +29,7 @@ impl Default for TreeConfig {
     fn default() -> Self {
         Self {
             persistence_threshold: DEFAULT_PERSISTENCE_THRESHOLD,
+            memory_block_buffer_target: DEFAUL_MEMORY_BLOCK_BUFFER_TARGET,
             block_buffer_limit: DEFAULT_BLOCK_BUFFER_LIMIT,
             max_invalid_header_cache_length: DEFAULT_MAX_INVALID_HEADER_CACHE_LENGTH,
             max_execute_block_batch_size: DEFAULT_MAX_EXECUTE_BLOCK_BATCH_SIZE,
@@ -35,12 +41,14 @@ impl TreeConfig {
     /// Create engine tree configuration.
     pub const fn new(
         persistence_threshold: u64,
+        memory_block_buffer_target: u64,
         block_buffer_limit: u32,
         max_invalid_header_cache_length: u32,
         max_execute_block_batch_size: usize,
     ) -> Self {
         Self {
             persistence_threshold,
+            memory_block_buffer_target,
             block_buffer_limit,
             max_invalid_header_cache_length,
             max_execute_block_batch_size,
@@ -52,6 +60,11 @@ impl TreeConfig {
         self.persistence_threshold
     }
 
+    /// Return the memory block buffer target.
+    pub const fn memory_block_buffer_target(&self) -> u64 {
+        self.memory_block_buffer_target
+    }
+
     /// Return the block buffer limit.
     pub const fn block_buffer_limit(&self) -> u32 {
         self.block_buffer_limit
@@ -73,6 +86,15 @@ impl TreeConfig {
         self
     }
 
+    /// Setter for memory block buffer target.
+    pub const fn with_memory_block_buffer_target(
+        mut self,
+        memory_block_buffer_target: u64,
+    ) -> Self {
+        self.memory_block_buffer_target = memory_block_buffer_target;
+        self
+    }
+
     /// Setter for block buffer limit.
     pub const fn with_block_buffer_limit(mut self, block_buffer_limit: u32) -> Self {
         self.block_buffer_limit = block_buffer_limit;
diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
index 9e9a05a8c98c4..f286bade375df 100644
--- a/crates/engine/tree/src/tree/mod.rs
+++ b/crates/engine/tree/src/tree/mod.rs
@@ -1124,7 +1124,7 @@ where
         let canonical_head_number = self.state.tree_state.canonical_block_number();
 
         let target_number =
-            canonical_head_number.saturating_sub(self.config.persistence_threshold());
+            canonical_head_number.saturating_sub(self.config.memory_block_buffer_target());
 
         while let Some(block) = self.state.tree_state.blocks_by_hash.get(&current_hash) {
             if block.block.number <= last_persisted_number {
@@ -2354,7 +2354,8 @@ mod tests {
         if let PersistenceAction::SaveBlocks(saved_blocks, _) = received_action {
             // only blocks.len() - tree_config.persistence_threshold() will be
             // persisted
-            let expected_persist_len = blocks.len() - tree_config.persistence_threshold() as usize;
+            let expected_persist_len =
+                blocks.len() - tree_config.memory_block_buffer_target() as usize;
             assert_eq!(saved_blocks.len(), expected_persist_len);
             assert_eq!(saved_blocks, blocks[..expected_persist_len]);
         } else {
@@ -2648,13 +2649,15 @@ mod tests {
             last_persisted_block_number;
 
         let persistence_threshold = 4;
-        test_harness.tree.config =
-            TreeConfig::default().with_persistence_threshold(persistence_threshold);
+        let memory_block_buffer_target = 3;
+        test_harness.tree.config = TreeConfig::default()
+            .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();
 
         let expected_blocks_to_persist_length: usize =
-            (canonical_head_number - persistence_threshold - last_persisted_block_number)
+            (canonical_head_number - memory_block_buffer_target - last_persisted_block_number)
                 .try_into()
                 .unwrap();