1use crate::{errors::DatabaseError, precompiles::PrecompileCache};
2use ethrex_common::seeded_hash::{SeededHashMap, SeededHashSet};
3use ethrex_common::{
4 Address, H256, U256,
5 types::{AccountState, ChainConfig, Code, CodeMetadata},
6};
7use std::sync::{
8 Arc, OnceLock, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard,
9 atomic::{AtomicU64, Ordering},
10};
11
12pub mod gen_db;
13
14pub const BLOATED_BATCH_THRESHOLD: usize = 16_384;
21
22type AccountCache = SeededHashMap<Address, AccountState>;
26type StorageCache = SeededHashMap<(Address, H256), U256>;
27type CodeCache = SeededHashMap<H256, Code>;
28pub struct TouchedKeys {
30 pub accounts: Vec<(Address, H256)>,
32 pub slots: Vec<(Address, H256)>,
34}
35
36pub trait Database: Send + Sync {
37 fn get_account_state(&self, address: Address) -> Result<AccountState, DatabaseError>;
38 fn get_storage_value(&self, address: Address, key: H256) -> Result<U256, DatabaseError>;
39 fn get_block_hash(&self, block_number: u64) -> Result<H256, DatabaseError>;
40 fn get_chain_config(&self) -> Result<ChainConfig, DatabaseError>;
41 fn get_account_code(&self, code_hash: H256) -> Result<Code, DatabaseError>;
42 fn get_code_metadata(&self, code_hash: H256) -> Result<CodeMetadata, DatabaseError>;
43 fn precompile_cache(&self) -> Option<&PrecompileCache> {
45 None
46 }
47 fn get_account_states_batch(
51 &self,
52 addresses: &[Address],
53 ) -> Result<Vec<AccountState>, DatabaseError> {
54 addresses
55 .iter()
56 .map(|a| self.get_account_state(*a))
57 .collect()
58 }
59 fn get_account_codes_batch(
64 &self,
65 code_hashes: &[H256],
66 ) -> Result<Vec<Option<Code>>, DatabaseError> {
67 code_hashes
68 .iter()
69 .map(|h| self.get_account_code(*h).map(Some))
70 .collect()
71 }
72 fn get_storage_values_batch(
76 &self,
77 keys: &[(Address, H256)],
78 ) -> Result<Vec<U256>, DatabaseError> {
79 keys.iter()
80 .map(|&(addr, key)| self.get_storage_value(addr, key))
81 .collect()
82 }
83
84 fn code_cache_budget_bytes(&self) -> u64 {
91 64 * 1024 * 1024
92 }
93
94 fn prefetch_codes(&self, code_hashes: &[H256]) -> Result<u64, DatabaseError> {
100 let mut warmed = 0u64;
101 for &hash in code_hashes {
102 if let Ok(code) = self.get_account_code(hash) {
103 warmed = warmed.saturating_add(u64::try_from(code.size()).unwrap_or(u64::MAX));
104 }
105 }
106 Ok(warmed)
107 }
108 fn prefetch_accounts(&self, addresses: &[Address]) -> Result<(), DatabaseError> {
110 for &addr in addresses {
111 self.get_account_state(addr)?;
112 }
113 Ok(())
114 }
115 fn prefetch_storage(&self, keys: &[(Address, H256)]) -> Result<(), DatabaseError> {
117 for &(addr, key) in keys {
118 self.get_storage_value(addr, key)?;
119 }
120 Ok(())
121 }
122}
123
124pub struct CachingDatabase {
150 inner: Arc<dyn Database>,
151 accounts: RwLock<AccountCache>,
153 storage: RwLock<StorageCache>,
155 code: RwLock<CodeCache>,
157 code_bytes: AtomicU64,
161 precompile_cache: Option<PrecompileCache>,
164 chain_config: OnceLock<ChainConfig>,
166}
167
168impl CachingDatabase {
169 pub fn new(inner: Arc<dyn Database>, precompile_cache_enabled: bool) -> Self {
170 Self {
171 inner,
172 accounts: RwLock::new(AccountCache::default()),
173 storage: RwLock::new(StorageCache::default()),
174 code: RwLock::new(CodeCache::default()),
175 code_bytes: AtomicU64::new(0),
176 precompile_cache: precompile_cache_enabled.then(PrecompileCache::new),
177 chain_config: OnceLock::new(),
178 }
179 }
180
181 fn read_accounts(&self) -> Result<RwLockReadGuard<'_, AccountCache>, DatabaseError> {
182 self.accounts.read().map_err(poison_error_to_db_error)
183 }
184
185 fn write_accounts(&self) -> Result<RwLockWriteGuard<'_, AccountCache>, DatabaseError> {
186 self.accounts.write().map_err(poison_error_to_db_error)
187 }
188
189 fn read_storage(&self) -> Result<RwLockReadGuard<'_, StorageCache>, DatabaseError> {
190 self.storage.read().map_err(poison_error_to_db_error)
191 }
192
193 fn write_storage(&self) -> Result<RwLockWriteGuard<'_, StorageCache>, DatabaseError> {
194 self.storage.write().map_err(poison_error_to_db_error)
195 }
196
197 pub fn code_bytes(&self) -> u64 {
199 self.code_bytes.load(Ordering::Relaxed)
200 }
201
202 fn keep_code(&self, cache: &mut CodeCache, code_hash: H256, code: Code) {
204 if let std::collections::hash_map::Entry::Vacant(entry) = cache.entry(code_hash) {
205 let size = u64::try_from(code.size()).unwrap_or(u64::MAX);
206 entry.insert(code);
207 self.code_bytes.fetch_add(size, Ordering::Relaxed);
208 }
209 }
210
211 fn read_code(&self) -> Result<RwLockReadGuard<'_, CodeCache>, DatabaseError> {
212 self.code.read().map_err(poison_error_to_db_error)
213 }
214
215 fn write_code(&self) -> Result<RwLockWriteGuard<'_, CodeCache>, DatabaseError> {
216 self.code.write().map_err(poison_error_to_db_error)
217 }
218
219 #[cfg(feature = "rayon")]
223 fn point_get_storage_many(
224 &self,
225 missing: &[(Address, H256)],
226 ) -> Result<Vec<U256>, DatabaseError> {
227 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
228 missing
229 .par_iter()
230 .map(|&(addr, key)| self.inner.get_storage_value(addr, key))
231 .collect()
232 }
233
234 #[cfg(not(feature = "rayon"))]
235 fn point_get_storage_many(
236 &self,
237 missing: &[(Address, H256)],
238 ) -> Result<Vec<U256>, DatabaseError> {
239 missing
240 .iter()
241 .map(|&(addr, key)| self.inner.get_storage_value(addr, key))
242 .collect()
243 }
244
245 pub fn touched_keys_where(
251 &self,
252 account_filter: &dyn Fn(&Address) -> bool,
253 slot_filter: &dyn Fn(&(Address, H256)) -> bool,
254 ) -> TouchedKeys {
255 let accounts = self
256 .accounts
257 .read()
258 .map(|a| {
259 a.iter()
260 .filter(|(addr, _)| account_filter(addr))
261 .map(|(addr, st)| (*addr, st.storage_root))
262 .collect()
263 })
264 .unwrap_or_default();
265 let storage = self
266 .storage
267 .read()
268 .map(|s| s.keys().filter(|k| slot_filter(k)).copied().collect())
269 .unwrap_or_default();
270 TouchedKeys {
271 accounts,
272 slots: storage,
273 }
274 }
275
276 #[cfg(feature = "rayon")]
280 fn point_get_accounts_many(
281 &self,
282 missing: &[Address],
283 ) -> Result<Vec<AccountState>, DatabaseError> {
284 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
285 missing
286 .par_iter()
287 .map(|&addr| self.inner.get_account_state(addr))
288 .collect()
289 }
290
291 #[cfg(not(feature = "rayon"))]
292 fn point_get_accounts_many(
293 &self,
294 missing: &[Address],
295 ) -> Result<Vec<AccountState>, DatabaseError> {
296 missing
297 .iter()
298 .map(|&addr| self.inner.get_account_state(addr))
299 .collect()
300 }
301}
302
303fn poison_error_to_db_error<T>(err: PoisonError<T>) -> DatabaseError {
304 DatabaseError::Custom(format!("Cache lock poisoned: {err}"))
305}
306
307impl Database for CachingDatabase {
308 fn get_account_state(&self, address: Address) -> Result<AccountState, DatabaseError> {
309 if let Some(state) = self.read_accounts()?.get(&address).copied() {
311 return Ok(state);
312 }
313
314 let state = self.inner.get_account_state(address)?;
316
317 self.write_accounts()?.insert(address, state);
319
320 Ok(state)
321 }
322
323 fn get_storage_value(&self, address: Address, key: H256) -> Result<U256, DatabaseError> {
324 if let Some(value) = self.read_storage()?.get(&(address, key)).copied() {
326 return Ok(value);
327 }
328
329 let value = self.inner.get_storage_value(address, key)?;
331
332 self.write_storage()?.insert((address, key), value);
334
335 Ok(value)
336 }
337
338 fn get_block_hash(&self, block_number: u64) -> Result<H256, DatabaseError> {
339 self.inner.get_block_hash(block_number)
342 }
343
344 fn get_chain_config(&self) -> Result<ChainConfig, DatabaseError> {
345 if let Some(cfg) = self.chain_config.get() {
346 return Ok(*cfg);
347 }
348 let cfg = self.inner.get_chain_config()?;
349 let _ = self.chain_config.set(cfg);
351 Ok(*self.chain_config.get().unwrap_or(&cfg))
352 }
353
354 fn get_account_code(&self, code_hash: H256) -> Result<Code, DatabaseError> {
355 if let Some(code) = self.read_code()?.get(&code_hash).cloned() {
357 return Ok(code);
358 }
359
360 let code = self.inner.get_account_code(code_hash)?;
362
363 self.keep_code(&mut *self.write_code()?, code_hash, code.clone());
365
366 Ok(code)
367 }
368
369 fn get_code_metadata(&self, code_hash: H256) -> Result<CodeMetadata, DatabaseError> {
370 if let Some(code) = self.read_code()?.get(&code_hash) {
376 return Ok(CodeMetadata {
377 length: u64::try_from(code.len()).unwrap_or(u64::MAX),
378 });
379 }
380 self.inner.get_code_metadata(code_hash)
381 }
382
383 fn precompile_cache(&self) -> Option<&PrecompileCache> {
384 self.precompile_cache.as_ref()
385 }
386
387 fn get_account_states_batch(
395 &self,
396 addresses: &[Address],
397 ) -> Result<Vec<AccountState>, DatabaseError> {
398 self.prefetch_accounts(addresses)?;
399 let cache = self.read_accounts()?;
400 addresses
401 .iter()
402 .map(|addr| {
403 cache.get(addr).copied().ok_or_else(|| {
404 DatabaseError::Custom(format!("account {addr:?} missing after prefetch"))
405 })
406 })
407 .collect()
408 }
409
410 fn code_cache_budget_bytes(&self) -> u64 {
411 self.inner.code_cache_budget_bytes()
412 }
413
414 fn prefetch_codes(&self, code_hashes: &[H256]) -> Result<u64, DatabaseError> {
422 let missing: Vec<H256> = {
423 let cache = self.read_code()?;
424 let mut seen: SeededHashSet<H256> = Default::default();
425 code_hashes
426 .iter()
427 .copied()
428 .filter(|h| !cache.contains_key(h) && seen.insert(*h))
429 .collect()
430 };
431
432 if missing.is_empty() {
433 return Ok(0);
434 }
435
436 let codes = self.inner.get_account_codes_batch(&missing)?;
437 let mut warmed = 0u64;
438 let mut cache = self.write_code()?;
439 for (hash, code) in missing.into_iter().zip(codes.into_iter()) {
440 if let Some(code) = code {
442 warmed = warmed.saturating_add(u64::try_from(code.size()).unwrap_or(u64::MAX));
443 self.keep_code(&mut cache, hash, code);
444 }
445 }
446 Ok(warmed)
447 }
448
449 fn prefetch_accounts(&self, addresses: &[Address]) -> Result<(), DatabaseError> {
450 let missing: Vec<Address> = {
452 let cache = self.read_accounts()?;
453 addresses
454 .iter()
455 .copied()
456 .filter(|a| !cache.contains_key(a))
457 .collect()
458 };
459 if missing.is_empty() {
460 return Ok(());
461 }
462 let states = if missing.len() >= BLOATED_BATCH_THRESHOLD {
471 self.inner.get_account_states_batch(&missing)?
472 } else {
473 self.point_get_accounts_many(&missing)?
474 };
475 let mut cache = self.write_accounts()?;
476 for (addr, state) in missing.into_iter().zip(states.into_iter()) {
477 cache.entry(addr).or_insert(state);
478 }
479 Ok(())
480 }
481
482 fn prefetch_storage(&self, keys: &[(Address, H256)]) -> Result<(), DatabaseError> {
483 let missing: Vec<(Address, H256)> = {
485 let cache = self.read_storage()?;
486 keys.iter()
487 .copied()
488 .filter(|k| !cache.contains_key(k))
489 .collect()
490 };
491 if missing.is_empty() {
492 return Ok(());
493 }
494 let values = if missing.len() >= BLOATED_BATCH_THRESHOLD {
511 self.inner.get_storage_values_batch(&missing)?
516 } else {
517 self.point_get_storage_many(&missing)?
518 };
519 let mut cache = self.write_storage()?;
520 for (key, value) in missing.into_iter().zip(values.into_iter()) {
521 cache.entry(key).or_insert(value);
522 }
523 Ok(())
524 }
525}
526
527#[cfg(test)]
528#[allow(
529 clippy::unwrap_used,
530 clippy::indexing_slicing,
531 clippy::arithmetic_side_effects,
532 clippy::as_conversions
533)]
534mod code_bytes_tests {
535 use super::*;
536 use bytes::Bytes;
537
538 struct BigCodeStore;
540
541 fn code_for(hash: H256) -> Code {
542 let mut bytecode = vec![0x5b; 4096];
543 bytecode[..32].copy_from_slice(hash.as_bytes());
544 Code::from_bytecode_unchecked(Bytes::from(bytecode), hash)
545 }
546
547 impl Database for BigCodeStore {
548 fn get_account_state(&self, _: Address) -> Result<AccountState, DatabaseError> {
549 Ok(AccountState::default())
550 }
551 fn get_storage_value(&self, _: Address, _: H256) -> Result<U256, DatabaseError> {
552 Ok(U256::zero())
553 }
554 fn get_block_hash(&self, _: u64) -> Result<H256, DatabaseError> {
555 Ok(H256::zero())
556 }
557 fn get_chain_config(&self) -> Result<ChainConfig, DatabaseError> {
558 Ok(ChainConfig::default())
559 }
560 fn get_account_code(&self, code_hash: H256) -> Result<Code, DatabaseError> {
561 Ok(code_for(code_hash))
562 }
563 fn get_code_metadata(&self, code_hash: H256) -> Result<CodeMetadata, DatabaseError> {
564 Ok(CodeMetadata {
565 length: code_for(code_hash).len() as u64,
566 })
567 }
568 }
569
570 #[test]
573 fn code_bytes_counts_each_distinct_code_once() {
574 let cache = CachingDatabase::new(Arc::new(BigCodeStore), false);
575 let size = code_for(H256::zero()).size() as u64;
576 let hashes: Vec<H256> = (1..=4u64).map(H256::from_low_u64_be).collect();
577
578 cache.get_account_code(hashes[0]).unwrap();
579 cache.get_account_code(hashes[0]).unwrap();
580 assert_eq!(cache.code_bytes(), size);
581
582 cache.prefetch_codes(&hashes).unwrap();
583 cache.prefetch_codes(&hashes).unwrap();
584 assert_eq!(cache.code_bytes(), 4 * size);
585 assert_eq!(cache.read_code().unwrap().len(), 4);
586 }
587}
588
589#[cfg(test)]
590#[allow(clippy::expect_used)]
591mod seeded_cache_tests {
592 use super::*;
593 use std::hash::BuildHasher;
594
595 #[track_caller]
598 fn assert_seeded<S: BuildHasher + Default>(state: &S) {
599 let key = H256::repeat_byte(0x5a);
600 assert_ne!(state.hash_one(key), S::default().hash_one(key));
601 }
602
603 struct EmptyStore;
604
605 impl Database for EmptyStore {
606 fn get_account_state(&self, _: Address) -> Result<AccountState, DatabaseError> {
607 Ok(AccountState::default())
608 }
609 fn get_storage_value(&self, _: Address, _: H256) -> Result<U256, DatabaseError> {
610 Ok(U256::zero())
611 }
612 fn get_block_hash(&self, _: u64) -> Result<H256, DatabaseError> {
613 Ok(H256::zero())
614 }
615 fn get_chain_config(&self) -> Result<ChainConfig, DatabaseError> {
616 Ok(ChainConfig::default())
617 }
618 fn get_account_code(&self, _: H256) -> Result<Code, DatabaseError> {
619 Ok(Code::default())
620 }
621 fn get_code_metadata(&self, _: H256) -> Result<CodeMetadata, DatabaseError> {
622 Ok(CodeMetadata { length: 0 })
623 }
624 }
625
626 #[test]
629 fn warming_caches_are_seeded() {
630 let cache = CachingDatabase::new(Arc::new(EmptyStore), false);
631 assert_seeded(cache.read_accounts().expect("account cache lock").hasher());
632 assert_seeded(cache.read_storage().expect("storage cache lock").hasher());
633 assert_seeded(cache.read_code().expect("code cache lock").hasher());
634 }
635}