From ac2853884a624ca112815ca958481bee6f9ed542 Mon Sep 17 00:00:00 2001 From: Federico Gimenez Date: Fri, 26 Jul 2024 11:51:11 +0200 Subject: [PATCH] feat: integrate memory overlay in canonical state (#9817) --- crates/chain-state/src/in_memory.rs | 246 +++++++++++++++++- crates/chain-state/src/memory_overlay.rs | 6 +- .../src/providers/blockchain_provider.rs | 41 +-- 3 files changed, 265 insertions(+), 28 deletions(-) diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs index 6ce742c69e..1ac01b57b1 100644 --- a/crates/chain-state/src/in_memory.rs +++ b/crates/chain-state/src/in_memory.rs @@ -1,7 +1,8 @@ //! Types for tracking the canonical chain state in memory. use crate::{ - CanonStateNotification, CanonStateNotificationSender, CanonStateNotifications, ChainInfoTracker, + CanonStateNotification, CanonStateNotificationSender, CanonStateNotifications, + ChainInfoTracker, MemoryOverlayStateProvider, }; use parking_lot::RwLock; use reth_chainspec::ChainInfo; @@ -10,6 +11,7 @@ use reth_primitives::{ Address, BlockNumHash, Header, Receipt, Receipts, SealedBlock, SealedBlockWithSenders, SealedHeader, B256, }; +use reth_storage_api::StateProviderBox; use reth_trie::{updates::TrieUpdates, HashedPostState}; use std::{collections::HashMap, ops::Deref, sync::Arc, time::Instant}; use tokio::sync::broadcast; @@ -354,6 +356,24 @@ impl CanonicalInMemoryState { pub fn notify_canon_state(&self, event: CanonStateNotification) { self.inner.canon_state_notification_sender.send(event).ok(); } + + /// Return state provider with reference to in-memory blocks that overlay database state. + /// + /// This merges the state of all blocks that are part of the chain that the requested block is + /// the head of. This includes all blocks that connect back to the canonical block on disk. + pub fn state_provider( + &self, + hash: B256, + historical: StateProviderBox, + ) -> MemoryOverlayStateProvider { + let in_memory = if let Some(state) = self.state_by_hash(hash) { + state.chain().into_iter().map(|block_state| block_state.block()).collect() + } else { + Vec::new() + }; + + MemoryOverlayStateProvider::new(in_memory, historical) + } } /// State after applying the given block, this block is part of the canonical chain that partially @@ -434,6 +454,27 @@ impl BlockState { }) .unwrap_or_default() } + + /// Returns a vector of parent `BlockStates` starting from the oldest one. + pub fn parent_state_chain(&self) -> Vec<&Self> { + let mut parents = Vec::new(); + let mut current = self.parent.as_deref(); + + while let Some(parent) = current { + parents.insert(0, parent); + current = parent.parent.as_deref(); + } + + parents + } + + /// Returns a vector of `BlockStates` representing the entire in memory chain, + /// including self as the last element. + pub fn chain(&self) -> Vec<&Self> { + let mut chain = self.parent_state_chain(); + chain.push(self); + chain + } } /// Represents an executed block stored in-memory. @@ -575,17 +616,104 @@ mod tests { use super::*; use crate::test_utils::{get_executed_block_with_number, get_executed_block_with_receipts}; use rand::Rng; - use reth_primitives::Receipt; + use reth_errors::ProviderResult; + use reth_primitives::{Account, BlockNumber, Bytecode, Receipt, StorageKey, StorageValue}; + use reth_storage_api::{ + AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider, + }; + use reth_trie::AccountProof; - fn create_mock_state(block_number: u64) -> BlockState { - BlockState::new(get_executed_block_with_number(block_number, B256::random())) + fn create_mock_state(block_number: u64, parent_hash: B256) -> BlockState { + BlockState::new(get_executed_block_with_number(block_number, parent_hash)) + } + + fn create_mock_state_chain(num_blocks: u64) -> Vec { + let mut chain = Vec::with_capacity(num_blocks as usize); + let mut parent_hash = B256::random(); + let mut parent_state: Option = None; + + for i in 1..=num_blocks { + let mut state = create_mock_state(i, parent_hash); + if let Some(parent) = parent_state { + state.parent = Some(Box::new(parent)); + } + parent_hash = state.hash(); + parent_state = Some(state.clone()); + chain.push(state); + } + + chain + } + + struct MockStateProvider; + + impl StateProvider for MockStateProvider { + fn storage( + &self, + _address: Address, + _storage_key: StorageKey, + ) -> ProviderResult> { + Ok(None) + } + + fn bytecode_by_hash(&self, _code_hash: B256) -> ProviderResult> { + Ok(None) + } + } + + impl BlockHashReader for MockStateProvider { + fn block_hash(&self, _number: BlockNumber) -> ProviderResult> { + Ok(None) + } + + fn canonical_hashes_range( + &self, + _start: BlockNumber, + _end: BlockNumber, + ) -> ProviderResult> { + Ok(vec![]) + } + } + + impl AccountReader for MockStateProvider { + fn basic_account(&self, _address: Address) -> ProviderResult> { + Ok(None) + } + } + + impl StateRootProvider for MockStateProvider { + fn hashed_state_root(&self, _hashed_state: &HashedPostState) -> ProviderResult { + Ok(B256::random()) + } + + fn hashed_state_root_with_updates( + &self, + _hashed_state: &HashedPostState, + ) -> ProviderResult<(B256, TrieUpdates)> { + Ok((B256::random(), TrieUpdates::default())) + } + + fn state_root(&self, _bundle_state: &revm::db::BundleState) -> ProviderResult { + Ok(B256::random()) + } + } + + impl StateProofProvider for MockStateProvider { + fn hashed_proof( + &self, + _hashed_state: &HashedPostState, + _address: Address, + _slots: &[B256], + ) -> ProviderResult { + Ok(AccountProof::new(Address::random())) + } } #[test] fn test_in_memory_state_impl_state_by_hash() { let mut state_by_hash = HashMap::new(); let number = rand::thread_rng().gen::(); - let state = Arc::new(create_mock_state(number)); + let state = Arc::new(create_mock_state(number, B256::random())); state_by_hash.insert(state.hash(), state.clone()); let in_memory_state = InMemoryState::new(state_by_hash, HashMap::new(), None); @@ -600,7 +728,7 @@ mod tests { let mut hash_by_number = HashMap::new(); let number = rand::thread_rng().gen::(); - let state = Arc::new(create_mock_state(number)); + let state = Arc::new(create_mock_state(number, B256::random())); let hash = state.hash(); state_by_hash.insert(hash, state.clone()); @@ -616,9 +744,9 @@ mod tests { fn test_in_memory_state_impl_head_state() { let mut state_by_hash = HashMap::new(); let mut hash_by_number = HashMap::new(); - let state1 = Arc::new(create_mock_state(1)); - let state2 = Arc::new(create_mock_state(2)); + let state1 = Arc::new(create_mock_state(1, B256::random())); let hash1 = state1.hash(); + let state2 = Arc::new(create_mock_state(2, hash1)); let hash2 = state2.hash(); hash_by_number.insert(1, hash1); hash_by_number.insert(2, hash2); @@ -635,7 +763,7 @@ mod tests { #[test] fn test_in_memory_state_impl_pending_state() { let pending_number = rand::thread_rng().gen::(); - let pending_state = create_mock_state(pending_number); + let pending_state = create_mock_state(pending_number, B256::random()); let pending_hash = pending_state.hash(); let in_memory_state = @@ -733,4 +861,104 @@ mod tests { assert_eq!(state.inner.in_memory_state.block_count(), 1); } + + #[test] + fn test_canonical_in_memory_state_state_provider() { + let block1 = get_executed_block_with_number(1, B256::random()); + let block2 = get_executed_block_with_number(2, block1.block().hash()); + let block3 = get_executed_block_with_number(3, block2.block().hash()); + + let state1 = BlockState::new(block1.clone()); + let state2 = BlockState::with_parent(block2.clone(), Some(state1.clone())); + let state3 = BlockState::with_parent(block3.clone(), Some(state2.clone())); + + let mut blocks = HashMap::new(); + blocks.insert(block1.block().hash(), Arc::new(state1)); + blocks.insert(block2.block().hash(), Arc::new(state2)); + blocks.insert(block3.block().hash(), Arc::new(state3)); + + let mut numbers = HashMap::new(); + numbers.insert(1, block1.block().hash()); + numbers.insert(2, block2.block().hash()); + numbers.insert(3, block3.block().hash()); + + let canonical_state = CanonicalInMemoryState::new(blocks, numbers, None); + + let historical: StateProviderBox = Box::new(MockStateProvider); + + let overlay_provider = canonical_state.state_provider(block3.block().hash(), historical); + + assert_eq!(overlay_provider.in_memory.len(), 3); + assert_eq!(overlay_provider.in_memory[0].block().number, 1); + assert_eq!(overlay_provider.in_memory[1].block().number, 2); + assert_eq!(overlay_provider.in_memory[2].block().number, 3); + + assert_eq!( + overlay_provider.in_memory[1].block().parent_hash, + overlay_provider.in_memory[0].block().hash() + ); + assert_eq!( + overlay_provider.in_memory[2].block().parent_hash, + overlay_provider.in_memory[1].block().hash() + ); + + let unknown_hash = B256::random(); + let empty_overlay_provider = + canonical_state.state_provider(unknown_hash, Box::new(MockStateProvider)); + assert_eq!(empty_overlay_provider.in_memory.len(), 0); + } + + #[test] + fn test_block_state_parent_blocks() { + let chain = create_mock_state_chain(4); + + let parents = chain[3].parent_state_chain(); + assert_eq!(parents.len(), 3); + assert_eq!(parents[0].block().block.number, 1); + assert_eq!(parents[1].block().block.number, 2); + assert_eq!(parents[2].block().block.number, 3); + + let parents = chain[2].parent_state_chain(); + assert_eq!(parents.len(), 2); + assert_eq!(parents[0].block().block.number, 1); + assert_eq!(parents[1].block().block.number, 2); + + let parents = chain[0].parent_state_chain(); + assert_eq!(parents.len(), 0); + } + + #[test] + fn test_block_state_single_block_state_chain() { + let single_block_number = 1; + let single_block = create_mock_state(single_block_number, B256::random()); + let single_block_hash = single_block.block().block.hash(); + + let parents = single_block.parent_state_chain(); + assert_eq!(parents.len(), 0); + + let block_state_chain = single_block.chain(); + assert_eq!(block_state_chain.len(), 1); + assert_eq!(block_state_chain[0].block().block.number, single_block_number); + assert_eq!(block_state_chain[0].block().block.hash(), single_block_hash); + } + + #[test] + fn test_block_state_chain() { + let chain = create_mock_state_chain(3); + + let block_state_chain = chain[2].chain(); + assert_eq!(block_state_chain.len(), 3); + assert_eq!(block_state_chain[0].block().block.number, 1); + assert_eq!(block_state_chain[1].block().block.number, 2); + assert_eq!(block_state_chain[2].block().block.number, 3); + + let block_state_chain = chain[1].chain(); + assert_eq!(block_state_chain.len(), 2); + assert_eq!(block_state_chain[0].block().block.number, 1); + assert_eq!(block_state_chain[1].block().block.number, 2); + + let block_state_chain = chain[0].chain(); + assert_eq!(block_state_chain.len(), 1); + assert_eq!(block_state_chain[0].block().block.number, 1); + } } diff --git a/crates/chain-state/src/memory_overlay.rs b/crates/chain-state/src/memory_overlay.rs index 65b3b89bb8..cba585018a 100644 --- a/crates/chain-state/src/memory_overlay.rs +++ b/crates/chain-state/src/memory_overlay.rs @@ -11,11 +11,11 @@ use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState}; #[allow(missing_debug_implementations)] pub struct MemoryOverlayStateProvider { /// The collection of executed parent blocks. - in_memory: Vec, + pub(crate) in_memory: Vec, /// The collection of hashed state from in-memory blocks. - hashed_post_state: HashedPostState, + pub(crate) hashed_post_state: HashedPostState, /// Historical state provider for state lookups that are not found in in-memory blocks. - historical: Box, + pub(crate) historical: Box, } impl MemoryOverlayStateProvider { diff --git a/crates/storage/provider/src/providers/blockchain_provider.rs b/crates/storage/provider/src/providers/blockchain_provider.rs index 45cc5aa072..952cdff4bc 100644 --- a/crates/storage/provider/src/providers/blockchain_provider.rs +++ b/crates/storage/provider/src/providers/blockchain_provider.rs @@ -1,12 +1,12 @@ use crate::{ - providers::BundleStateProvider, AccountReader, BlockHashReader, BlockIdReader, BlockNumReader, - BlockReader, BlockReaderIdExt, BlockSource, BlockchainTreePendingStateProvider, - CanonChainTracker, CanonStateNotifications, CanonStateSubscriptions, ChainSpecProvider, - ChangeSetReader, DatabaseProviderFactory, DatabaseProviderRO, EvmEnvProvider, - FullExecutionDataProvider, HeaderProvider, ProviderError, ProviderFactory, - PruneCheckpointReader, ReceiptProvider, ReceiptProviderIdExt, RequestsProvider, - StageCheckpointReader, StateProviderBox, StateProviderFactory, StaticFileProviderFactory, - TransactionVariant, TransactionsProvider, WithdrawalsProvider, + providers::{BundleStateProvider, StaticFileProvider}, + AccountReader, BlockHashReader, BlockIdReader, BlockNumReader, BlockReader, BlockReaderIdExt, + BlockSource, BlockchainTreePendingStateProvider, CanonChainTracker, CanonStateNotifications, + CanonStateSubscriptions, ChainSpecProvider, ChangeSetReader, DatabaseProviderFactory, + DatabaseProviderRO, EvmEnvProvider, FullExecutionDataProvider, HeaderProvider, ProviderError, + ProviderFactory, PruneCheckpointReader, ReceiptProvider, ReceiptProviderIdExt, + RequestsProvider, StageCheckpointReader, StateProviderBox, StateProviderFactory, + StaticFileProviderFactory, TransactionVariant, TransactionsProvider, WithdrawalsProvider, }; use alloy_rpc_types_engine::ForkchoiceState; use reth_chain_state::CanonicalInMemoryState; @@ -33,8 +33,6 @@ use std::{ }; use tracing::trace; -use super::StaticFileProvider; - /// The main type for interacting with the blockchain. /// /// This type serves as the main entry point for interacting with the blockchain and provides data @@ -624,15 +622,28 @@ where fn pending(&self) -> ProviderResult { trace!(target: "providers::blockchain", "Getting provider for pending state"); - // TODO: check in memory overlay https://github.com/paradigmxyz/reth/issues/9614 + if let Some(block) = self.canonical_in_memory_state.pending_block_num_hash() { + let historical = self.database.history_by_block_hash(block.hash)?; + let pending_provider = + self.canonical_in_memory_state.state_provider(block.hash, historical); + + return Ok(Box::new(pending_provider)); + } // fallback to latest state if the pending block is not available self.latest() } - fn pending_state_by_hash(&self, _block_hash: B256) -> ProviderResult> { - // TODO: check in memory overlay https://github.com/paradigmxyz/reth/issues/9614 + fn pending_state_by_hash(&self, block_hash: B256) -> ProviderResult> { + let historical = self.database.history_by_block_hash(block_hash)?; + if let Some(block) = self.canonical_in_memory_state.pending_block_num_hash() { + if block.hash == block_hash { + let pending_provider = + self.canonical_in_memory_state.state_provider(block_hash, historical); + return Ok(Some(Box::new(pending_provider))) + } + } Ok(None) } @@ -640,10 +651,8 @@ where &self, bundle_state_data: Box, ) -> ProviderResult { - let canonical_fork = bundle_state_data.canonical_fork(); - trace!(target: "providers::blockchain", ?canonical_fork, "Returning post state provider"); + let state_provider = self.pending()?; - let state_provider = self.history_by_block_hash(canonical_fork.hash)?; let bundle_state_provider = BundleStateProvider::new(state_provider, bundle_state_data); Ok(Box::new(bundle_state_provider)) }