From 74d7fe3075aff34e63ce64ca24b7188fc7c7c854 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A9a=20Narzis?=
 <78718413+lean-apple@users.noreply.github.com>
Date: Fri, 8 Nov 2024 21:50:48 +0700
Subject: [PATCH] feat(rpc): optimize block validation with state caching
 (#12299)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
---
 crates/rpc/rpc/src/validation.rs | 50 +++++++++++++++++++++++++++-----
 1 file changed, 43 insertions(+), 7 deletions(-)

diff --git a/crates/rpc/rpc/src/validation.rs b/crates/rpc/rpc/src/validation.rs
index 919fe2d85..c3f2aab70 100644
--- a/crates/rpc/rpc/src/validation.rs
+++ b/crates/rpc/rpc/src/validation.rs
@@ -20,7 +20,7 @@ use reth_provider::{
     AccountReader, BlockExecutionInput, BlockExecutionOutput, BlockReaderIdExt, HeaderProvider,
     StateProviderFactory, WithdrawalsProvider,
 };
-use reth_revm::database::StateProviderDatabase;
+use reth_revm::{cached::CachedReads, database::StateProviderDatabase};
 use reth_rpc_api::BlockSubmissionValidationApiServer;
 use reth_rpc_eth_types::EthApiError;
 use reth_rpc_server_types::{result::internal_rpc_err, ToRpcResult};
@@ -28,6 +28,7 @@ use reth_trie::HashedPostState;
 use revm_primitives::{Address, B256, U256};
 use serde::{Deserialize, Serialize};
 use std::{collections::HashSet, sync::Arc};
+use tokio::sync::RwLock;
 
 /// Configuration for validation API.
 #[derive(Debug, Clone, Default, Eq, PartialEq, Serialize, Deserialize)]
@@ -64,7 +65,7 @@ pub enum ValidationApiError {
     Execution(#[from] BlockExecutionError),
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct ValidationApiInner<Provider: ChainSpecProvider, E> {
     /// The provider that can interact with the chain.
     provider: Provider,
@@ -76,10 +77,15 @@ pub struct ValidationApiInner<Provider: ChainSpecProvider, E> {
     executor_provider: E,
     /// Set of disallowed addresses
     disallow: HashSet<Address>,
+    /// Cached state reads to avoid redundant disk I/O across multiple validation attempts
+    /// targeting the same state. Stores a tuple of (`block_hash`, `cached_reads`) for the
+    /// latest head block state. Uses async `RwLock` to safely handle concurrent validation
+    /// requests.
+    cached_state: RwLock<(B256, CachedReads)>,
 }
 
 /// The type that implements the `validation` rpc namespace trait
-#[derive(Clone, Debug, derive_more::Deref)]
+#[derive(Debug, derive_more::Deref)]
 pub struct ValidationApi<Provider: ChainSpecProvider, E> {
     #[deref]
     inner: Arc<ValidationApiInner<Provider, E>>,
@@ -105,10 +111,31 @@ where
             payload_validator,
             executor_provider,
             disallow,
+            cached_state: Default::default(),
         });
 
         Self { inner }
     }
+
+    /// Returns the cached reads for the given head hash.
+    async fn cached_reads(&self, head: B256) -> CachedReads {
+        let cache = self.inner.cached_state.read().await;
+        if cache.0 == head {
+            cache.1.clone()
+        } else {
+            Default::default()
+        }
+    }
+
+    /// Updates the cached state for the given head hash.
+    async fn update_cached_reads(&self, head: B256, cached_state: CachedReads) {
+        let mut cache = self.inner.cached_state.write().await;
+        if cache.0 == head {
+            cache.1.extend(cached_state);
+        } else {
+            *cache = (head, cached_state)
+        }
+    }
 }
 
 impl<Provider, E> ValidationApi<Provider, E>
@@ -119,12 +146,11 @@ where
         + HeaderProvider
         + AccountReader
         + WithdrawalsProvider
-        + Clone
         + 'static,
     E: BlockExecutorProvider,
 {
     /// Validates the given block and a [`BidTrace`] against it.
-    pub fn validate_message_against_block(
+    pub async fn validate_message_against_block(
         &self,
         block: SealedBlockWithSenders,
         message: BidTrace,
@@ -168,8 +194,13 @@ where
         self.consensus.validate_header_against_parent(&block.header, &latest_header)?;
         self.validate_gas_limit(registered_gas_limit, &latest_header, &block.header)?;
 
-        let state_provider = self.provider.state_by_block_hash(latest_header.hash())?;
-        let executor = self.executor_provider.executor(StateProviderDatabase::new(&state_provider));
+        let latest_header_hash = latest_header.hash();
+        let state_provider = self.provider.state_by_block_hash(latest_header_hash)?;
+
+        let mut request_cache = self.cached_reads(latest_header_hash).await;
+
+        let cached_db = request_cache.as_db_mut(StateProviderDatabase::new(&state_provider));
+        let executor = self.executor_provider.executor(cached_db);
 
         let block = block.unseal();
         let mut accessed_blacklisted = None;
@@ -186,6 +217,9 @@ where
             },
         )?;
 
+        // update the cached reads
+        self.update_cached_reads(latest_header_hash, request_cache).await;
+
         if let Some(account) = accessed_blacklisted {
             return Err(ValidationApiError::Blacklist(account))
         }
@@ -413,6 +447,7 @@ where
             request.request.message,
             request.registered_gas_limit,
         )
+        .await
         .map_err(|e| RethError::Other(e.into()))
         .to_rpc_result()
     }
@@ -446,6 +481,7 @@ where
             request.request.message,
             request.registered_gas_limit,
         )
+        .await
         .map_err(|e| RethError::Other(e.into()))
         .to_rpc_result()
     }