use crate::{errors::DatabaseError, precompiles::PrecompileCache};
use ethrex_common::seeded_hash::{SeededHashMap, SeededHashSet};
use ethrex_common::{
Address, H256, U256,
types::{AccountState, ChainConfig, Code, CodeMetadata},
};
use std::sync::{
Arc, OnceLock, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard,
atomic::{AtomicU64, Ordering},
};
pub mod gen_db;
pub const BLOATED_BATCH_THRESHOLD: usize = 16_384;
type AccountCache = SeededHashMap<Address, AccountState>;
type StorageCache = SeededHashMap<(Address, H256), U256>;
type CodeCache = SeededHashMap<H256, Code>;
pub struct TouchedKeys {
pub accounts: Vec<(Address, H256)>,
pub slots: Vec<(Address, H256)>,
}
pub trait Database: Send + Sync {
fn get_account_state(&self, address: Address) -> Result<AccountState, DatabaseError>;
fn get_storage_value(&self, address: Address, key: H256) -> Result<U256, DatabaseError>;
fn get_block_hash(&self, block_number: u64) -> Result<H256, DatabaseError>;
fn get_chain_config(&self) -> Result<ChainConfig, DatabaseError>;
fn get_account_code(&self, code_hash: H256) -> Result<Code, DatabaseError>;
fn get_code_metadata(&self, code_hash: H256) -> Result<CodeMetadata, DatabaseError>;
fn precompile_cache(&self) -> Option<&PrecompileCache> {
None
}
fn get_account_states_batch(
&self,
addresses: &[Address],
) -> Result<Vec<AccountState>, DatabaseError> {
addresses
.iter()
.map(|a| self.get_account_state(*a))
.collect()
}
fn get_account_codes_batch(
&self,
code_hashes: &[H256],
) -> Result<Vec<Option<Code>>, DatabaseError> {
code_hashes
.iter()
.map(|h| self.get_account_code(*h).map(Some))
.collect()
}
fn get_storage_values_batch(
&self,
keys: &[(Address, H256)],
) -> Result<Vec<U256>, DatabaseError> {
keys.iter()
.map(|&(addr, key)| self.get_storage_value(addr, key))
.collect()
}
fn code_cache_budget_bytes(&self) -> u64 {
64 * 1024 * 1024
}
fn prefetch_codes(&self, code_hashes: &[H256]) -> Result<u64, DatabaseError> {
let mut warmed = 0u64;
for &hash in code_hashes {
if let Ok(code) = self.get_account_code(hash) {
warmed = warmed.saturating_add(u64::try_from(code.size()).unwrap_or(u64::MAX));
}
}
Ok(warmed)
}
fn prefetch_accounts(&self, addresses: &[Address]) -> Result<(), DatabaseError> {
for &addr in addresses {
self.get_account_state(addr)?;
}
Ok(())
}
fn prefetch_storage(&self, keys: &[(Address, H256)]) -> Result<(), DatabaseError> {
for &(addr, key) in keys {
self.get_storage_value(addr, key)?;
}
Ok(())
}
}
pub struct CachingDatabase {
inner: Arc<dyn Database>,
accounts: RwLock<AccountCache>,
storage: RwLock<StorageCache>,
code: RwLock<CodeCache>,
code_bytes: AtomicU64,
precompile_cache: Option<PrecompileCache>,
chain_config: OnceLock<ChainConfig>,
}
impl CachingDatabase {
pub fn new(inner: Arc<dyn Database>, precompile_cache_enabled: bool) -> Self {
Self {
inner,
accounts: RwLock::new(AccountCache::default()),
storage: RwLock::new(StorageCache::default()),
code: RwLock::new(CodeCache::default()),
code_bytes: AtomicU64::new(0),
precompile_cache: precompile_cache_enabled.then(PrecompileCache::new),
chain_config: OnceLock::new(),
}
}
fn read_accounts(&self) -> Result<RwLockReadGuard<'_, AccountCache>, DatabaseError> {
self.accounts.read().map_err(poison_error_to_db_error)
}
fn write_accounts(&self) -> Result<RwLockWriteGuard<'_, AccountCache>, DatabaseError> {
self.accounts.write().map_err(poison_error_to_db_error)
}
fn read_storage(&self) -> Result<RwLockReadGuard<'_, StorageCache>, DatabaseError> {
self.storage.read().map_err(poison_error_to_db_error)
}
fn write_storage(&self) -> Result<RwLockWriteGuard<'_, StorageCache>, DatabaseError> {
self.storage.write().map_err(poison_error_to_db_error)
}
pub fn code_bytes(&self) -> u64 {
self.code_bytes.load(Ordering::Relaxed)
}
fn keep_code(&self, cache: &mut CodeCache, code_hash: H256, code: Code) {
if let std::collections::hash_map::Entry::Vacant(entry) = cache.entry(code_hash) {
let size = u64::try_from(code.size()).unwrap_or(u64::MAX);
entry.insert(code);
self.code_bytes.fetch_add(size, Ordering::Relaxed);
}
}
fn read_code(&self) -> Result<RwLockReadGuard<'_, CodeCache>, DatabaseError> {
self.code.read().map_err(poison_error_to_db_error)
}
fn write_code(&self) -> Result<RwLockWriteGuard<'_, CodeCache>, DatabaseError> {
self.code.write().map_err(poison_error_to_db_error)
}
#[cfg(feature = "rayon")]
fn point_get_storage_many(
&self,
missing: &[(Address, H256)],
) -> Result<Vec<U256>, DatabaseError> {
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
missing
.par_iter()
.map(|&(addr, key)| self.inner.get_storage_value(addr, key))
.collect()
}
#[cfg(not(feature = "rayon"))]
fn point_get_storage_many(
&self,
missing: &[(Address, H256)],
) -> Result<Vec<U256>, DatabaseError> {
missing
.iter()
.map(|&(addr, key)| self.inner.get_storage_value(addr, key))
.collect()
}
pub fn touched_keys_where(
&self,
account_filter: &dyn Fn(&Address) -> bool,
slot_filter: &dyn Fn(&(Address, H256)) -> bool,
) -> TouchedKeys {
let accounts = self
.accounts
.read()
.map(|a| {
a.iter()
.filter(|(addr, _)| account_filter(addr))
.map(|(addr, st)| (*addr, st.storage_root))
.collect()
})
.unwrap_or_default();
let storage = self
.storage
.read()
.map(|s| s.keys().filter(|k| slot_filter(k)).copied().collect())
.unwrap_or_default();
TouchedKeys {
accounts,
slots: storage,
}
}
#[cfg(feature = "rayon")]
fn point_get_accounts_many(
&self,
missing: &[Address],
) -> Result<Vec<AccountState>, DatabaseError> {
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
missing
.par_iter()
.map(|&addr| self.inner.get_account_state(addr))
.collect()
}
#[cfg(not(feature = "rayon"))]
fn point_get_accounts_many(
&self,
missing: &[Address],
) -> Result<Vec<AccountState>, DatabaseError> {
missing
.iter()
.map(|&addr| self.inner.get_account_state(addr))
.collect()
}
}
fn poison_error_to_db_error<T>(err: PoisonError<T>) -> DatabaseError {
DatabaseError::Custom(format!("Cache lock poisoned: {err}"))
}
impl Database for CachingDatabase {
fn get_account_state(&self, address: Address) -> Result<AccountState, DatabaseError> {
if let Some(state) = self.read_accounts()?.get(&address).copied() {
return Ok(state);
}
let state = self.inner.get_account_state(address)?;
self.write_accounts()?.insert(address, state);
Ok(state)
}
fn get_storage_value(&self, address: Address, key: H256) -> Result<U256, DatabaseError> {
if let Some(value) = self.read_storage()?.get(&(address, key)).copied() {
return Ok(value);
}
let value = self.inner.get_storage_value(address, key)?;
self.write_storage()?.insert((address, key), value);
Ok(value)
}
fn get_block_hash(&self, block_number: u64) -> Result<H256, DatabaseError> {
self.inner.get_block_hash(block_number)
}
fn get_chain_config(&self) -> Result<ChainConfig, DatabaseError> {
if let Some(cfg) = self.chain_config.get() {
return Ok(*cfg);
}
let cfg = self.inner.get_chain_config()?;
let _ = self.chain_config.set(cfg);
Ok(*self.chain_config.get().unwrap_or(&cfg))
}
fn get_account_code(&self, code_hash: H256) -> Result<Code, DatabaseError> {
if let Some(code) = self.read_code()?.get(&code_hash).cloned() {
return Ok(code);
}
let code = self.inner.get_account_code(code_hash)?;
self.keep_code(&mut *self.write_code()?, code_hash, code.clone());
Ok(code)
}
fn get_code_metadata(&self, code_hash: H256) -> Result<CodeMetadata, DatabaseError> {
if let Some(code) = self.read_code()?.get(&code_hash) {
return Ok(CodeMetadata {
length: u64::try_from(code.len()).unwrap_or(u64::MAX),
});
}
self.inner.get_code_metadata(code_hash)
}
fn precompile_cache(&self) -> Option<&PrecompileCache> {
self.precompile_cache.as_ref()
}
fn get_account_states_batch(
&self,
addresses: &[Address],
) -> Result<Vec<AccountState>, DatabaseError> {
self.prefetch_accounts(addresses)?;
let cache = self.read_accounts()?;
addresses
.iter()
.map(|addr| {
cache.get(addr).copied().ok_or_else(|| {
DatabaseError::Custom(format!("account {addr:?} missing after prefetch"))
})
})
.collect()
}
fn code_cache_budget_bytes(&self) -> u64 {
self.inner.code_cache_budget_bytes()
}
fn prefetch_codes(&self, code_hashes: &[H256]) -> Result<u64, DatabaseError> {
let missing: Vec<H256> = {
let cache = self.read_code()?;
let mut seen: SeededHashSet<H256> = Default::default();
code_hashes
.iter()
.copied()
.filter(|h| !cache.contains_key(h) && seen.insert(*h))
.collect()
};
if missing.is_empty() {
return Ok(0);
}
let codes = self.inner.get_account_codes_batch(&missing)?;
let mut warmed = 0u64;
let mut cache = self.write_code()?;
for (hash, code) in missing.into_iter().zip(codes.into_iter()) {
if let Some(code) = code {
warmed = warmed.saturating_add(u64::try_from(code.size()).unwrap_or(u64::MAX));
self.keep_code(&mut cache, hash, code);
}
}
Ok(warmed)
}
fn prefetch_accounts(&self, addresses: &[Address]) -> Result<(), DatabaseError> {
let missing: Vec<Address> = {
let cache = self.read_accounts()?;
addresses
.iter()
.copied()
.filter(|a| !cache.contains_key(a))
.collect()
};
if missing.is_empty() {
return Ok(());
}
let states = if missing.len() >= BLOATED_BATCH_THRESHOLD {
self.inner.get_account_states_batch(&missing)?
} else {
self.point_get_accounts_many(&missing)?
};
let mut cache = self.write_accounts()?;
for (addr, state) in missing.into_iter().zip(states.into_iter()) {
cache.entry(addr).or_insert(state);
}
Ok(())
}
fn prefetch_storage(&self, keys: &[(Address, H256)]) -> Result<(), DatabaseError> {
let missing: Vec<(Address, H256)> = {
let cache = self.read_storage()?;
keys.iter()
.copied()
.filter(|k| !cache.contains_key(k))
.collect()
};
if missing.is_empty() {
return Ok(());
}
let values = if missing.len() >= BLOATED_BATCH_THRESHOLD {
self.inner.get_storage_values_batch(&missing)?
} else {
self.point_get_storage_many(&missing)?
};
let mut cache = self.write_storage()?;
for (key, value) in missing.into_iter().zip(values.into_iter()) {
cache.entry(key).or_insert(value);
}
Ok(())
}
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::indexing_slicing,
clippy::arithmetic_side_effects,
clippy::as_conversions
)]
mod code_bytes_tests {
use super::*;
use bytes::Bytes;
struct BigCodeStore;
fn code_for(hash: H256) -> Code {
let mut bytecode = vec![0x5b; 4096];
bytecode[..32].copy_from_slice(hash.as_bytes());
Code::from_bytecode_unchecked(Bytes::from(bytecode), hash)
}
impl Database for BigCodeStore {
fn get_account_state(&self, _: Address) -> Result<AccountState, DatabaseError> {
Ok(AccountState::default())
}
fn get_storage_value(&self, _: Address, _: H256) -> Result<U256, DatabaseError> {
Ok(U256::zero())
}
fn get_block_hash(&self, _: u64) -> Result<H256, DatabaseError> {
Ok(H256::zero())
}
fn get_chain_config(&self) -> Result<ChainConfig, DatabaseError> {
Ok(ChainConfig::default())
}
fn get_account_code(&self, code_hash: H256) -> Result<Code, DatabaseError> {
Ok(code_for(code_hash))
}
fn get_code_metadata(&self, code_hash: H256) -> Result<CodeMetadata, DatabaseError> {
Ok(CodeMetadata {
length: code_for(code_hash).len() as u64,
})
}
}
#[test]
fn code_bytes_counts_each_distinct_code_once() {
let cache = CachingDatabase::new(Arc::new(BigCodeStore), false);
let size = code_for(H256::zero()).size() as u64;
let hashes: Vec<H256> = (1..=4u64).map(H256::from_low_u64_be).collect();
cache.get_account_code(hashes[0]).unwrap();
cache.get_account_code(hashes[0]).unwrap();
assert_eq!(cache.code_bytes(), size);
cache.prefetch_codes(&hashes).unwrap();
cache.prefetch_codes(&hashes).unwrap();
assert_eq!(cache.code_bytes(), 4 * size);
assert_eq!(cache.read_code().unwrap().len(), 4);
}
}
#[cfg(test)]
#[allow(clippy::expect_used)]
mod seeded_cache_tests {
use super::*;
use std::hash::BuildHasher;
#[track_caller]
fn assert_seeded<S: BuildHasher + Default>(state: &S) {
let key = H256::repeat_byte(0x5a);
assert_ne!(state.hash_one(key), S::default().hash_one(key));
}
struct EmptyStore;
impl Database for EmptyStore {
fn get_account_state(&self, _: Address) -> Result<AccountState, DatabaseError> {
Ok(AccountState::default())
}
fn get_storage_value(&self, _: Address, _: H256) -> Result<U256, DatabaseError> {
Ok(U256::zero())
}
fn get_block_hash(&self, _: u64) -> Result<H256, DatabaseError> {
Ok(H256::zero())
}
fn get_chain_config(&self) -> Result<ChainConfig, DatabaseError> {
Ok(ChainConfig::default())
}
fn get_account_code(&self, _: H256) -> Result<Code, DatabaseError> {
Ok(Code::default())
}
fn get_code_metadata(&self, _: H256) -> Result<CodeMetadata, DatabaseError> {
Ok(CodeMetadata { length: 0 })
}
}
#[test]
fn warming_caches_are_seeded() {
let cache = CachingDatabase::new(Arc::new(EmptyStore), false);
assert_seeded(cache.read_accounts().expect("account cache lock").hasher());
assert_seeded(cache.read_storage().expect("storage cache lock").hasher());
assert_seeded(cache.read_code().expect("code cache lock").hasher());
}
}