reth_engine_tree/tree/
cached_state.rs

1//! Execution cache implementation for block processing.
2use alloy_primitives::{Address, StorageKey, StorageValue, B256};
3use metrics::Gauge;
4use mini_moka::sync::CacheBuilder;
5use reth_errors::ProviderResult;
6use reth_metrics::Metrics;
7use reth_primitives_traits::{Account, Bytecode};
8use reth_provider::{
9    AccountReader, BlockHashReader, BytecodeReader, HashedPostStateProvider, StateProofProvider,
10    StateProvider, StateRootProvider, StorageRootProvider,
11};
12use reth_revm::db::BundleState;
13use reth_trie::{
14    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
15    MultiProofTargets, StorageMultiProof, StorageProof, TrieInput,
16};
17use revm_primitives::map::DefaultHashBuilder;
18use std::{sync::Arc, time::Duration};
19use tracing::{debug_span, instrument, trace};
20
21pub(crate) type Cache<K, V> =
22    mini_moka::sync::Cache<K, V, alloy_primitives::map::DefaultHashBuilder>;
23
24/// A wrapper of a state provider and a shared cache.
25pub(crate) struct CachedStateProvider<S> {
26    /// The state provider
27    state_provider: S,
28
29    /// The caches used for the provider
30    caches: ExecutionCache,
31
32    /// Metrics for the cached state provider
33    metrics: CachedStateMetrics,
34}
35
36impl<S> CachedStateProvider<S>
37where
38    S: StateProvider,
39{
40    /// Creates a new [`CachedStateProvider`] from an [`ExecutionCache`], state provider, and
41    /// [`CachedStateMetrics`].
42    pub(crate) const fn new_with_caches(
43        state_provider: S,
44        caches: ExecutionCache,
45        metrics: CachedStateMetrics,
46    ) -> Self {
47        Self { state_provider, caches, metrics }
48    }
49}
50
51/// Metrics for the cached state provider, showing hits / misses for each cache
52#[derive(Metrics, Clone)]
53#[metrics(scope = "sync.caching")]
54pub(crate) struct CachedStateMetrics {
55    /// Code cache hits
56    code_cache_hits: Gauge,
57
58    /// Code cache misses
59    code_cache_misses: Gauge,
60
61    /// Code cache size
62    ///
63    /// NOTE: this uses the moka caches' `entry_count`, NOT the `weighted_size` method to calculate
64    /// size.
65    code_cache_size: Gauge,
66
67    /// Storage cache hits
68    storage_cache_hits: Gauge,
69
70    /// Storage cache misses
71    storage_cache_misses: Gauge,
72
73    /// Storage cache size
74    ///
75    /// NOTE: this uses the moka caches' `entry_count`, NOT the `weighted_size` method to calculate
76    /// size.
77    storage_cache_size: Gauge,
78
79    /// Account cache hits
80    account_cache_hits: Gauge,
81
82    /// Account cache misses
83    account_cache_misses: Gauge,
84
85    /// Account cache size
86    ///
87    /// NOTE: this uses the moka caches' `entry_count`, NOT the `weighted_size` method to calculate
88    /// size.
89    account_cache_size: Gauge,
90}
91
92impl CachedStateMetrics {
93    /// Sets all values to zero, indicating that a new block is being executed.
94    pub(crate) fn reset(&self) {
95        // code cache
96        self.code_cache_hits.set(0);
97        self.code_cache_misses.set(0);
98
99        // storage cache
100        self.storage_cache_hits.set(0);
101        self.storage_cache_misses.set(0);
102
103        // account cache
104        self.account_cache_hits.set(0);
105        self.account_cache_misses.set(0);
106    }
107
108    /// Returns a new zeroed-out instance of [`CachedStateMetrics`].
109    pub(crate) fn zeroed() -> Self {
110        let zeroed = Self::default();
111        zeroed.reset();
112        zeroed
113    }
114}
115
116impl<S: AccountReader> AccountReader for CachedStateProvider<S> {
117    fn basic_account(&self, address: &Address) -> ProviderResult<Option<Account>> {
118        if let Some(res) = self.caches.account_cache.get(address) {
119            self.metrics.account_cache_hits.increment(1);
120            return Ok(res)
121        }
122
123        self.metrics.account_cache_misses.increment(1);
124
125        let res = self.state_provider.basic_account(address)?;
126        self.caches.account_cache.insert(*address, res);
127        Ok(res)
128    }
129}
130
131/// Represents the status of a storage slot in the cache.
132#[derive(Debug, Clone, PartialEq, Eq)]
133pub(crate) enum SlotStatus {
134    /// The account's storage cache doesn't exist.
135    NotCached,
136    /// The storage slot exists in cache and is empty (value is zero).
137    Empty,
138    /// The storage slot exists in cache and has a specific non-zero value.
139    Value(StorageValue),
140}
141
142impl<S: StateProvider> StateProvider for CachedStateProvider<S> {
143    fn storage(
144        &self,
145        account: Address,
146        storage_key: StorageKey,
147    ) -> ProviderResult<Option<StorageValue>> {
148        match self.caches.get_storage(&account, &storage_key) {
149            (SlotStatus::NotCached, maybe_cache) => {
150                let final_res = self.state_provider.storage(account, storage_key)?;
151                let account_cache = maybe_cache.unwrap_or_default();
152                account_cache.insert_storage(storage_key, final_res);
153                // we always need to insert the value to update the weights.
154                // Note: there exists a race when the storage cache did not exist yet and two
155                // consumers looking up the a storage value for this account for the first time,
156                // however we can assume that this will only happen for the very first (mostlikely
157                // the same) value, and don't expect that this will accidentally
158                // replace an account storage cache with additional values.
159                self.caches.insert_storage_cache(account, account_cache);
160
161                self.metrics.storage_cache_misses.increment(1);
162                Ok(final_res)
163            }
164            (SlotStatus::Empty, _) => {
165                self.metrics.storage_cache_hits.increment(1);
166                Ok(None)
167            }
168            (SlotStatus::Value(value), _) => {
169                self.metrics.storage_cache_hits.increment(1);
170                Ok(Some(value))
171            }
172        }
173    }
174}
175
176impl<S: BytecodeReader> BytecodeReader for CachedStateProvider<S> {
177    fn bytecode_by_hash(&self, code_hash: &B256) -> ProviderResult<Option<Bytecode>> {
178        if let Some(res) = self.caches.code_cache.get(code_hash) {
179            self.metrics.code_cache_hits.increment(1);
180            return Ok(res)
181        }
182
183        self.metrics.code_cache_misses.increment(1);
184
185        let final_res = self.state_provider.bytecode_by_hash(code_hash)?;
186        self.caches.code_cache.insert(*code_hash, final_res.clone());
187        Ok(final_res)
188    }
189}
190
191impl<S: StateRootProvider> StateRootProvider for CachedStateProvider<S> {
192    fn state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
193        self.state_provider.state_root(hashed_state)
194    }
195
196    fn state_root_from_nodes(&self, input: TrieInput) -> ProviderResult<B256> {
197        self.state_provider.state_root_from_nodes(input)
198    }
199
200    fn state_root_with_updates(
201        &self,
202        hashed_state: HashedPostState,
203    ) -> ProviderResult<(B256, TrieUpdates)> {
204        self.state_provider.state_root_with_updates(hashed_state)
205    }
206
207    fn state_root_from_nodes_with_updates(
208        &self,
209        input: TrieInput,
210    ) -> ProviderResult<(B256, TrieUpdates)> {
211        self.state_provider.state_root_from_nodes_with_updates(input)
212    }
213}
214
215impl<S: StateProofProvider> StateProofProvider for CachedStateProvider<S> {
216    fn proof(
217        &self,
218        input: TrieInput,
219        address: Address,
220        slots: &[B256],
221    ) -> ProviderResult<AccountProof> {
222        self.state_provider.proof(input, address, slots)
223    }
224
225    fn multiproof(
226        &self,
227        input: TrieInput,
228        targets: MultiProofTargets,
229    ) -> ProviderResult<MultiProof> {
230        self.state_provider.multiproof(input, targets)
231    }
232
233    fn witness(
234        &self,
235        input: TrieInput,
236        target: HashedPostState,
237    ) -> ProviderResult<Vec<alloy_primitives::Bytes>> {
238        self.state_provider.witness(input, target)
239    }
240}
241
242impl<S: StorageRootProvider> StorageRootProvider for CachedStateProvider<S> {
243    fn storage_root(
244        &self,
245        address: Address,
246        hashed_storage: HashedStorage,
247    ) -> ProviderResult<B256> {
248        self.state_provider.storage_root(address, hashed_storage)
249    }
250
251    fn storage_proof(
252        &self,
253        address: Address,
254        slot: B256,
255        hashed_storage: HashedStorage,
256    ) -> ProviderResult<StorageProof> {
257        self.state_provider.storage_proof(address, slot, hashed_storage)
258    }
259
260    /// Generate a storage multiproof for multiple storage slots.
261    ///
262    /// A **storage multiproof** is a cryptographic proof that can verify the values
263    /// of multiple storage slots for a single account in a single verification step.
264    /// Instead of generating separate proofs for each slot (which would be inefficient),
265    /// a multiproof bundles the necessary trie nodes to prove all requested slots.
266    ///
267    /// ## How it works:
268    /// 1. Takes an account address and a list of storage slot keys
269    /// 2. Traverses the account's storage trie to collect proof nodes
270    /// 3. Returns a [`StorageMultiProof`] containing the minimal set of trie nodes needed to verify
271    ///    all the requested storage slots
272    fn storage_multiproof(
273        &self,
274        address: Address,
275        slots: &[B256],
276        hashed_storage: HashedStorage,
277    ) -> ProviderResult<StorageMultiProof> {
278        self.state_provider.storage_multiproof(address, slots, hashed_storage)
279    }
280}
281
282impl<S: BlockHashReader> BlockHashReader for CachedStateProvider<S> {
283    fn block_hash(&self, number: alloy_primitives::BlockNumber) -> ProviderResult<Option<B256>> {
284        self.state_provider.block_hash(number)
285    }
286
287    fn canonical_hashes_range(
288        &self,
289        start: alloy_primitives::BlockNumber,
290        end: alloy_primitives::BlockNumber,
291    ) -> ProviderResult<Vec<B256>> {
292        self.state_provider.canonical_hashes_range(start, end)
293    }
294}
295
296impl<S: HashedPostStateProvider> HashedPostStateProvider for CachedStateProvider<S> {
297    fn hashed_post_state(&self, bundle_state: &reth_revm::db::BundleState) -> HashedPostState {
298        self.state_provider.hashed_post_state(bundle_state)
299    }
300}
301
302/// Execution cache used during block processing.
303///
304/// Optimizes state access by maintaining in-memory copies of frequently accessed
305/// accounts, storage slots, and bytecode. Works in conjunction with prewarming
306/// to reduce database I/O during block execution.
307#[derive(Debug, Clone)]
308pub(crate) struct ExecutionCache {
309    /// Cache for contract bytecode, keyed by code hash.
310    code_cache: Cache<B256, Option<Bytecode>>,
311
312    /// Per-account storage cache: outer cache keyed by Address, inner cache tracks that account’s
313    /// storage slots.
314    storage_cache: Cache<Address, Arc<AccountStorageCache>>,
315
316    /// Cache for basic account information (nonce, balance, code hash).
317    account_cache: Cache<Address, Option<Account>>,
318}
319
320impl ExecutionCache {
321    /// Get storage value from hierarchical cache.
322    ///
323    /// Returns a tuple of:
324    /// - `SlotStatus` indicating whether:
325    ///   - `NotCached`: The account's storage cache doesn't exist
326    ///   - `Empty`: The slot exists in the account's cache but is empty
327    ///   - `Value`: The slot exists and has a specific value
328    /// - `Option<Arc<AccountStorageCache>>`: The account's storage cache if it exists
329    pub(crate) fn get_storage(
330        &self,
331        address: &Address,
332        key: &StorageKey,
333    ) -> (SlotStatus, Option<Arc<AccountStorageCache>>) {
334        match self.storage_cache.get(address) {
335            None => (SlotStatus::NotCached, None),
336            Some(account_cache) => {
337                let status = account_cache.get_storage(key);
338                (status, Some(account_cache))
339            }
340        }
341    }
342
343    /// Insert storage value into hierarchical cache
344    #[cfg(test)]
345    pub(crate) fn insert_storage(
346        &self,
347        address: Address,
348        key: StorageKey,
349        value: Option<StorageValue>,
350    ) {
351        self.insert_storage_bulk(address, [(key, value)]);
352    }
353
354    /// Insert multiple storage values into hierarchical cache for a single account
355    ///
356    /// This method is optimized for inserting multiple storage values for the same address
357    /// by doing the account cache lookup only once instead of for each key-value pair.
358    pub(crate) fn insert_storage_bulk<I>(&self, address: Address, storage_entries: I)
359    where
360        I: IntoIterator<Item = (StorageKey, Option<StorageValue>)>,
361    {
362        let account_cache = self.storage_cache.get(&address).unwrap_or_default();
363
364        for (key, value) in storage_entries {
365            account_cache.insert_storage(key, value);
366        }
367
368        // Insert to the cache so that moka picks up on the changed size, even though the actual
369        // value (the Arc<AccountStorageCache>) is the same
370        self.storage_cache.insert(address, account_cache);
371    }
372
373    /// Inserts the [`AccountStorageCache`].
374    pub(crate) fn insert_storage_cache(
375        &self,
376        address: Address,
377        storage_cache: Arc<AccountStorageCache>,
378    ) {
379        self.storage_cache.insert(address, storage_cache);
380    }
381
382    /// Invalidate storage for specific account
383    pub(crate) fn invalidate_account_storage(&self, address: &Address) {
384        self.storage_cache.invalidate(address);
385    }
386
387    /// Returns the total number of storage slots cached across all accounts
388    pub(crate) fn total_storage_slots(&self) -> usize {
389        self.storage_cache.iter().map(|addr| addr.len()).sum()
390    }
391
392    /// Inserts the post-execution state changes into the cache.
393    ///
394    /// This method is called after transaction execution to update the cache with
395    /// the touched and modified state. The insertion order is critical:
396    ///
397    /// 1. Bytecodes: Insert contract code first
398    /// 2. Storage slots: Update storage values for each account
399    /// 3. Accounts: Update account info (nonce, balance, code hash)
400    ///
401    /// ## Why This Order Matters
402    ///
403    /// Account information references bytecode via code hash. If we update accounts
404    /// before bytecode, we might create cache entries pointing to non-existent code.
405    /// The current order ensures cache consistency.
406    ///
407    /// ## Error Handling
408    ///
409    /// Returns an error if the state updates are inconsistent and should be discarded.
410    #[instrument(level = "debug", target = "engine::caching", skip_all)]
411    pub(crate) fn insert_state(&self, state_updates: &BundleState) -> Result<(), ()> {
412        let _enter =
413            debug_span!(target: "engine::tree", "contracts", len = state_updates.contracts.len())
414                .entered();
415        // Insert bytecodes
416        for (code_hash, bytecode) in &state_updates.contracts {
417            self.code_cache.insert(*code_hash, Some(Bytecode(bytecode.clone())));
418        }
419        drop(_enter);
420
421        let _enter = debug_span!(
422            target: "engine::tree",
423            "accounts",
424            accounts = state_updates.state.len(),
425            storages =
426                state_updates.state.values().map(|account| account.storage.len()).sum::<usize>()
427        )
428        .entered();
429        for (addr, account) in &state_updates.state {
430            // If the account was not modified, as in not changed and not destroyed, then we have
431            // nothing to do w.r.t. this particular account and can move on
432            if account.status.is_not_modified() {
433                continue
434            }
435
436            // If the account was destroyed, invalidate from the account / storage caches
437            if account.was_destroyed() {
438                // Invalidate the account cache entry if destroyed
439                self.account_cache.invalidate(addr);
440
441                self.invalidate_account_storage(addr);
442                continue
443            }
444
445            // If we have an account that was modified, but it has a `None` account info, some wild
446            // error has occurred because this state should be unrepresentable. An account with
447            // `None` current info, should be destroyed.
448            let Some(ref account_info) = account.info else {
449                trace!(target: "engine::caching", ?account, "Account with None account info found in state updates");
450                return Err(())
451            };
452
453            // Now we iterate over all storage and make updates to the cached storage values
454            // Use bulk insertion to optimize cache lookups - only lookup the account cache once
455            // instead of for each storage key
456            let storage_entries = account.storage.iter().map(|(storage_key, slot)| {
457                // We convert the storage key from U256 to B256 because that is how it's represented
458                // in the cache
459                ((*storage_key).into(), Some(slot.present_value))
460            });
461            self.insert_storage_bulk(*addr, storage_entries);
462
463            // Insert will update if present, so we just use the new account info as the new value
464            // for the account cache
465            self.account_cache.insert(*addr, Some(Account::from(account_info)));
466        }
467
468        Ok(())
469    }
470}
471
472/// A builder for [`ExecutionCache`].
473#[derive(Debug)]
474pub(crate) struct ExecutionCacheBuilder {
475    /// Code cache entries
476    code_cache_entries: u64,
477
478    /// Storage cache entries
479    storage_cache_entries: u64,
480
481    /// Account cache entries
482    account_cache_entries: u64,
483}
484
485impl ExecutionCacheBuilder {
486    /// Build an [`ExecutionCache`] struct, so that execution caches can be easily cloned.
487    pub(crate) fn build_caches(self, total_cache_size: u64) -> ExecutionCache {
488        let storage_cache_size = (total_cache_size * 8888) / 10000; // 88.88% of total
489        let account_cache_size = (total_cache_size * 556) / 10000; // 5.56% of total
490        let code_cache_size = (total_cache_size * 556) / 10000; // 5.56% of total
491
492        const EXPIRY_TIME: Duration = Duration::from_secs(7200); // 2 hours
493        const TIME_TO_IDLE: Duration = Duration::from_secs(3600); // 1 hour
494
495        let storage_cache = CacheBuilder::new(self.storage_cache_entries)
496            .weigher(|_key: &Address, value: &Arc<AccountStorageCache>| -> u32 {
497                // values based on results from measure_storage_cache_overhead test
498                let base_weight = 39_000;
499                let slots_weight = value.len() * 218;
500                (base_weight + slots_weight) as u32
501            })
502            .max_capacity(storage_cache_size)
503            .time_to_live(EXPIRY_TIME)
504            .time_to_idle(TIME_TO_IDLE)
505            .build_with_hasher(DefaultHashBuilder::default());
506
507        let account_cache = CacheBuilder::new(self.account_cache_entries)
508            .weigher(|_key: &Address, value: &Option<Account>| -> u32 {
509                // Account has a fixed size (none, balance,code_hash)
510                20 + size_of_val(value) as u32
511            })
512            .max_capacity(account_cache_size)
513            .time_to_live(EXPIRY_TIME)
514            .time_to_idle(TIME_TO_IDLE)
515            .build_with_hasher(DefaultHashBuilder::default());
516
517        let code_cache = CacheBuilder::new(self.code_cache_entries)
518            .weigher(|_key: &B256, value: &Option<Bytecode>| -> u32 {
519                let code_size = match value {
520                    Some(bytecode) => {
521                        // base weight + actual (padded) bytecode size + size of the jump table
522                        (size_of_val(value) +
523                            bytecode.bytecode().len() +
524                            bytecode
525                                .legacy_jump_table()
526                                .map(|table| table.as_slice().len())
527                                .unwrap_or_default()) as u32
528                    }
529                    None => size_of_val(value) as u32,
530                };
531                32 + code_size
532            })
533            .max_capacity(code_cache_size)
534            .time_to_live(EXPIRY_TIME)
535            .time_to_idle(TIME_TO_IDLE)
536            .build_with_hasher(DefaultHashBuilder::default());
537
538        ExecutionCache { code_cache, storage_cache, account_cache }
539    }
540}
541
542impl Default for ExecutionCacheBuilder {
543    fn default() -> Self {
544        // With weigher and max_capacity in place, these numbers represent
545        // the maximum number of entries that can be stored, not the actual
546        // memory usage which is controlled by max_capacity.
547        //
548        // Code cache: up to 10M entries but limited to 0.5GB
549        // Storage cache: up to 10M accounts but limited to 8GB
550        // Account cache: up to 10M accounts but limited to 0.5GB
551        Self {
552            code_cache_entries: 10_000_000,
553            storage_cache_entries: 10_000_000,
554            account_cache_entries: 10_000_000,
555        }
556    }
557}
558
559/// A saved cache that has been used for executing a specific block, which has been updated for its
560/// execution.
561#[derive(Debug, Clone)]
562pub(crate) struct SavedCache {
563    /// The hash of the block these caches were used to execute.
564    hash: B256,
565
566    /// The caches used for the provider.
567    caches: ExecutionCache,
568
569    /// Metrics for the cached state provider
570    metrics: CachedStateMetrics,
571
572    /// A guard to track in-flight usage of this cache.
573    /// The cache is considered available if the strong count is 1.
574    usage_guard: Arc<()>,
575}
576
577impl SavedCache {
578    /// Creates a new instance with the internals
579    pub(super) fn new(hash: B256, caches: ExecutionCache, metrics: CachedStateMetrics) -> Self {
580        Self { hash, caches, metrics, usage_guard: Arc::new(()) }
581    }
582
583    /// Returns the hash for this cache
584    pub(crate) const fn executed_block_hash(&self) -> B256 {
585        self.hash
586    }
587
588    /// Splits the cache into its caches and metrics, consuming it.
589    pub(crate) fn split(self) -> (ExecutionCache, CachedStateMetrics) {
590        (self.caches, self.metrics)
591    }
592
593    /// Returns true if the cache is available for use (no other tasks are currently using it).
594    pub(crate) fn is_available(&self) -> bool {
595        Arc::strong_count(&self.usage_guard) == 1
596    }
597
598    /// Returns the [`ExecutionCache`] belonging to the tracked hash.
599    pub(crate) const fn cache(&self) -> &ExecutionCache {
600        &self.caches
601    }
602
603    /// Returns the metrics associated with this cache.
604    pub(crate) const fn metrics(&self) -> &CachedStateMetrics {
605        &self.metrics
606    }
607
608    /// Updates the metrics for the [`ExecutionCache`].
609    pub(crate) fn update_metrics(&self) {
610        self.metrics.storage_cache_size.set(self.caches.total_storage_slots() as f64);
611        self.metrics.account_cache_size.set(self.caches.account_cache.entry_count() as f64);
612        self.metrics.code_cache_size.set(self.caches.code_cache.entry_count() as f64);
613    }
614}
615
616#[cfg(test)]
617impl SavedCache {
618    fn clone_guard_for_test(&self) -> Arc<()> {
619        self.usage_guard.clone()
620    }
621}
622
623/// Cache for an individual account's storage slots.
624///
625/// This represents the second level of the hierarchical storage cache.
626/// Each account gets its own `AccountStorageCache` to store accessed storage slots.
627#[derive(Debug, Clone)]
628pub(crate) struct AccountStorageCache {
629    /// Map of storage keys to their cached values.
630    slots: Cache<StorageKey, Option<StorageValue>>,
631}
632
633impl AccountStorageCache {
634    /// Create a new [`AccountStorageCache`]
635    pub(crate) fn new(max_slots: u64) -> Self {
636        Self {
637            slots: CacheBuilder::new(max_slots).build_with_hasher(DefaultHashBuilder::default()),
638        }
639    }
640
641    /// Get a storage value from this account's cache.
642    /// - `NotCached`: The slot is not in the cache
643    /// - `Empty`: The slot is empty
644    /// - `Value`: The slot has a specific value
645    pub(crate) fn get_storage(&self, key: &StorageKey) -> SlotStatus {
646        match self.slots.get(key) {
647            None => SlotStatus::NotCached,
648            Some(None) => SlotStatus::Empty,
649            Some(Some(value)) => SlotStatus::Value(value),
650        }
651    }
652
653    /// Insert a storage value
654    pub(crate) fn insert_storage(&self, key: StorageKey, value: Option<StorageValue>) {
655        self.slots.insert(key, value);
656    }
657
658    /// Returns the number of slots in the cache
659    pub(crate) fn len(&self) -> usize {
660        self.slots.entry_count() as usize
661    }
662}
663
664impl Default for AccountStorageCache {
665    fn default() -> Self {
666        // With weigher and max_capacity in place, this number represents
667        // the maximum number of entries that can be stored, not the actual
668        // memory usage which is controlled by storage cache's max_capacity.
669        Self::new(1_000_000)
670    }
671}
672
673#[cfg(test)]
674mod tests {
675    use super::*;
676    use alloy_primitives::{B256, U256};
677    use rand::Rng;
678    use reth_provider::test_utils::{ExtendedAccount, MockEthProvider};
679    use std::mem::size_of;
680
681    mod tracking_allocator {
682        use std::{
683            alloc::{GlobalAlloc, Layout, System},
684            sync::atomic::{AtomicUsize, Ordering},
685        };
686
687        #[derive(Debug)]
688        pub(crate) struct TrackingAllocator {
689            allocated: AtomicUsize,
690            total_allocated: AtomicUsize,
691            inner: System,
692        }
693
694        impl TrackingAllocator {
695            pub(crate) const fn new() -> Self {
696                Self {
697                    allocated: AtomicUsize::new(0),
698                    total_allocated: AtomicUsize::new(0),
699                    inner: System,
700                }
701            }
702
703            pub(crate) fn reset(&self) {
704                self.allocated.store(0, Ordering::SeqCst);
705                self.total_allocated.store(0, Ordering::SeqCst);
706            }
707
708            pub(crate) fn total_allocated(&self) -> usize {
709                self.total_allocated.load(Ordering::SeqCst)
710            }
711        }
712
713        unsafe impl GlobalAlloc for TrackingAllocator {
714            unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
715                let ret = unsafe { self.inner.alloc(layout) };
716                if !ret.is_null() {
717                    self.allocated.fetch_add(layout.size(), Ordering::SeqCst);
718                    self.total_allocated.fetch_add(layout.size(), Ordering::SeqCst);
719                }
720                ret
721            }
722
723            unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
724                self.allocated.fetch_sub(layout.size(), Ordering::SeqCst);
725                unsafe { self.inner.dealloc(ptr, layout) }
726            }
727        }
728    }
729
730    use tracking_allocator::TrackingAllocator;
731
732    #[global_allocator]
733    static ALLOCATOR: TrackingAllocator = TrackingAllocator::new();
734
735    fn measure_allocation<T, F>(f: F) -> (usize, T)
736    where
737        F: FnOnce() -> T,
738    {
739        ALLOCATOR.reset();
740        let result = f();
741        let total = ALLOCATOR.total_allocated();
742        (total, result)
743    }
744
745    #[test]
746    fn measure_storage_cache_overhead() {
747        let (base_overhead, cache) = measure_allocation(|| AccountStorageCache::new(1000));
748        println!("Base AccountStorageCache overhead: {base_overhead} bytes");
749        let mut rng = rand::rng();
750
751        let key = StorageKey::random();
752        let value = StorageValue::from(rng.random::<u128>());
753        let (first_slot, _) = measure_allocation(|| {
754            cache.insert_storage(key, Some(value));
755        });
756        println!("First slot insertion overhead: {first_slot} bytes");
757
758        const TOTAL_SLOTS: usize = 10_000;
759        let (test_slots, _) = measure_allocation(|| {
760            for _ in 0..TOTAL_SLOTS {
761                let key = StorageKey::random();
762                let value = StorageValue::from(rng.random::<u128>());
763                cache.insert_storage(key, Some(value));
764            }
765        });
766        println!("Average overhead over {} slots: {} bytes", TOTAL_SLOTS, test_slots / TOTAL_SLOTS);
767
768        println!("\nTheoretical sizes:");
769        println!("StorageKey size: {} bytes", size_of::<StorageKey>());
770        println!("StorageValue size: {} bytes", size_of::<StorageValue>());
771        println!("Option<StorageValue> size: {} bytes", size_of::<Option<StorageValue>>());
772        println!("Option<B256> size: {} bytes", size_of::<Option<B256>>());
773    }
774
775    #[test]
776    fn test_empty_storage_cached_state_provider() {
777        // make sure when we have an empty value in storage, we return `Empty` and not `NotCached`
778        let address = Address::random();
779        let storage_key = StorageKey::random();
780        let account = ExtendedAccount::new(0, U256::ZERO);
781
782        // note there is no storage here
783        let provider = MockEthProvider::default();
784        provider.extend_accounts(vec![(address, account)]);
785
786        let caches = ExecutionCacheBuilder::default().build_caches(1000);
787        let state_provider =
788            CachedStateProvider::new_with_caches(provider, caches, CachedStateMetrics::zeroed());
789
790        // check that the storage is empty
791        let res = state_provider.storage(address, storage_key);
792        assert!(res.is_ok());
793        assert_eq!(res.unwrap(), None);
794    }
795
796    #[test]
797    fn test_uncached_storage_cached_state_provider() {
798        // make sure when we have something uncached, we get the cached value
799        let address = Address::random();
800        let storage_key = StorageKey::random();
801        let storage_value = U256::from(1);
802        let account =
803            ExtendedAccount::new(0, U256::ZERO).extend_storage(vec![(storage_key, storage_value)]);
804
805        // note that we extend storage here with one value
806        let provider = MockEthProvider::default();
807        provider.extend_accounts(vec![(address, account)]);
808
809        let caches = ExecutionCacheBuilder::default().build_caches(1000);
810        let state_provider =
811            CachedStateProvider::new_with_caches(provider, caches, CachedStateMetrics::zeroed());
812
813        // check that the storage returns the expected value
814        let res = state_provider.storage(address, storage_key);
815        assert!(res.is_ok());
816        assert_eq!(res.unwrap(), Some(storage_value));
817    }
818
819    #[test]
820    fn test_get_storage_populated() {
821        // make sure when we have something cached, we get the cached value in the `SlotStatus`
822        let address = Address::random();
823        let storage_key = StorageKey::random();
824        let storage_value = U256::from(1);
825
826        // insert into caches directly
827        let caches = ExecutionCacheBuilder::default().build_caches(1000);
828        caches.insert_storage(address, storage_key, Some(storage_value));
829
830        // check that the storage returns the cached value
831        let (slot_status, _) = caches.get_storage(&address, &storage_key);
832        assert_eq!(slot_status, SlotStatus::Value(storage_value));
833    }
834
835    #[test]
836    fn test_get_storage_not_cached() {
837        // make sure when we have nothing cached, we get the `NotCached` value in the `SlotStatus`
838        let storage_key = StorageKey::random();
839        let address = Address::random();
840
841        // just create empty caches
842        let caches = ExecutionCacheBuilder::default().build_caches(1000);
843
844        // check that the storage is not cached
845        let (slot_status, _) = caches.get_storage(&address, &storage_key);
846        assert_eq!(slot_status, SlotStatus::NotCached);
847    }
848
849    #[test]
850    fn test_get_storage_empty() {
851        // make sure when we insert an empty value to the cache, we get the `Empty` value in the
852        // `SlotStatus`
853        let address = Address::random();
854        let storage_key = StorageKey::random();
855
856        // insert into caches directly
857        let caches = ExecutionCacheBuilder::default().build_caches(1000);
858        caches.insert_storage(address, storage_key, None);
859
860        // check that the storage is empty
861        let (slot_status, _) = caches.get_storage(&address, &storage_key);
862        assert_eq!(slot_status, SlotStatus::Empty);
863    }
864
865    // Tests for SavedCache locking mechanism
866    #[test]
867    fn test_saved_cache_is_available() {
868        let execution_cache = ExecutionCacheBuilder::default().build_caches(1000);
869        let cache = SavedCache::new(B256::ZERO, execution_cache, CachedStateMetrics::zeroed());
870
871        // Initially, the cache should be available (only one reference)
872        assert!(cache.is_available(), "Cache should be available initially");
873
874        // Clone the usage guard (simulating it being handed out)
875        let _guard = cache.clone_guard_for_test();
876
877        // Now the cache should not be available (two references)
878        assert!(!cache.is_available(), "Cache should not be available with active guard");
879    }
880
881    #[test]
882    fn test_saved_cache_multiple_references() {
883        let execution_cache = ExecutionCacheBuilder::default().build_caches(1000);
884        let cache =
885            SavedCache::new(B256::from([2u8; 32]), execution_cache, CachedStateMetrics::zeroed());
886
887        // Create multiple references to the usage guard
888        let guard1 = cache.clone_guard_for_test();
889        let guard2 = cache.clone_guard_for_test();
890        let guard3 = guard1.clone();
891
892        // Cache should not be available with multiple guards
893        assert!(!cache.is_available());
894
895        // Drop guards one by one
896        drop(guard1);
897        assert!(!cache.is_available()); // Still not available
898
899        drop(guard2);
900        assert!(!cache.is_available()); // Still not available
901
902        drop(guard3);
903        assert!(cache.is_available()); // Now available
904    }
905}