From 4ce5a60c5876986647fe5040a82036cc3ce6d2fb Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Thu, 30 Jan 2025 19:48:35 -0500
Subject: [PATCH] feat(cli): add use-caching-and-prewarming flag (#14114)

---
 book/cli/reth/node.md                 |  3 +++
 crates/engine/tree/src/tree/config.rs | 17 +++++++++++++++++
 crates/node/core/src/args/engine.rs   |  5 +++++
 3 files changed, 25 insertions(+)

diff --git a/book/cli/reth/node.md b/book/cli/reth/node.md
index 210a3506e..b813c07e3 100644
--- a/book/cli/reth/node.md
+++ b/book/cli/reth/node.md
@@ -695,6 +695,9 @@ Engine:
       --engine.state-root-task
           Enable state root task
 
+      --engine.caching-and-prewarming
+          Enable cross-block caching and parallel prewarming
+
       --engine.state-root-task-compare-updates
           Enable comparing trie updates from the state root task to the trie updates from the regular state root calculation
 
diff --git a/crates/engine/tree/src/tree/config.rs b/crates/engine/tree/src/tree/config.rs
index cec8bb4e8..59e7b87db 100644
--- a/crates/engine/tree/src/tree/config.rs
+++ b/crates/engine/tree/src/tree/config.rs
@@ -46,6 +46,8 @@ pub struct TreeConfig {
     /// Whether to always compare trie updates from the state root task to the trie updates from
     /// the regular state root calculation.
     always_compare_trie_updates: bool,
+    /// Whether to use cross-block caching and parallel prewarming
+    use_caching_and_prewarming: bool,
 }
 
 impl Default for TreeConfig {
@@ -58,12 +60,14 @@ impl Default for TreeConfig {
             max_execute_block_batch_size: DEFAULT_MAX_EXECUTE_BLOCK_BATCH_SIZE,
             use_state_root_task: false,
             always_compare_trie_updates: false,
+            use_caching_and_prewarming: false,
         }
     }
 }
 
 impl TreeConfig {
     /// Create engine tree configuration.
+    #[allow(clippy::too_many_arguments)]
     pub const fn new(
         persistence_threshold: u64,
         memory_block_buffer_target: u64,
@@ -72,6 +76,7 @@ impl TreeConfig {
         max_execute_block_batch_size: usize,
         use_state_root_task: bool,
         always_compare_trie_updates: bool,
+        use_caching_and_prewarming: bool,
     ) -> Self {
         Self {
             persistence_threshold,
@@ -81,6 +86,7 @@ impl TreeConfig {
             max_execute_block_batch_size,
             use_state_root_task,
             always_compare_trie_updates,
+            use_caching_and_prewarming,
         }
     }
 
@@ -114,6 +120,11 @@ impl TreeConfig {
         self.use_state_root_task
     }
 
+    /// Returns whether or not cross-block caching and parallel prewarming should be used.
+    pub const fn use_caching_and_prewarming(&self) -> bool {
+        self.use_caching_and_prewarming
+    }
+
     /// Returns whether to always compare trie updates from the state root task to the trie updates
     /// from the regular state root calculation.
     pub const fn always_compare_trie_updates(&self) -> bool {
@@ -165,6 +176,12 @@ impl TreeConfig {
         self
     }
 
+    /// Setter for whether to use the new state root task calculation method.
+    pub const fn with_caching_and_prewarming(mut self, use_caching_and_prewarming: bool) -> Self {
+        self.use_caching_and_prewarming = use_caching_and_prewarming;
+        self
+    }
+
     /// Setter for whether to always compare trie updates from the state root task to the trie
     /// updates from the regular state root calculation.
     pub const fn with_always_compare_trie_updates(
diff --git a/crates/node/core/src/args/engine.rs b/crates/node/core/src/args/engine.rs
index 411074d28..8de89b3a6 100644
--- a/crates/node/core/src/args/engine.rs
+++ b/crates/node/core/src/args/engine.rs
@@ -20,6 +20,10 @@ pub struct EngineArgs {
     #[arg(long = "engine.state-root-task")]
     pub state_root_task_enabled: bool,
 
+    /// Enable cross-block caching and parallel prewarming
+    #[arg(long = "engine.caching-and-prewarming")]
+    pub caching_and_prewarming_enabled: bool,
+
     /// Enable comparing trie updates from the state root task to the trie updates from the regular
     /// state root calculation.
     #[arg(long = "engine.state-root-task-compare-updates")]
@@ -33,6 +37,7 @@ impl Default for EngineArgs {
             memory_block_buffer_target: DEFAULT_MEMORY_BLOCK_BUFFER_TARGET,
             state_root_task_enabled: false,
             state_root_task_compare_updates: false,
+            caching_and_prewarming_enabled: false,
         }
     }
 }