perf(engine): lazy load overlay trie state (#10541)

This commit is contained in:
Roman Krasiuk
2024-08-26 06:11:54 -07:00
committed by GitHub
parent 1ac5336303
commit c4a231335a
3 changed files with 52 additions and 34 deletions

View File

@ -493,7 +493,7 @@ impl CanonicalInMemoryState {
Vec::new()
};
MemoryOverlayStateProvider::new(in_memory, historical)
MemoryOverlayStateProvider::new(historical, in_memory)
}
/// Returns an iterator over all canonical blocks in the in-memory state, from newest to oldest.

View File

@ -11,20 +11,18 @@ use reth_trie::{
prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
HashedStorage,
};
use std::collections::HashMap;
use std::{collections::HashMap, sync::OnceLock};
/// A state provider that stores references to in-memory blocks along with their state as well as
/// the historical state provider for fallback lookups.
#[allow(missing_debug_implementations)]
pub struct MemoryOverlayStateProvider {
/// The collection of executed parent blocks. Expected order is newest to oldest.
pub(crate) in_memory: Vec<ExecutedBlock>,
/// The collection of hashed state from in-memory blocks.
pub(crate) hashed_post_state: HashedPostState,
/// The collection of aggregated in-memory trie updates.
pub(crate) trie_updates: TrieUpdates,
/// Historical state provider for state lookups that are not found in in-memory blocks.
pub(crate) historical: Box<dyn StateProvider>,
/// The collection of executed parent blocks. Expected order is newest to oldest.
pub(crate) in_memory: Vec<ExecutedBlock>,
/// Lazy-loaded in-memory trie data.
pub(crate) trie_state: OnceLock<MemoryOverlayTrieState>,
}
impl MemoryOverlayStateProvider {
@ -35,20 +33,29 @@ impl MemoryOverlayStateProvider {
/// - `in_memory` - the collection of executed ancestor blocks in reverse.
/// - `historical` - a historical state provider for the latest ancestor block stored in the
/// database.
pub fn new(in_memory: Vec<ExecutedBlock>, historical: Box<dyn StateProvider>) -> Self {
let mut hashed_post_state = HashedPostState::default();
let mut trie_updates = TrieUpdates::default();
for block in in_memory.iter().rev() {
hashed_post_state.extend_ref(block.hashed_state.as_ref());
trie_updates.extend_ref(block.trie.as_ref());
}
Self { in_memory, hashed_post_state, trie_updates, historical }
pub fn new(historical: Box<dyn StateProvider>, in_memory: Vec<ExecutedBlock>) -> Self {
Self { historical, in_memory, trie_state: OnceLock::new() }
}
/// Turn this state provider into a [`StateProviderBox`]
pub fn boxed(self) -> StateProviderBox {
Box::new(self)
}
/// Return lazy-loaded trie state aggregated from in-memory blocks.
fn trie_state(&self) -> MemoryOverlayTrieState {
self.trie_state
.get_or_init(|| {
let mut hashed_state = HashedPostState::default();
let mut trie_nodes = TrieUpdates::default();
for block in self.in_memory.iter().rev() {
hashed_state.extend_ref(block.hashed_state.as_ref());
trie_nodes.extend_ref(block.trie.as_ref());
}
MemoryOverlayTrieState { trie_nodes, hashed_state }
})
.clone()
}
}
impl BlockHashReader for MemoryOverlayStateProvider {
@ -105,14 +112,13 @@ impl StateRootProvider for MemoryOverlayStateProvider {
fn hashed_state_root_from_nodes(
&self,
nodes: TrieUpdates,
hashed_state: HashedPostState,
state: HashedPostState,
prefix_sets: TriePrefixSetsMut,
) -> ProviderResult<B256> {
let mut trie_nodes = self.trie_updates.clone();
let MemoryOverlayTrieState { mut trie_nodes, mut hashed_state } = self.trie_state();
trie_nodes.extend(nodes);
let mut state = self.hashed_post_state.clone();
state.extend(hashed_state);
self.historical.hashed_state_root_from_nodes(trie_nodes, state, prefix_sets)
hashed_state.extend(state);
self.historical.hashed_state_root_from_nodes(trie_nodes, hashed_state, prefix_sets)
}
fn hashed_state_root_with_updates(
@ -130,14 +136,17 @@ impl StateRootProvider for MemoryOverlayStateProvider {
fn hashed_state_root_from_nodes_with_updates(
&self,
nodes: TrieUpdates,
hashed_state: HashedPostState,
state: HashedPostState,
prefix_sets: TriePrefixSetsMut,
) -> ProviderResult<(B256, TrieUpdates)> {
let mut trie_nodes = self.trie_updates.clone();
let MemoryOverlayTrieState { mut trie_nodes, mut hashed_state } = self.trie_state();
trie_nodes.extend(nodes);
let mut state = self.hashed_post_state.clone();
state.extend(hashed_state);
self.historical.hashed_state_root_from_nodes_with_updates(trie_nodes, state, prefix_sets)
hashed_state.extend(state);
self.historical.hashed_state_root_from_nodes_with_updates(
trie_nodes,
hashed_state,
prefix_sets,
)
}
// TODO: Currently this does not reuse available in-memory trie nodes.
@ -153,13 +162,13 @@ impl StateRootProvider for MemoryOverlayStateProvider {
impl StateProofProvider for MemoryOverlayStateProvider {
fn hashed_proof(
&self,
hashed_state: HashedPostState,
state: HashedPostState,
address: Address,
slots: &[B256],
) -> ProviderResult<AccountProof> {
let mut state = self.hashed_post_state.clone();
state.extend(hashed_state);
self.historical.hashed_proof(state, address, slots)
let MemoryOverlayTrieState { mut hashed_state, .. } = self.trie_state();
hashed_state.extend(state);
self.historical.hashed_proof(hashed_state, address, slots)
}
// TODO: Currently this does not reuse available in-memory trie nodes.
@ -168,9 +177,9 @@ impl StateProofProvider for MemoryOverlayStateProvider {
overlay: HashedPostState,
target: HashedPostState,
) -> ProviderResult<HashMap<B256, Bytes>> {
let mut state = self.hashed_post_state.clone();
state.extend(overlay);
self.historical.witness(state, target)
let MemoryOverlayTrieState { mut hashed_state, .. } = self.trie_state();
hashed_state.extend(overlay);
self.historical.witness(hashed_state, target)
}
}
@ -199,3 +208,12 @@ impl StateProvider for MemoryOverlayStateProvider {
self.historical.bytecode_by_hash(code_hash)
}
}
/// The collection of data necessary for trie-related operations for [`MemoryOverlayStateProvider`].
#[derive(Clone, Debug)]
pub(crate) struct MemoryOverlayTrieState {
/// The collection of aggregated in-memory trie updates.
pub(crate) trie_nodes: TrieUpdates,
/// The collection of hashed state from in-memory blocks.
pub(crate) hashed_state: HashedPostState,
}