1use 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
24pub(crate) struct CachedStateProvider<S> {
26    state_provider: S,
28
29    caches: ExecutionCache,
31
32    metrics: CachedStateMetrics,
34}
35
36impl<S> CachedStateProvider<S>
37where
38    S: StateProvider,
39{
40    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#[derive(Metrics, Clone)]
53#[metrics(scope = "sync.caching")]
54pub(crate) struct CachedStateMetrics {
55    code_cache_hits: Gauge,
57
58    code_cache_misses: Gauge,
60
61    code_cache_size: Gauge,
66
67    storage_cache_hits: Gauge,
69
70    storage_cache_misses: Gauge,
72
73    storage_cache_size: Gauge,
78
79    account_cache_hits: Gauge,
81
82    account_cache_misses: Gauge,
84
85    account_cache_size: Gauge,
90}
91
92impl CachedStateMetrics {
93    pub(crate) fn reset(&self) {
95        self.code_cache_hits.set(0);
97        self.code_cache_misses.set(0);
98
99        self.storage_cache_hits.set(0);
101        self.storage_cache_misses.set(0);
102
103        self.account_cache_hits.set(0);
105        self.account_cache_misses.set(0);
106    }
107
108    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#[derive(Debug, Clone, PartialEq, Eq)]
133pub(crate) enum SlotStatus {
134    NotCached,
136    Empty,
138    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 => {
150                self.metrics.storage_cache_misses.increment(1);
151                let final_res = self.state_provider.storage(account, storage_key)?;
152                self.caches.insert_storage(account, storage_key, final_res);
153                Ok(final_res)
154            }
155            SlotStatus::Empty => {
156                self.metrics.storage_cache_hits.increment(1);
157                Ok(None)
158            }
159            SlotStatus::Value(value) => {
160                self.metrics.storage_cache_hits.increment(1);
161                Ok(Some(value))
162            }
163        }
164    }
165}
166
167impl<S: BytecodeReader> BytecodeReader for CachedStateProvider<S> {
168    fn bytecode_by_hash(&self, code_hash: &B256) -> ProviderResult<Option<Bytecode>> {
169        if let Some(res) = self.caches.code_cache.get(code_hash) {
170            self.metrics.code_cache_hits.increment(1);
171            return Ok(res)
172        }
173
174        self.metrics.code_cache_misses.increment(1);
175
176        let final_res = self.state_provider.bytecode_by_hash(code_hash)?;
177        self.caches.code_cache.insert(*code_hash, final_res.clone());
178        Ok(final_res)
179    }
180}
181
182impl<S: StateRootProvider> StateRootProvider for CachedStateProvider<S> {
183    fn state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
184        self.state_provider.state_root(hashed_state)
185    }
186
187    fn state_root_from_nodes(&self, input: TrieInput) -> ProviderResult<B256> {
188        self.state_provider.state_root_from_nodes(input)
189    }
190
191    fn state_root_with_updates(
192        &self,
193        hashed_state: HashedPostState,
194    ) -> ProviderResult<(B256, TrieUpdates)> {
195        self.state_provider.state_root_with_updates(hashed_state)
196    }
197
198    fn state_root_from_nodes_with_updates(
199        &self,
200        input: TrieInput,
201    ) -> ProviderResult<(B256, TrieUpdates)> {
202        self.state_provider.state_root_from_nodes_with_updates(input)
203    }
204}
205
206impl<S: StateProofProvider> StateProofProvider for CachedStateProvider<S> {
207    fn proof(
208        &self,
209        input: TrieInput,
210        address: Address,
211        slots: &[B256],
212    ) -> ProviderResult<AccountProof> {
213        self.state_provider.proof(input, address, slots)
214    }
215
216    fn multiproof(
217        &self,
218        input: TrieInput,
219        targets: MultiProofTargets,
220    ) -> ProviderResult<MultiProof> {
221        self.state_provider.multiproof(input, targets)
222    }
223
224    fn witness(
225        &self,
226        input: TrieInput,
227        target: HashedPostState,
228    ) -> ProviderResult<Vec<alloy_primitives::Bytes>> {
229        self.state_provider.witness(input, target)
230    }
231}
232
233impl<S: StorageRootProvider> StorageRootProvider for CachedStateProvider<S> {
234    fn storage_root(
235        &self,
236        address: Address,
237        hashed_storage: HashedStorage,
238    ) -> ProviderResult<B256> {
239        self.state_provider.storage_root(address, hashed_storage)
240    }
241
242    fn storage_proof(
243        &self,
244        address: Address,
245        slot: B256,
246        hashed_storage: HashedStorage,
247    ) -> ProviderResult<StorageProof> {
248        self.state_provider.storage_proof(address, slot, hashed_storage)
249    }
250
251    fn storage_multiproof(
264        &self,
265        address: Address,
266        slots: &[B256],
267        hashed_storage: HashedStorage,
268    ) -> ProviderResult<StorageMultiProof> {
269        self.state_provider.storage_multiproof(address, slots, hashed_storage)
270    }
271}
272
273impl<S: BlockHashReader> BlockHashReader for CachedStateProvider<S> {
274    fn block_hash(&self, number: alloy_primitives::BlockNumber) -> ProviderResult<Option<B256>> {
275        self.state_provider.block_hash(number)
276    }
277
278    fn canonical_hashes_range(
279        &self,
280        start: alloy_primitives::BlockNumber,
281        end: alloy_primitives::BlockNumber,
282    ) -> ProviderResult<Vec<B256>> {
283        self.state_provider.canonical_hashes_range(start, end)
284    }
285}
286
287impl<S: HashedPostStateProvider> HashedPostStateProvider for CachedStateProvider<S> {
288    fn hashed_post_state(&self, bundle_state: &reth_revm::db::BundleState) -> HashedPostState {
289        self.state_provider.hashed_post_state(bundle_state)
290    }
291}
292
293#[derive(Debug, Clone)]
299pub(crate) struct ExecutionCache {
300    code_cache: Cache<B256, Option<Bytecode>>,
302
303    storage_cache: Cache<Address, Arc<AccountStorageCache>>,
306
307    account_cache: Cache<Address, Option<Account>>,
309}
310
311impl ExecutionCache {
312    pub(crate) fn get_storage(&self, address: &Address, key: &StorageKey) -> SlotStatus {
319        match self.storage_cache.get(address) {
320            None => SlotStatus::NotCached,
321            Some(account_cache) => account_cache.get_storage(key),
322        }
323    }
324
325    pub(crate) fn insert_storage(
327        &self,
328        address: Address,
329        key: StorageKey,
330        value: Option<StorageValue>,
331    ) {
332        self.insert_storage_bulk(address, [(key, value)]);
333    }
334
335    pub(crate) fn insert_storage_bulk<I>(&self, address: Address, storage_entries: I)
340    where
341        I: IntoIterator<Item = (StorageKey, Option<StorageValue>)>,
342    {
343        let account_cache = self.storage_cache.get(&address).unwrap_or_default();
344
345        for (key, value) in storage_entries {
346            account_cache.insert_storage(key, value);
347        }
348
349        self.storage_cache.insert(address, account_cache);
352    }
353
354    pub(crate) fn invalidate_account_storage(&self, address: &Address) {
356        self.storage_cache.invalidate(address);
357    }
358
359    pub(crate) fn total_storage_slots(&self) -> usize {
361        self.storage_cache.iter().map(|addr| addr.len()).sum()
362    }
363
364    #[instrument(level = "debug", target = "engine::caching", skip_all)]
383    pub(crate) fn insert_state(&self, state_updates: &BundleState) -> Result<(), ()> {
384        let _enter =
385            debug_span!(target: "engine::tree", "contracts", len = state_updates.contracts.len())
386                .entered();
387        for (code_hash, bytecode) in &state_updates.contracts {
389            self.code_cache.insert(*code_hash, Some(Bytecode(bytecode.clone())));
390        }
391        drop(_enter);
392
393        let _enter = debug_span!(
394            target: "engine::tree",
395            "accounts",
396            accounts = state_updates.state.len(),
397            storages =
398                state_updates.state.values().map(|account| account.storage.len()).sum::<usize>()
399        )
400        .entered();
401        for (addr, account) in &state_updates.state {
402            if account.status.is_not_modified() {
405                continue
406            }
407
408            if account.was_destroyed() {
410                self.account_cache.invalidate(addr);
412
413                self.invalidate_account_storage(addr);
414                continue
415            }
416
417            let Some(ref account_info) = account.info else {
421                trace!(target: "engine::caching", ?account, "Account with None account info found in state updates");
422                return Err(())
423            };
424
425            let storage_entries = account.storage.iter().map(|(storage_key, slot)| {
429                ((*storage_key).into(), Some(slot.present_value))
432            });
433            self.insert_storage_bulk(*addr, storage_entries);
434
435            self.account_cache.insert(*addr, Some(Account::from(account_info)));
438        }
439
440        Ok(())
441    }
442}
443
444#[derive(Debug)]
446pub(crate) struct ExecutionCacheBuilder {
447    code_cache_entries: u64,
449
450    storage_cache_entries: u64,
452
453    account_cache_entries: u64,
455}
456
457impl ExecutionCacheBuilder {
458    pub(crate) fn build_caches(self, total_cache_size: u64) -> ExecutionCache {
460        let storage_cache_size = (total_cache_size * 8888) / 10000; let account_cache_size = (total_cache_size * 556) / 10000; let code_cache_size = (total_cache_size * 556) / 10000; const EXPIRY_TIME: Duration = Duration::from_secs(7200); const TIME_TO_IDLE: Duration = Duration::from_secs(3600); let storage_cache = CacheBuilder::new(self.storage_cache_entries)
468            .weigher(|_key: &Address, value: &Arc<AccountStorageCache>| -> u32 {
469                let base_weight = 39_000;
471                let slots_weight = value.len() * 218;
472                (base_weight + slots_weight) as u32
473            })
474            .max_capacity(storage_cache_size)
475            .time_to_live(EXPIRY_TIME)
476            .time_to_idle(TIME_TO_IDLE)
477            .build_with_hasher(DefaultHashBuilder::default());
478
479        let account_cache = CacheBuilder::new(self.account_cache_entries)
480            .weigher(|_key: &Address, value: &Option<Account>| -> u32 {
481                20 + size_of_val(value) as u32
483            })
484            .max_capacity(account_cache_size)
485            .time_to_live(EXPIRY_TIME)
486            .time_to_idle(TIME_TO_IDLE)
487            .build_with_hasher(DefaultHashBuilder::default());
488
489        let code_cache = CacheBuilder::new(self.code_cache_entries)
490            .weigher(|_key: &B256, value: &Option<Bytecode>| -> u32 {
491                let code_size = match value {
492                    Some(bytecode) => {
493                        (size_of_val(value) +
495                            bytecode.bytecode().len() +
496                            bytecode
497                                .legacy_jump_table()
498                                .map(|table| table.as_slice().len())
499                                .unwrap_or_default()) as u32
500                    }
501                    None => size_of_val(value) as u32,
502                };
503                32 + code_size
504            })
505            .max_capacity(code_cache_size)
506            .time_to_live(EXPIRY_TIME)
507            .time_to_idle(TIME_TO_IDLE)
508            .build_with_hasher(DefaultHashBuilder::default());
509
510        ExecutionCache { code_cache, storage_cache, account_cache }
511    }
512}
513
514impl Default for ExecutionCacheBuilder {
515    fn default() -> Self {
516        Self {
524            code_cache_entries: 10_000_000,
525            storage_cache_entries: 10_000_000,
526            account_cache_entries: 10_000_000,
527        }
528    }
529}
530
531#[derive(Debug, Clone)]
534pub(crate) struct SavedCache {
535    hash: B256,
537
538    caches: ExecutionCache,
540
541    metrics: CachedStateMetrics,
543
544    usage_guard: Arc<()>,
547}
548
549impl SavedCache {
550    pub(super) fn new(hash: B256, caches: ExecutionCache, metrics: CachedStateMetrics) -> Self {
552        Self { hash, caches, metrics, usage_guard: Arc::new(()) }
553    }
554
555    pub(crate) const fn executed_block_hash(&self) -> B256 {
557        self.hash
558    }
559
560    pub(crate) fn split(self) -> (ExecutionCache, CachedStateMetrics) {
562        (self.caches, self.metrics)
563    }
564
565    pub(crate) fn is_available(&self) -> bool {
567        Arc::strong_count(&self.usage_guard) == 1
568    }
569
570    pub(crate) const fn cache(&self) -> &ExecutionCache {
572        &self.caches
573    }
574
575    pub(crate) const fn metrics(&self) -> &CachedStateMetrics {
577        &self.metrics
578    }
579
580    pub(crate) fn update_metrics(&self) {
582        self.metrics.storage_cache_size.set(self.caches.total_storage_slots() as f64);
583        self.metrics.account_cache_size.set(self.caches.account_cache.entry_count() as f64);
584        self.metrics.code_cache_size.set(self.caches.code_cache.entry_count() as f64);
585    }
586}
587
588#[cfg(test)]
589impl SavedCache {
590    fn clone_guard_for_test(&self) -> Arc<()> {
591        self.usage_guard.clone()
592    }
593}
594
595#[derive(Debug, Clone)]
600pub(crate) struct AccountStorageCache {
601    slots: Cache<StorageKey, Option<StorageValue>>,
603}
604
605impl AccountStorageCache {
606    pub(crate) fn new(max_slots: u64) -> Self {
608        Self {
609            slots: CacheBuilder::new(max_slots).build_with_hasher(DefaultHashBuilder::default()),
610        }
611    }
612
613    pub(crate) fn get_storage(&self, key: &StorageKey) -> SlotStatus {
618        match self.slots.get(key) {
619            None => SlotStatus::NotCached,
620            Some(None) => SlotStatus::Empty,
621            Some(Some(value)) => SlotStatus::Value(value),
622        }
623    }
624
625    pub(crate) fn insert_storage(&self, key: StorageKey, value: Option<StorageValue>) {
627        self.slots.insert(key, value);
628    }
629
630    pub(crate) fn len(&self) -> usize {
632        self.slots.entry_count() as usize
633    }
634}
635
636impl Default for AccountStorageCache {
637    fn default() -> Self {
638        Self::new(1_000_000)
642    }
643}
644
645#[cfg(test)]
646mod tests {
647    use super::*;
648    use alloy_primitives::{B256, U256};
649    use rand::Rng;
650    use reth_provider::test_utils::{ExtendedAccount, MockEthProvider};
651    use std::mem::size_of;
652
653    mod tracking_allocator {
654        use std::{
655            alloc::{GlobalAlloc, Layout, System},
656            sync::atomic::{AtomicUsize, Ordering},
657        };
658
659        #[derive(Debug)]
660        pub(crate) struct TrackingAllocator {
661            allocated: AtomicUsize,
662            total_allocated: AtomicUsize,
663            inner: System,
664        }
665
666        impl TrackingAllocator {
667            pub(crate) const fn new() -> Self {
668                Self {
669                    allocated: AtomicUsize::new(0),
670                    total_allocated: AtomicUsize::new(0),
671                    inner: System,
672                }
673            }
674
675            pub(crate) fn reset(&self) {
676                self.allocated.store(0, Ordering::SeqCst);
677                self.total_allocated.store(0, Ordering::SeqCst);
678            }
679
680            pub(crate) fn total_allocated(&self) -> usize {
681                self.total_allocated.load(Ordering::SeqCst)
682            }
683        }
684
685        unsafe impl GlobalAlloc for TrackingAllocator {
686            unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
687                let ret = unsafe { self.inner.alloc(layout) };
688                if !ret.is_null() {
689                    self.allocated.fetch_add(layout.size(), Ordering::SeqCst);
690                    self.total_allocated.fetch_add(layout.size(), Ordering::SeqCst);
691                }
692                ret
693            }
694
695            unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
696                self.allocated.fetch_sub(layout.size(), Ordering::SeqCst);
697                unsafe { self.inner.dealloc(ptr, layout) }
698            }
699        }
700    }
701
702    use tracking_allocator::TrackingAllocator;
703
704    #[global_allocator]
705    static ALLOCATOR: TrackingAllocator = TrackingAllocator::new();
706
707    fn measure_allocation<T, F>(f: F) -> (usize, T)
708    where
709        F: FnOnce() -> T,
710    {
711        ALLOCATOR.reset();
712        let result = f();
713        let total = ALLOCATOR.total_allocated();
714        (total, result)
715    }
716
717    #[test]
718    fn measure_storage_cache_overhead() {
719        let (base_overhead, cache) = measure_allocation(|| AccountStorageCache::new(1000));
720        println!("Base AccountStorageCache overhead: {base_overhead} bytes");
721        let mut rng = rand::rng();
722
723        let key = StorageKey::random();
724        let value = StorageValue::from(rng.random::<u128>());
725        let (first_slot, _) = measure_allocation(|| {
726            cache.insert_storage(key, Some(value));
727        });
728        println!("First slot insertion overhead: {first_slot} bytes");
729
730        const TOTAL_SLOTS: usize = 10_000;
731        let (test_slots, _) = measure_allocation(|| {
732            for _ in 0..TOTAL_SLOTS {
733                let key = StorageKey::random();
734                let value = StorageValue::from(rng.random::<u128>());
735                cache.insert_storage(key, Some(value));
736            }
737        });
738        println!("Average overhead over {} slots: {} bytes", TOTAL_SLOTS, test_slots / TOTAL_SLOTS);
739
740        println!("\nTheoretical sizes:");
741        println!("StorageKey size: {} bytes", size_of::<StorageKey>());
742        println!("StorageValue size: {} bytes", size_of::<StorageValue>());
743        println!("Option<StorageValue> size: {} bytes", size_of::<Option<StorageValue>>());
744        println!("Option<B256> size: {} bytes", size_of::<Option<B256>>());
745    }
746
747    #[test]
748    fn test_empty_storage_cached_state_provider() {
749        let address = Address::random();
751        let storage_key = StorageKey::random();
752        let account = ExtendedAccount::new(0, U256::ZERO);
753
754        let provider = MockEthProvider::default();
756        provider.extend_accounts(vec![(address, account)]);
757
758        let caches = ExecutionCacheBuilder::default().build_caches(1000);
759        let state_provider =
760            CachedStateProvider::new_with_caches(provider, caches, CachedStateMetrics::zeroed());
761
762        let res = state_provider.storage(address, storage_key);
764        assert!(res.is_ok());
765        assert_eq!(res.unwrap(), None);
766    }
767
768    #[test]
769    fn test_uncached_storage_cached_state_provider() {
770        let address = Address::random();
772        let storage_key = StorageKey::random();
773        let storage_value = U256::from(1);
774        let account =
775            ExtendedAccount::new(0, U256::ZERO).extend_storage(vec![(storage_key, storage_value)]);
776
777        let provider = MockEthProvider::default();
779        provider.extend_accounts(vec![(address, account)]);
780
781        let caches = ExecutionCacheBuilder::default().build_caches(1000);
782        let state_provider =
783            CachedStateProvider::new_with_caches(provider, caches, CachedStateMetrics::zeroed());
784
785        let res = state_provider.storage(address, storage_key);
787        assert!(res.is_ok());
788        assert_eq!(res.unwrap(), Some(storage_value));
789    }
790
791    #[test]
792    fn test_get_storage_populated() {
793        let address = Address::random();
795        let storage_key = StorageKey::random();
796        let storage_value = U256::from(1);
797
798        let caches = ExecutionCacheBuilder::default().build_caches(1000);
800        caches.insert_storage(address, storage_key, Some(storage_value));
801
802        let slot_status = caches.get_storage(&address, &storage_key);
804        assert_eq!(slot_status, SlotStatus::Value(storage_value));
805    }
806
807    #[test]
808    fn test_get_storage_not_cached() {
809        let storage_key = StorageKey::random();
811        let address = Address::random();
812
813        let caches = ExecutionCacheBuilder::default().build_caches(1000);
815
816        let slot_status = caches.get_storage(&address, &storage_key);
818        assert_eq!(slot_status, SlotStatus::NotCached);
819    }
820
821    #[test]
822    fn test_get_storage_empty() {
823        let address = Address::random();
826        let storage_key = StorageKey::random();
827
828        let caches = ExecutionCacheBuilder::default().build_caches(1000);
830        caches.insert_storage(address, storage_key, None);
831
832        let slot_status = caches.get_storage(&address, &storage_key);
834        assert_eq!(slot_status, SlotStatus::Empty);
835    }
836
837    #[test]
839    fn test_saved_cache_is_available() {
840        let execution_cache = ExecutionCacheBuilder::default().build_caches(1000);
841        let cache = SavedCache::new(B256::ZERO, execution_cache, CachedStateMetrics::zeroed());
842
843        assert!(cache.is_available(), "Cache should be available initially");
845
846        let _guard = cache.clone_guard_for_test();
848
849        assert!(!cache.is_available(), "Cache should not be available with active guard");
851    }
852
853    #[test]
854    fn test_saved_cache_multiple_references() {
855        let execution_cache = ExecutionCacheBuilder::default().build_caches(1000);
856        let cache =
857            SavedCache::new(B256::from([2u8; 32]), execution_cache, CachedStateMetrics::zeroed());
858
859        let guard1 = cache.clone_guard_for_test();
861        let guard2 = cache.clone_guard_for_test();
862        let guard3 = guard1.clone();
863
864        assert!(!cache.is_available());
866
867        drop(guard1);
869        assert!(!cache.is_available()); drop(guard2);
872        assert!(!cache.is_available()); drop(guard3);
875        assert!(cache.is_available()); }
877}