feat: perf improvements for build_pending_state (#184)

* wip: limit max depth of pending state
* track pending snapshot height accurately
* clean up unused imports
* move PendingBlocksAPI impl to pending_blocks.rs
* rustfmt
* update err msg

Author: haardikk21

crates/flashblocks-rpc/src/pending_blocks.rs

@@ -1,19 +1,23 @@
+use std::sync::Arc;
+
 use alloy_consensus::{Header, Sealed};
 use alloy_eips::BlockNumberOrTag;
 use alloy_primitives::{
+    Address, B256, BlockNumber, TxHash, U256,
     map::foldhash::{HashMap, HashMapExt},
-    Address, BlockNumber, TxHash, B256, U256,
 };
 use alloy_provider::network::TransactionResponse;
-use alloy_rpc_types::{state::StateOverride, BlockTransactions};
+use alloy_rpc_types::{BlockTransactions, state::StateOverride};
 use alloy_rpc_types_eth::{Filter, Header as RPCHeader, Log};
+use arc_swap::Guard;
 use eyre::eyre;
 use op_alloy_network::Optimism;
 use op_alloy_rpc_types::{OpTransactionReceipt, Transaction};
 use reth::revm::{db::Cache, state::EvmState};
-use reth_rpc_eth_api::RpcBlock;
+use reth_rpc_convert::RpcTransaction;
+use reth_rpc_eth_api::{RpcBlock, RpcReceipt};
 
-use crate::subscription::Flashblock;
+use crate::{rpc::PendingBlocksAPI, subscription::Flashblock};
 
 pub struct PendingBlocksBuilder {
     flashblocks: Vec<Flashblock>,
@@ -25,6 +29,7 @@ pub struct PendingBlocksBuilder {
     transaction_receipts: HashMap<B256, OpTransactionReceipt>,
     transactions_by_hash: HashMap<B256, Transaction>,
     transaction_state: HashMap<B256, EvmState>,
+    transaction_senders: HashMap<B256, Address>,
     state_overrides: Option<StateOverride>,
 
     db_cache: Cache,
@@ -41,6 +46,7 @@ impl PendingBlocksBuilder {
             transaction_receipts: HashMap::new(),
             transactions_by_hash: HashMap::new(),
             transaction_state: HashMap::new(),
+            transaction_senders: HashMap::new(),
             state_overrides: None,
             db_cache: Cache::default(),
         }
@@ -77,6 +83,12 @@ impl PendingBlocksBuilder {
         self
     }
 
+    #[inline]
+    pub(crate) fn with_transaction_sender(&mut self, hash: B256, sender: Address) -> &Self {
+        self.transaction_senders.insert(hash, sender);
+        self
+    }
+
     #[inline]
     pub(crate) fn increment_nonce(&mut self, sender: Address) -> &Self {
         let zero = U256::from(0);
@@ -122,6 +134,7 @@ impl PendingBlocksBuilder {
             transaction_receipts: self.transaction_receipts,
             transactions_by_hash: self.transactions_by_hash,
             transaction_state: self.transaction_state,
+            transaction_senders: self.transaction_senders,
             state_overrides: self.state_overrides,
             db_cache: self.db_cache,
         })
@@ -139,6 +152,7 @@ pub struct PendingBlocks {
     transaction_receipts: HashMap<B256, OpTransactionReceipt>,
     transactions_by_hash: HashMap<B256, Transaction>,
     transaction_state: HashMap<B256, EvmState>,
+    transaction_senders: HashMap<B256, Address>,
     state_overrides: Option<StateOverride>,
 
     db_cache: Cache,
@@ -153,6 +167,10 @@ impl PendingBlocks {
         BlockNumberOrTag::Number(self.headers.first().unwrap().number - 1)
     }
 
+    pub fn earliest_block_number(&self) -> BlockNumber {
+        self.headers.first().unwrap().number
+    }
+
     pub fn latest_flashblock_index(&self) -> u64 {
         self.flashblocks.last().unwrap().index
     }
@@ -165,8 +183,12 @@ impl PendingBlocks {
         self.flashblocks.clone()
     }
 
-    pub fn get_transaction_state(&self, hash: B256) -> Option<EvmState> {
-        self.transaction_state.get(&hash).cloned()
+    pub fn get_transaction_state(&self, hash: &B256) -> Option<EvmState> {
+        self.transaction_state.get(hash).cloned()
+    }
+
+    pub fn get_transaction_sender(&self, tx_hash: &B256) -> Option<Address> {
+        self.transaction_senders.get(tx_hash).cloned()
     }
 
     pub fn get_db_cache(&self) -> Cache {
@@ -236,3 +258,43 @@ impl PendingBlocks {
         logs
     }
 }
+
+impl PendingBlocksAPI for Guard<Option<Arc<PendingBlocks>>> {
+    fn get_canonical_block_number(&self) -> BlockNumberOrTag {
+        self.as_ref().map(|pb| pb.canonical_block_number()).unwrap_or(BlockNumberOrTag::Latest)
+    }
+
+    fn get_transaction_count(&self, address: Address) -> U256 {
+        self.as_ref().map(|pb| pb.get_transaction_count(address)).unwrap_or_else(|| U256::from(0))
+    }
+
+    fn get_block(&self, full: bool) -> Option<RpcBlock<Optimism>> {
+        self.as_ref().map(|pb| pb.get_latest_block(full))
+    }
+
+    fn get_transaction_receipt(
+        &self,
+        tx_hash: alloy_primitives::TxHash,
+    ) -> Option<RpcReceipt<Optimism>> {
+        self.as_ref().and_then(|pb| pb.get_receipt(tx_hash))
+    }
+
+    fn get_transaction_by_hash(
+        &self,
+        tx_hash: alloy_primitives::TxHash,
+    ) -> Option<RpcTransaction<Optimism>> {
+        self.as_ref().and_then(|pb| pb.get_transaction_by_hash(tx_hash))
+    }
+
+    fn get_balance(&self, address: Address) -> Option<U256> {
+        self.as_ref().and_then(|pb| pb.get_balance(address))
+    }
+
+    fn get_state_overrides(&self) -> Option<StateOverride> {
+        self.as_ref().map(|pb| pb.get_state_overrides()).unwrap_or_default()
+    }
+
+    fn get_pending_logs(&self, filter: &Filter) -> Vec<Log> {
+        self.as_ref().map(|pb| pb.get_pending_logs(filter)).unwrap_or_default()
+    }
+}

crates/flashblocks-rpc/src/rpc.rs

@@ -2,36 +2,36 @@ use std::{sync::Arc, time::Duration};
 
 use alloy_eips::{BlockId, BlockNumberOrTag};
 use alloy_primitives::{
-    map::foldhash::{HashSet, HashSetExt},
     Address, TxHash, U256,
+    map::foldhash::{HashSet, HashSetExt},
 };
 use alloy_rpc_types::{
+    BlockOverrides,
     simulate::{SimBlock, SimulatePayload, SimulatedBlock},
     state::{EvmOverrides, StateOverride, StateOverridesBuilder},
-    BlockOverrides,
 };
 use alloy_rpc_types_eth::{Filter, Log};
 use arc_swap::Guard;
 use jsonrpsee::{
-    core::{async_trait, RpcResult},
+    core::{RpcResult, async_trait},
     proc_macros::rpc,
 };
-use jsonrpsee_types::{error::INVALID_PARAMS_CODE, ErrorObjectOwned};
+use jsonrpsee_types::{ErrorObjectOwned, error::INVALID_PARAMS_CODE};
 use op_alloy_network::Optimism;
 use op_alloy_rpc_types::OpTransactionRequest;
 use reth::{
     providers::CanonStateSubscriptions,
     rpc::{eth::EthFilter, server_types::eth::EthApiError},
 };
 use reth_rpc_eth_api::{
-    helpers::{EthBlocks, EthCall, EthState, EthTransactions, FullEthApi},
     EthApiTypes, EthFilterApiServer, RpcBlock, RpcReceipt, RpcTransaction,
+    helpers::{EthBlocks, EthCall, EthState, EthTransactions, FullEthApi},
 };
 use tokio::{
     sync::{broadcast, broadcast::error::RecvError},
     time,
 };
-use tokio_stream::{wrappers::BroadcastStream, StreamExt};
+use tokio_stream::{StreamExt, wrappers::BroadcastStream};
 use tracing::{debug, trace, warn};
 
 use crate::{metrics::Metrics, pending_blocks::PendingBlocks};
@@ -91,7 +91,7 @@ pub trait EthApiOverride {
 
     #[method(name = "getBalance")]
     async fn get_balance(&self, address: Address, block_number: Option<BlockId>)
-        -> RpcResult<U256>;
+    -> RpcResult<U256>;
 
     #[method(name = "getTransactionCount")]
     async fn get_transaction_count(