Cache flashblock trie nodes to optimize bundle metering

Adds a single-entry cache for the latest flashblock's trie nodes, allowing
bundle metering operations to reuse the cached trie computation instead of
recalculating from scratch each time.

Key changes:
- Add FlashblockTrieCache: thread-safe single-entry cache for latest flashblock
- Add FlashblockTrieData: contains trie updates and hashed state for reuse
- Cache keyed by block hash and flashblock index for accuracy
- Compute flashblock trie once per flashblock, reuse for all bundle operations
- Extract flashblock trie calculation outside timing metrics
- Use TrieInput.prepend_cached() to combine flashblock + bundle tries

The cache replaces previous entries when a new flashblock is cached, as it's
designed only for the current/latest flashblock, not historical ones.

Author: niran

Cargo.lock

@@ -75,6 +75,7 @@ reth-db = { git = "https://github.com/paradigmxyz/reth", tag = "v1.8.2" }
 reth-testing-utils = { git = "https://github.com/paradigmxyz/reth", tag = "v1.8.2" }
 reth-db-common = { git = "https://github.com/paradigmxyz/reth", tag = "v1.8.2" }
 reth-ipc = { git = "https://github.com/paradigmxyz/reth", tag = "v1.8.2" }
+reth-trie-common = { git = "https://github.com/paradigmxyz/reth", tag = "v1.8.2" }
 
 # revm
 revm = { version = "29.0.0", default-features = false }

Cargo.toml

@@ -26,6 +26,7 @@ reth-optimism-chainspec.workspace = true
 reth-optimism-primitives.workspace = true
 reth-transaction-pool.workspace = true
 reth-optimism-cli.workspace = true  # Enables serde & codec traits for OpReceipt/OpTxEnvelope
+reth-trie-common.workspace = true
 
 # alloy
 alloy-primitives.workspace = true
@@ -49,6 +50,7 @@ jsonrpsee.workspace = true
 tracing.workspace = true
 serde.workspace = true
 eyre.workspace = true
+arc-swap.workspace = true
 
 [dev-dependencies]
 alloy-genesis.workspace = true

crates/metering/Cargo.toml

@@ -0,0 +1,121 @@
+use alloy_primitives::B256;
+use arc_swap::ArcSwap;
+use eyre::Result as EyreResult;
+use reth_provider::StateProvider;
+use std::sync::Arc;
+
+use crate::FlashblocksState;
+
+/// Trie nodes and hashed state from computing a flashblock state root.
+///
+/// These cached nodes can be reused when computing a bundle's state root
+/// to avoid recalculating the flashblock portion of the trie.
+#[derive(Debug, Clone)]
+pub struct FlashblockTrieData {
+    pub trie_updates: reth_trie_common::updates::TrieUpdates,
+    pub hashed_state: reth_trie_common::HashedPostState,
+}
+
+/// Internal cache entry for a single flashblock.
+#[derive(Debug, Clone)]
+struct CachedFlashblockTrie {
+    block_hash: B256,
+    flashblock_index: u64,
+    /// The cached trie data
+    trie_data: FlashblockTrieData,
+}
+
+/// Thread-safe single-entry cache for the latest flashblock's trie nodes.
+///
+/// This cache stores the intermediate trie nodes computed when calculating
+/// the latest flashblock's state root. Subsequent bundle metering operations
+/// on the same flashblock can reuse these cached nodes instead of recalculating
+/// them, significantly improving performance.
+///
+/// **Important**: This cache holds only ONE flashblock's trie at a time.
+/// When a new flashblock is cached, it replaces any previously cached flashblock.
+/// This design assumes that bundle metering operations are performed on the
+/// current/latest flashblock, not historical ones.
+#[derive(Debug, Clone)]
+pub struct FlashblockTrieCache {
+    /// Single-entry cache for the latest flashblock's trie
+    cache: Arc<ArcSwap<Option<CachedFlashblockTrie>>>,
+}
+
+impl FlashblockTrieCache {
+    /// Creates a new empty flashblock trie cache.
+    pub fn new() -> Self {
+        Self {
+            cache: Arc::new(ArcSwap::from_pointee(None)),
+        }
+    }
+
+    /// Ensures the trie for the given flashblock is cached and returns it.
+    ///
+    /// If the cache already contains an entry for the given block_hash and flashblock_index,
+    /// this returns the cached data immediately without recomputation. Otherwise, it computes
+    /// the flashblock's state root, caches the resulting trie nodes, **replacing any previously
+    /// cached flashblock**, and returns the new data.
+    ///
+    /// # Single-Entry Cache Behavior
+    ///
+    /// This cache only stores one flashblock's trie at a time. Calling this method with a different
+    /// flashblock will evict the previous entry. This is by design, as the cache is intended for
+    /// the latest/current flashblock only.
+    ///
+    /// # Arguments
+    ///
+    /// * `block_hash` - Hash of the block containing the flashblock
+    /// * `flashblock_index` - Index of the flashblock within the block
+    /// * `flashblocks_state` - The accumulated state from pending flashblocks
+    /// * `canonical_state_provider` - State provider for the canonical chain
+    ///
+    /// # Returns
+    ///
+    /// The cached `FlashblockTrieData` containing the trie updates and hashed state.
+    pub fn ensure_cached(
+        &self,
+        block_hash: B256,
+        flashblock_index: u64,
+        flashblocks_state: &FlashblocksState,
+        canonical_state_provider: &dyn StateProvider,
+    ) -> EyreResult<FlashblockTrieData> {
+        // Check if we already have a cached trie for this exact flashblock
+        let cached = self.cache.load();
+        if let Some(ref cache) = **cached {
+            if cache.block_hash == block_hash && cache.flashblock_index == flashblock_index {
+                // Cache is still valid for this flashblock, return it
+                return Ok(cache.trie_data.clone());
+            }
+        }
+
+        // Need to compute the flashblock trie (this will replace any existing cache entry)
+
+        // Compute hashed post state from the bundle
+        let hashed_state = canonical_state_provider.hashed_post_state(&flashblocks_state.bundle_state);
+
+        // Calculate state root with updates to get the trie nodes
+        let (_state_root, trie_updates) = canonical_state_provider.state_root_with_updates(hashed_state.clone())?;
+
+        // Create the trie data
+        let trie_data = FlashblockTrieData {
+            trie_updates,
+            hashed_state,
+        };
+
+        // Store the new entry, replacing any previous flashblock's cached trie
+        self.cache.store(Arc::new(Some(CachedFlashblockTrie {
+            block_hash,
+            flashblock_index,
+            trie_data: trie_data.clone(),
+        })));
+
+        Ok(trie_data)
+    }
+}
+
+impl Default for FlashblockTrieCache {
+    fn default() -> Self {
+        Self::new()
+    }
+}

crates/metering/src/flashblock_trie_cache.rs

@@ -1,8 +1,10 @@
+mod flashblock_trie_cache;
 mod meter;
 mod rpc;
 #[cfg(test)]
 mod tests;
 
+pub use flashblock_trie_cache::{FlashblockTrieCache, FlashblockTrieData};
 pub use meter::{meter_bundle, FlashblocksState, MeterBundleOutput};
 pub use rpc::{MeteringApiImpl, MeteringApiServer};
 pub use tips_core::types::{Bundle, MeterBundleResponse, TransactionResult};

crates/metering/src/lib.rs

@@ -2,12 +2,13 @@ use alloy_consensus::{transaction::SignerRecoverable, BlockHeader, Transaction a
 use alloy_primitives::{B256, U256};
 use eyre::{eyre, Result as EyreResult};
 use reth::revm::db::{BundleState, Cache, CacheDB, State};
-use revm_database::states::bundle_state::BundleRetention;
 use reth_evm::execute::BlockBuilder;
 use reth_evm::ConfigureEvm;
 use reth_optimism_chainspec::OpChainSpec;
 use reth_optimism_evm::{OpEvmConfig, OpNextBlockEnvAttributes};
 use reth_primitives_traits::SealedHeader;
+use reth_trie_common::TrieInput;
+use revm_database::states::bundle_state::BundleRetention;
 use std::sync::Arc;
 use std::time::Instant;
 
@@ -54,13 +55,25 @@ pub fn meter_bundle<SP>(
     header: &SealedHeader,
     bundle_with_metadata: &tips_core::types::BundleWithMetadata,
     flashblocks_state: Option<FlashblocksState>,
+    cached_flashblock_trie: Option<crate::FlashblockTrieData>,
 ) -> EyreResult<MeterBundleOutput>
 where
     SP: reth_provider::StateProvider,
 {
     // Get bundle hash from BundleWithMetadata
     let bundle_hash = bundle_with_metadata.bundle_hash();
 
+    // If we have flashblocks but no cached trie, compute the flashblock trie first
+    // (before starting any timers, since we only want to time the bundle's execution and state root)
+    let flashblock_trie = if cached_flashblock_trie.is_none() && flashblocks_state.is_some() {
+        let fb_state = flashblocks_state.as_ref().unwrap();
+        let fb_hashed_state = state_provider.hashed_post_state(&fb_state.bundle_state);
+        let (_fb_state_root, fb_trie_updates) = state_provider.state_root_with_updates(fb_hashed_state.clone())?;
+        Some((fb_trie_updates, fb_hashed_state))
+    } else {
+        None
+    };
+
     // Create state database
     let state_db = reth::revm::database::StateProviderDatabase::new(state_provider);
 
@@ -152,14 +165,28 @@ where
     }
 
     // Calculate state root and measure its calculation time
-    // The bundle already includes flashblocks state if it was provided via with_bundle_prestate
     db.merge_transitions(BundleRetention::Reverts);
     let bundle = db.take_bundle();
     let state_provider = db.database.db.as_ref();
-    let state_root_start = Instant::now();
 
+    let state_root_start = Instant::now();
     let hashed_post_state = state_provider.hashed_post_state(&bundle);
-    let _ = state_provider.state_root_with_updates(hashed_post_state);
+
+    if let Some(cached) = cached_flashblock_trie {
+        // We have cached flashblock trie nodes, use them
+        let mut trie_input = TrieInput::from_state(hashed_post_state);
+        trie_input.prepend_cached(cached.trie_updates, cached.hashed_state);
+        let _ = state_provider.state_root_from_nodes_with_updates(trie_input)?;
+    } else if let Some((fb_trie_updates, fb_hashed_state)) = flashblock_trie {
+        // We computed the flashblock trie above, now use it
+        let mut trie_input = TrieInput::from_state(hashed_post_state);
+        trie_input.prepend_cached(fb_trie_updates, fb_hashed_state);
+        let _ = state_provider.state_root_from_nodes_with_updates(trie_input)?;
+    } else {
+        // No flashblocks, just calculate bundle state root
+        let _ = state_provider.state_root_with_updates(hashed_post_state)?;
+    }
+
     let state_root_time = state_root_start.elapsed().as_micros();
 
     let total_execution_time = execution_start.elapsed().as_micros();

crates/metering/src/meter.rs

@@ -14,7 +14,7 @@ use std::sync::Arc;
 use tips_core::types::{Bundle, BundleWithMetadata, MeterBundleResponse};
 use tracing::{error, info};
 
-use crate::meter_bundle;
+use crate::{meter_bundle, FlashblockTrieCache};
 
 /// RPC API for transaction metering
 #[rpc(server, namespace = "base")]
@@ -28,6 +28,8 @@ pub trait MeteringApi {
 pub struct MeteringApiImpl<Provider, FB> {
     provider: Provider,
     flashblocks_state: Arc<FB>,
+    /// Single-entry cache for the latest flashblock's trie nodes
+    trie_cache: FlashblockTrieCache,
 }
 
 impl<Provider, FB> MeteringApiImpl<Provider, FB>
@@ -43,6 +45,7 @@ where
         Self {
             provider,
             flashblocks_state,
+            trie_cache: FlashblockTrieCache::new(),
         }
     }
 }
@@ -161,6 +164,32 @@ where
         // Get the flashblock index if we have pending flashblocks
         let state_flashblock_index = pending_blocks.as_ref().map(|pb| pb.latest_flashblock_index());
 
+        // If we have flashblocks, ensure the trie is cached and get it
+        let cached_trie = if let Some(ref fb_state) = flashblocks_state {
+            let fb_index = state_flashblock_index.unwrap();
+
+            // Ensure the flashblock trie is cached and return it
+            Some(
+                self.trie_cache
+                    .ensure_cached(
+                        header.hash(),
+                        fb_index,
+                        fb_state,
+                        &*state_provider,
+                    )
+                    .map_err(|e| {
+                        error!(error = %e, "Failed to cache flashblock trie");
+                        jsonrpsee::types::ErrorObjectOwned::owned(
+                            jsonrpsee::types::ErrorCode::InternalError.code(),
+                            format!("Failed to cache flashblock trie: {}", e),
+                            None::<()>,
+                        )
+                    })?,
+            )
+        } else {
+            None
+        };
+
         // Meter bundle using utility function
         let result = meter_bundle(
             state_provider,
@@ -169,6 +198,7 @@ where
             &header,
             &bundle_with_metadata,
             flashblocks_state,
+            cached_trie,
         )
         .map_err(|e| {
             error!(error = %e, "Bundle metering failed");

crates/metering/src/rpc.rs

@@ -162,6 +162,7 @@ fn meter_bundle_empty_transactions() -> eyre::Result<()> {
         &harness.header,
         &bundle_with_metadata,
         None,
+        None,
     )?;
 
     assert!(output.results.is_empty());
@@ -211,6 +212,7 @@ fn meter_bundle_single_transaction() -> eyre::Result<()> {
         &harness.header,
         &bundle_with_metadata,
         None,
+        None,
     )?;
 
     assert_eq!(output.results.len(), 1);
@@ -308,6 +310,7 @@ fn meter_bundle_multiple_transactions() -> eyre::Result<()> {
         &harness.header,
         &bundle_with_metadata,
         None,
+        None,
     )?;
 
     assert_eq!(output.results.len(), 2);
@@ -397,6 +400,7 @@ fn meter_bundle_state_root_time_invariant() -> eyre::Result<()> {
         &harness.header,
         &bundle_with_metadata,
         None,
+        None,
     )?;
 
     // Verify invariant: total execution time must include state root time