1use crate::{errors::DatabaseError, precompiles::PrecompileCache};
2use ethrex_common::{
3 Address, H256, U256,
4 types::{AccountState, ChainConfig, Code, CodeMetadata},
5};
6#[cfg(all(feature = "rayon", not(feature = "eip-8025")))]
7use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
8use rustc_hash::FxHashMap;
9use std::sync::{Arc, OnceLock, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard};
10
11pub mod gen_db;
12
13type AccountCache = FxHashMap<Address, AccountState>;
15type StorageCache = FxHashMap<(Address, H256), U256>;
16type CodeCache = FxHashMap<H256, Code>;
17
18pub trait Database: Send + Sync {
19 fn get_account_state(&self, address: Address) -> Result<AccountState, DatabaseError>;
20 fn get_storage_value(&self, address: Address, key: H256) -> Result<U256, DatabaseError>;
21 fn get_block_hash(&self, block_number: u64) -> Result<H256, DatabaseError>;
22 fn get_chain_config(&self) -> Result<ChainConfig, DatabaseError>;
23 fn get_account_code(&self, code_hash: H256) -> Result<Code, DatabaseError>;
24 fn get_code_metadata(&self, code_hash: H256) -> Result<CodeMetadata, DatabaseError>;
25 fn precompile_cache(&self) -> Option<&PrecompileCache> {
27 None
28 }
29 fn prefetch_accounts(&self, addresses: &[Address]) -> Result<(), DatabaseError> {
31 for &addr in addresses {
32 self.get_account_state(addr)?;
33 }
34 Ok(())
35 }
36 fn prefetch_storage(&self, keys: &[(Address, H256)]) -> Result<(), DatabaseError> {
38 for &(addr, key) in keys {
39 self.get_storage_value(addr, key)?;
40 }
41 Ok(())
42 }
43}
44
45pub struct CachingDatabase {
55 inner: Arc<dyn Database>,
56 accounts: RwLock<AccountCache>,
58 storage: RwLock<StorageCache>,
60 code: RwLock<CodeCache>,
62 precompile_cache: Option<PrecompileCache>,
65 chain_config: OnceLock<ChainConfig>,
67}
68
69impl CachingDatabase {
70 pub fn new(inner: Arc<dyn Database>, precompile_cache_enabled: bool) -> Self {
71 Self {
72 inner,
73 accounts: RwLock::new(FxHashMap::default()),
74 storage: RwLock::new(FxHashMap::default()),
75 code: RwLock::new(FxHashMap::default()),
76 precompile_cache: precompile_cache_enabled.then(PrecompileCache::new),
77 chain_config: OnceLock::new(),
78 }
79 }
80
81 fn read_accounts(&self) -> Result<RwLockReadGuard<'_, AccountCache>, DatabaseError> {
82 self.accounts.read().map_err(poison_error_to_db_error)
83 }
84
85 fn write_accounts(&self) -> Result<RwLockWriteGuard<'_, AccountCache>, DatabaseError> {
86 self.accounts.write().map_err(poison_error_to_db_error)
87 }
88
89 fn read_storage(&self) -> Result<RwLockReadGuard<'_, StorageCache>, DatabaseError> {
90 self.storage.read().map_err(poison_error_to_db_error)
91 }
92
93 fn write_storage(&self) -> Result<RwLockWriteGuard<'_, StorageCache>, DatabaseError> {
94 self.storage.write().map_err(poison_error_to_db_error)
95 }
96
97 fn read_code(&self) -> Result<RwLockReadGuard<'_, CodeCache>, DatabaseError> {
98 self.code.read().map_err(poison_error_to_db_error)
99 }
100
101 fn write_code(&self) -> Result<RwLockWriteGuard<'_, CodeCache>, DatabaseError> {
102 self.code.write().map_err(poison_error_to_db_error)
103 }
104}
105
106fn poison_error_to_db_error<T>(err: PoisonError<T>) -> DatabaseError {
107 DatabaseError::Custom(format!("Cache lock poisoned: {err}"))
108}
109
110impl Database for CachingDatabase {
111 fn get_account_state(&self, address: Address) -> Result<AccountState, DatabaseError> {
112 if let Some(state) = self.read_accounts()?.get(&address).copied() {
114 return Ok(state);
115 }
116
117 let state = self.inner.get_account_state(address)?;
119
120 self.write_accounts()?.insert(address, state);
122
123 Ok(state)
124 }
125
126 fn get_storage_value(&self, address: Address, key: H256) -> Result<U256, DatabaseError> {
127 if let Some(value) = self.read_storage()?.get(&(address, key)).copied() {
129 return Ok(value);
130 }
131
132 let value = self.inner.get_storage_value(address, key)?;
134
135 self.write_storage()?.insert((address, key), value);
137
138 Ok(value)
139 }
140
141 fn get_block_hash(&self, block_number: u64) -> Result<H256, DatabaseError> {
142 self.inner.get_block_hash(block_number)
145 }
146
147 fn get_chain_config(&self) -> Result<ChainConfig, DatabaseError> {
148 if let Some(cfg) = self.chain_config.get() {
149 return Ok(*cfg);
150 }
151 let cfg = self.inner.get_chain_config()?;
152 let _ = self.chain_config.set(cfg);
154 Ok(*self.chain_config.get().unwrap_or(&cfg))
155 }
156
157 fn get_account_code(&self, code_hash: H256) -> Result<Code, DatabaseError> {
158 if let Some(code) = self.read_code()?.get(&code_hash).cloned() {
160 return Ok(code);
161 }
162
163 let code = self.inner.get_account_code(code_hash)?;
165
166 self.write_code()?.insert(code_hash, code.clone());
168
169 Ok(code)
170 }
171
172 fn get_code_metadata(&self, code_hash: H256) -> Result<CodeMetadata, DatabaseError> {
173 self.inner.get_code_metadata(code_hash)
177 }
178
179 fn precompile_cache(&self) -> Option<&PrecompileCache> {
180 self.precompile_cache.as_ref()
181 }
182
183 #[cfg(all(feature = "rayon", not(feature = "eip-8025")))]
184 fn prefetch_accounts(&self, addresses: &[Address]) -> Result<(), DatabaseError> {
185 let fetched: Vec<(Address, AccountState)> = addresses
187 .par_iter()
188 .map(|&addr| self.inner.get_account_state(addr).map(|s| (addr, s)))
189 .collect::<Result<_, _>>()?;
190 let mut cache = self.write_accounts()?;
191 for (addr, state) in fetched {
192 cache.entry(addr).or_insert(state);
193 }
194 Ok(())
195 }
196
197 #[cfg(all(feature = "rayon", not(feature = "eip-8025")))]
198 fn prefetch_storage(&self, keys: &[(Address, H256)]) -> Result<(), DatabaseError> {
199 let fetched: Vec<((Address, H256), U256)> = keys
201 .par_iter()
202 .map(|&(addr, key)| {
203 self.inner
204 .get_storage_value(addr, key)
205 .map(|v| ((addr, key), v))
206 })
207 .collect::<Result<_, _>>()?;
208 let mut cache = self.write_storage()?;
209 for (key, value) in fetched {
210 cache.entry(key).or_insert(value);
211 }
212 Ok(())
213 }
214}