Skip to main content

reth_storage_overlay/
manager.rs

1//! State trie and execution overlays for in-memory blocks.
2//!
3//! Payload validation needs a view of the state trie as of an in-memory parent block even when that
4//! parent has not been persisted yet. [`OverlayManager`] tracks those in-memory blocks and builds
5//! reusable state trie and execution overlays on demand.
6
7use crate::{
8    changeset_cache::compute_block_trie_updates,
9    database_state_frontiers,
10    manager_metrics::{ExecutionOverlayMetrics, OverlayCacheMetrics, StateTrieOverlayMetrics},
11    ChangesetCache, ExecutionOverlay, OverlayBuilder,
12};
13use alloy_eips::BlockNumHash;
14use alloy_primitives::{BlockNumber, B256};
15use parking_lot::Mutex;
16use reth_chain_state::{BlockState, ExecutedBlock, PreservedSparseTrie};
17use reth_errors::ProviderResult;
18use reth_ethereum_primitives::EthPrimitives;
19use reth_primitives_traits::{
20    dashmap::{mapref::entry::Entry, DashMap},
21    AlloyBlockHeader, FastInstant, NodePrimitives,
22};
23use reth_storage_api::{
24    BlockNumReader, ChangeSetReader, DBProvider, PruneCheckpointReader, StageCheckpointReader,
25    StorageChangeSetReader, StorageSettingsCache,
26};
27#[cfg(feature = "rayon")]
28use reth_tasks::WorkerPool;
29use reth_trie::{updates::TrieUpdatesSorted, HashedPostStateSorted, TrieInputSorted};
30use std::{
31    fmt,
32    ops::RangeInclusive,
33    sync::{Arc, OnceLock},
34    time::Instant,
35};
36use tracing::{debug, trace};
37
38/// Manages state trie and execution overlays for in-memory blocks.
39///
40/// The manager owns the in-memory block graph, changeset cache, and caches keyed by
41/// `(anchor_hash, tip_hash)`.
42#[derive(Clone)]
43pub struct OverlayManager<N: NodePrimitives = EthPrimitives> {
44    blocks: Arc<DashMap<B256, ExecutedBlock<N>>>,
45    state_trie_overlays: OverlayCache<TrieInputSorted>,
46    execution_overlays: OverlayCache<ExecutionOverlay>,
47    changeset_cache: ChangesetCache,
48    preserved_sparse_trie: Arc<Mutex<Option<PreservedSparseTrie>>>,
49    #[cfg(feature = "rayon")]
50    worker_pool: Option<Arc<WorkerPool>>,
51    metrics: StateTrieOverlayMetrics,
52    execution_metrics: ExecutionOverlayMetrics,
53}
54
55impl<N: NodePrimitives> Default for OverlayManager<N> {
56    fn default() -> Self {
57        Self {
58            blocks: Default::default(),
59            state_trie_overlays: Default::default(),
60            execution_overlays: Default::default(),
61            changeset_cache: Default::default(),
62            preserved_sparse_trie: Default::default(),
63            #[cfg(feature = "rayon")]
64            worker_pool: None,
65            metrics: Default::default(),
66            execution_metrics: Default::default(),
67        }
68    }
69}
70
71impl<N: NodePrimitives> std::fmt::Debug for OverlayManager<N> {
72    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73        f.debug_struct("OverlayManager")
74            .field("blocks", &self.blocks.len())
75            .field("state_trie_overlays", &self.state_trie_overlays.len())
76            .field("execution_overlays", &self.execution_overlays.len())
77            .finish()
78    }
79}
80
81impl<N: NodePrimitives> OverlayManager<N> {
82    /// Create a new [`OverlayManager`] backed by the given worker pool.
83    #[cfg(feature = "rayon")]
84    pub fn new(worker_pool: Arc<WorkerPool>) -> Self {
85        Self {
86            blocks: Default::default(),
87            state_trie_overlays: Default::default(),
88            execution_overlays: Default::default(),
89            changeset_cache: Default::default(),
90            preserved_sparse_trie: Default::default(),
91            worker_pool: Some(worker_pool),
92            metrics: Default::default(),
93            execution_metrics: Default::default(),
94        }
95    }
96
97    /// Creates an overlay builder for `parent_hash`.
98    ///
99    /// This rebuilds the in-memory chain ending at `parent_hash` from the manager's block graph.
100    /// Prefer [`Self::overlay_builder_for_state`] whenever the caller already holds the chain.
101    pub fn overlay_builder(&self, parent_hash: B256) -> OverlayBuilder<N> {
102        OverlayBuilder::new(parent_hash, self.block_state(parent_hash).map(Arc::new), self.clone())
103    }
104
105    /// Creates an overlay builder for an already materialized in-memory chain.
106    ///
107    /// The chain tip is used as the parent hash, so no block graph lookup or chain rebuild is
108    /// performed. The builder only reads `state`, which makes it safe to share the same
109    /// [`BlockState`] with other holders.
110    pub fn overlay_builder_for_state(&self, state: Arc<BlockState<N>>) -> OverlayBuilder<N> {
111        OverlayBuilder::new(state.hash(), Some(state), self.clone())
112    }
113
114    pub(crate) fn block_state(&self, parent_hash: B256) -> Option<BlockState<N>> {
115        let mut blocks = self.parent_chain(parent_hash).collect::<Vec<_>>();
116        blocks.pop().map(|oldest| {
117            blocks.into_iter().rev().fold(BlockState::new(oldest), |parent, block| {
118                BlockState::with_parent(block, Some(Arc::new(parent)))
119            })
120        })
121    }
122
123    pub(crate) const fn changeset_cache(&self) -> &ChangesetCache {
124        &self.changeset_cache
125    }
126
127    /// Gets complete trie before-values for an inclusive block range, oldest values taking
128    /// precedence. No updates from outside the requested range are included.
129    pub fn get_or_compute_cached_changesets_range<P>(
130        &self,
131        provider: &P,
132        range: RangeInclusive<BlockNumber>,
133    ) -> ProviderResult<Arc<TrieUpdatesSorted>>
134    where
135        P: DBProvider
136            + ChangeSetReader
137            + StorageChangeSetReader
138            + StageCheckpointReader
139            + PruneCheckpointReader
140            + BlockNumReader
141            + StorageSettingsCache,
142    {
143        let (partial_state_trie, finish) = database_state_frontiers(provider)?;
144        self.get_or_compute_cached_changesets_range_at_frontiers(
145            provider,
146            range,
147            partial_state_trie,
148            finish,
149        )
150    }
151
152    pub(crate) fn get_or_compute_cached_changesets_range_at_frontiers<P>(
153        &self,
154        provider: &P,
155        range: RangeInclusive<BlockNumber>,
156        partial_state_trie: BlockNumHash,
157        finish: BlockNumHash,
158    ) -> ProviderResult<Arc<TrieUpdatesSorted>>
159    where
160        P: DBProvider
161            + ChangeSetReader
162            + StorageChangeSetReader
163            + StageCheckpointReader
164            + PruneCheckpointReader
165            + BlockNumReader
166            + StorageSettingsCache,
167    {
168        self.changeset_cache.get_or_compute_range(self, provider, range, partial_state_trie, finish)
169    }
170
171    /// Evicts cached changesets for blocks below `up_to_block`.
172    pub fn evict_cached_changesets(&self, up_to_block: BlockNumber) {
173        self.changeset_cache.evict(up_to_block);
174    }
175
176    /// Computes the trie updates produced by `block_number`.
177    pub fn compute_block_trie_updates<P>(
178        &self,
179        provider: &P,
180        block_number: BlockNumber,
181    ) -> ProviderResult<TrieUpdatesSorted>
182    where
183        P: DBProvider
184            + ChangeSetReader
185            + StorageChangeSetReader
186            + PruneCheckpointReader
187            + StageCheckpointReader
188            + BlockNumReader
189            + StorageSettingsCache,
190    {
191        compute_block_trie_updates(self, provider, block_number)
192    }
193
194    /// Takes the preserved sparse trie if present.
195    pub fn take_sparse_trie(&self) -> Option<PreservedSparseTrie> {
196        self.preserved_sparse_trie.lock().take()
197    }
198
199    /// Stores a preserved sparse trie for later reuse.
200    pub fn store_sparse_trie(&self, trie: PreservedSparseTrie) {
201        *self.preserved_sparse_trie.lock() = Some(trie);
202    }
203
204    /// Clears any preserved sparse trie state.
205    pub fn clear_sparse_trie(&self) {
206        *self.preserved_sparse_trie.lock() = None;
207    }
208
209    /// Waits until the sparse trie lock becomes available.
210    ///
211    /// This acquires and immediately releases the lock, ensuring that any ongoing operations
212    /// complete before returning. Returns the time spent waiting for the lock.
213    pub fn wait_for_sparse_trie_availability(&self) -> std::time::Duration {
214        let start = FastInstant::now();
215        let _guard = self.preserved_sparse_trie.lock();
216        let elapsed = start.elapsed();
217        if elapsed.as_millis() > 5 {
218            debug!(
219                target: "storage::overlay::manager",
220                blocked_for=?elapsed,
221                "Waited for preserved sparse trie to become available"
222            );
223        }
224        elapsed
225    }
226
227    /// Inserts an executed in-memory block into the state trie overlay manager.
228    #[tracing::instrument(
229        level = "trace",
230        target = "storage::overlay::manager",
231        skip_all,
232        fields(
233            block_hash = %block.recovered_block().hash(),
234            parent_hash = %block.recovered_block().parent_hash(),
235            duplicate = false,
236        )
237    )]
238    pub fn insert_block(&self, block: ExecutedBlock<N>) {
239        let hash = block.recovered_block().hash();
240        let parent_hash = block.recovered_block().parent_hash();
241        let span = tracing::Span::current();
242
243        // First add the block to the live graph; duplicate inserts do not need cache work.
244        match self.blocks.entry(hash) {
245            Entry::Occupied(_) => {
246                span.record("duplicate", true);
247                debug!(
248                    target: "storage::overlay::manager",
249                    %hash,
250                    %parent_hash,
251                    "state trie overlay block already inserted"
252                );
253                return
254            }
255            Entry::Vacant(entry) => {
256                entry.insert(block);
257            }
258        }
259
260        // Snapshot matching parent overlays before spawning so DashMap iteration guards are
261        // dropped.
262        let cached_parent_overlays = self
263            .execution_overlays
264            .entries
265            .iter()
266            .filter_map(|entry| {
267                let key = *entry.key();
268                (key.tip_hash == parent_hash).then_some(key.anchor_hash)
269            })
270            .collect::<Vec<_>>();
271
272        debug!(
273            target: "storage::overlay::manager",
274            %hash,
275            %parent_hash,
276            "inserted block into state trie overlay manager"
277        );
278        if cached_parent_overlays.is_empty() {
279            return
280        }
281
282        #[cfg(not(feature = "rayon"))]
283        let _ = cached_parent_overlays;
284
285        // When a new block is inserted we optimistically and asynchronously flatten an execution
286        // overlay for it
287        #[cfg(feature = "rayon")]
288        {
289            for anchor_hash in cached_parent_overlays {
290                self.precompute_execution_overlay(hash, anchor_hash);
291            }
292        }
293    }
294
295    /// Optimistically computes an execution overlay from `anchor_hash` to `tip_hash`.
296    ///
297    /// This returns without waiting for computation. No work is scheduled without a worker pool
298    /// or when the tip is already persisted at the anchor. A concurrent persistence or reorg may
299    /// make the requested range unavailable, in which case the background task skips it.
300    #[cfg(feature = "rayon")]
301    pub fn precompute_execution_overlay(&self, tip_hash: B256, anchor_hash: B256) {
302        if tip_hash == anchor_hash {
303            return
304        }
305        let Some(worker_pool) = &self.worker_pool else { return };
306        let manager = self.clone();
307        let parent_span = tracing::Span::current();
308        worker_pool.spawn(move || {
309            let _span = tracing::trace_span!(
310                target: "storage::overlay::manager",
311                parent: parent_span,
312                "precompute_execution_overlay",
313                %tip_hash,
314                %anchor_hash,
315            )
316            .entered();
317            if let Err(err) = manager.precompute_execution_overlay_for_parent(tip_hash, anchor_hash) {
318                debug!(target: "storage::overlay::manager", %err, "Skipping execution overlay precompute");
319            }
320        });
321    }
322
323    /// Removes blocks from the live block graph and prunes cached overlays that can no longer be
324    /// built from the remaining blocks.
325    #[tracing::instrument(
326        level = "trace",
327        target = "storage::overlay::manager",
328        skip_all,
329        fields(
330            block_count = tracing::field::Empty,
331            removed_blocks = tracing::field::Empty,
332            pruned_overlays = tracing::field::Empty,
333        )
334    )]
335    pub fn remove_blocks(&self, hashes: impl IntoIterator<Item = B256>) {
336        let span = tracing::Span::current();
337
338        // Remove blocks first, then prune overlays against the remaining block graph.
339        let mut block_count = 0usize;
340        let mut removed_blocks = 0usize;
341        let mut pruned_overlays = 0usize;
342        for hash in hashes {
343            block_count += 1;
344            removed_blocks += self.blocks.remove(&hash).is_some() as usize;
345        }
346        span.record("block_count", block_count);
347        span.record("removed_blocks", removed_blocks);
348
349        if removed_blocks > 0 {
350            let overlays_before = self.state_trie_overlays.len() + self.execution_overlays.len();
351            self.state_trie_overlays.retain(|key, _| {
352                self.contains_hash(key.tip_hash, key.anchor_hash, key.anchor_hash)
353            });
354            self.execution_overlays.retain(|key, _| {
355                self.contains_hash(key.tip_hash, key.anchor_hash, key.anchor_hash)
356            });
357            pruned_overlays = overlays_before
358                .saturating_sub(self.state_trie_overlays.len() + self.execution_overlays.len());
359            span.record("pruned_overlays", pruned_overlays);
360        }
361        debug!(
362            target: "storage::overlay::manager",
363            block_count,
364            removed_blocks,
365            pruned_overlays,
366            "removed blocks from state trie overlay manager"
367        );
368    }
369
370    /// Returns the flattened overlay from `anchor_hash` to `parent_hash`.
371    #[tracing::instrument(
372        level = "trace",
373        target = "storage::overlay::manager",
374        skip_all,
375        fields(tip_hash = %parent_state.hash(), anchor_hash = %anchor_hash)
376    )]
377    pub(crate) fn overlay_for_parent(
378        &self,
379        parent_state: &BlockState<N>,
380        anchor_hash: B256,
381        cache_config: OverlayCacheConfig,
382    ) -> Result<(Arc<TrieUpdatesSorted>, Arc<HashedPostStateSorted>), StateTrieOverlayError> {
383        let parent_hash = parent_state.hash();
384        if parent_hash == anchor_hash {
385            return Ok((
386                Arc::new(TrieUpdatesSorted::default()),
387                Arc::new(HashedPostStateSorted::default()),
388            ))
389        }
390        debug!(
391            target: "storage::overlay::manager",
392            tip_hash = %parent_hash,
393            %anchor_hash,
394            "loading state trie overlay for parent"
395        );
396        let input = self
397            .get_or_compute_overlay(
398                &self.state_trie_overlays,
399                &self.metrics,
400                anchor_hash,
401                parent_state,
402                cache_config,
403                |input, span| {
404                    self.compute_state_trie_overlay(
405                        input,
406                        anchor_hash,
407                        span,
408                        cache_config.write_to_cache,
409                    )
410                },
411            )?
412            .expect("required overlay lookup cannot skip an in-progress computation");
413        Ok((Arc::clone(&input.nodes), Arc::clone(&input.state)))
414    }
415
416    /// Returns execution data for the in-memory chain from `anchor_hash` to `parent_hash`.
417    #[tracing::instrument(
418        level = "trace",
419        target = "storage::overlay::manager",
420        skip_all,
421        fields(tip_hash = %parent_state.hash(), anchor_hash = %anchor_hash)
422    )]
423    pub(crate) fn execution_overlay_for_block_state(
424        &self,
425        parent_state: &BlockState<N>,
426        anchor_hash: B256,
427        cache_config: OverlayCacheConfig,
428    ) -> Result<Arc<ExecutionOverlay>, StateTrieOverlayError> {
429        Ok(self
430            .execution_overlay_for_parent_inner(parent_state, anchor_hash, cache_config)?
431            .expect("required overlay lookup cannot skip an in-progress computation"))
432    }
433
434    #[cfg(feature = "rayon")]
435    fn precompute_execution_overlay_for_parent(
436        &self,
437        parent_hash: B256,
438        anchor_hash: B256,
439    ) -> Result<(), StateTrieOverlayError> {
440        let parent_state = self
441            .block_state(parent_hash)
442            .ok_or(StateTrieOverlayError { tip_hash: parent_hash, anchor_hash })?;
443        self.execution_overlay_for_parent_inner(
444            &parent_state,
445            anchor_hash,
446            OverlayCacheConfig { precompute: true, write_to_cache: true },
447        )
448        .map(drop)
449    }
450
451    fn execution_overlay_for_parent_inner(
452        &self,
453        parent_state: &BlockState<N>,
454        anchor_hash: B256,
455        cache_config: OverlayCacheConfig,
456    ) -> Result<Option<Arc<ExecutionOverlay>>, StateTrieOverlayError> {
457        let parent_hash = parent_state.hash();
458        if parent_hash == anchor_hash {
459            return Ok(Some(Arc::new(ExecutionOverlay::default())))
460        }
461
462        self.get_or_compute_overlay(
463            &self.execution_overlays,
464            &self.execution_metrics,
465            anchor_hash,
466            parent_state,
467            cache_config,
468            |input, span| {
469                self.compute_execution_overlay(
470                    input,
471                    anchor_hash,
472                    span,
473                    cache_config.write_to_cache,
474                )
475            },
476        )
477    }
478
479    #[tracing::instrument(
480        level = "trace",
481        target = "storage::overlay::manager",
482        skip_all,
483        fields(
484            tip_hash = %parent_state.hash(),
485            anchor_hash = %anchor_hash,
486            cache_reused = tracing::field::Empty,
487            block_count = tracing::field::Empty,
488            parent_overlay_reused = tracing::field::Empty,
489        )
490    )]
491    fn get_or_compute_overlay<T, M>(
492        &self,
493        cache: &OverlayCache<T>,
494        metrics: &M,
495        anchor_hash: B256,
496        parent_state: &BlockState<N>,
497        cache_config: OverlayCacheConfig,
498        compute: impl FnOnce(ComputeOverlayInput<N, T>, tracing::Span) -> T,
499    ) -> Result<Option<Arc<T>>, StateTrieOverlayError>
500    where
501        M: OverlayCacheMetrics,
502    {
503        let tip_hash = parent_state.hash();
504        let key = OverlayCacheKey { anchor_hash, tip_hash };
505        let span = tracing::Span::current();
506        if let Some(entry) = cache.entries.get(&key).map(|entry| entry.value().clone()) {
507            metrics.record_cache_reuse();
508            span.record("cache_reused", true);
509            return match entry {
510                OverlayCacheEntry::Ready(input) => Ok(Some(input)),
511                OverlayCacheEntry::Computing(_) if cache_config.precompute => Ok(None),
512                OverlayCacheEntry::Computing(waiter) => Ok(Some(waiter.wait(metrics))),
513            }
514        }
515        span.record("cache_reused", false);
516
517        // Resolve the block path and any cached parent overlay before locking the child entry.
518        let mut blocks = Self::blocks_from_parent_state(parent_state, anchor_hash)?;
519        span.record("block_count", blocks.len());
520
521        if !cache_config.write_to_cache {
522            let parent_input = blocks.first().and_then(|block| {
523                let parent_hash = block.recovered_block().parent_hash();
524                (parent_hash != anchor_hash)
525                    .then(|| cache.ready(&OverlayCacheKey { anchor_hash, tip_hash: parent_hash }))
526                    .flatten()
527            });
528            span.record("parent_overlay_reused", parent_input.is_some());
529            let compute_input = match parent_input {
530                Some(parent_input) => {
531                    ComputeOverlayInput::ExtendCached { block: blocks.swap_remove(0), parent_input }
532                }
533                None => ComputeOverlayInput::MergeBlocks(blocks),
534            };
535            return Ok(Some(Arc::new(compute(compute_input, span))))
536        }
537
538        enum CacheAction<T> {
539            Ready(Arc<T>),
540            Wait(Arc<OverlayWaiter<T>>),
541            Compute(Arc<OverlayWaiter<T>>),
542        }
543
544        let parent_hash = parent_state.block_ref().recovered_block().parent_hash();
545        cache.retain(|sibling_key, entry| {
546            sibling_key.tip_hash == tip_hash ||
547                !matches!(entry, OverlayCacheEntry::Ready(_)) ||
548                self.blocks
549                    .get(&sibling_key.tip_hash)
550                    .is_none_or(|block| block.recovered_block().parent_hash() != parent_hash)
551        });
552
553        let action = match cache.entries.entry(key) {
554            Entry::Occupied(entry) => {
555                let entry = entry.get().clone();
556                metrics.record_cache_reuse();
557                span.record("cache_reused", true);
558                match entry {
559                    OverlayCacheEntry::Ready(input) => CacheAction::Ready(input),
560                    OverlayCacheEntry::Computing(_) if cache_config.precompute => return Ok(None),
561                    OverlayCacheEntry::Computing(waiter) => CacheAction::Wait(waiter),
562                }
563            }
564            Entry::Vacant(entry) => {
565                metrics.record_cache_fill();
566                let waiter = Arc::new(OverlayWaiter::new());
567                entry.insert(OverlayCacheEntry::Computing(Arc::clone(&waiter)));
568                CacheAction::Compute(waiter)
569            }
570        };
571
572        match action {
573            CacheAction::Ready(input) => Ok(Some(input)),
574            CacheAction::Wait(waiter) => Ok(Some(waiter.wait(metrics))),
575            CacheAction::Compute(waiter) => {
576                let parent_input = blocks.first().and_then(|block| {
577                    let parent_hash = block.recovered_block().parent_hash();
578                    (parent_hash != anchor_hash)
579                        .then(|| {
580                            cache
581                                .take_ready(&OverlayCacheKey { anchor_hash, tip_hash: parent_hash })
582                        })
583                        .flatten()
584                });
585                span.record("parent_overlay_reused", parent_input.is_some());
586                let compute_input = match parent_input {
587                    Some(parent_input) => ComputeOverlayInput::ExtendCached {
588                        block: blocks.swap_remove(0),
589                        parent_input,
590                    },
591                    None => ComputeOverlayInput::MergeBlocks(blocks),
592                };
593                let input = Arc::new(compute(compute_input, span));
594                waiter.finish(Arc::clone(&input));
595
596                if let Entry::Occupied(mut entry) = cache.entries.entry(key) {
597                    // The entry may have been pruned while the overlay was computing. Only cache
598                    // the result if the map still points at the waiter installed by this task.
599                    let should_publish = match entry.get() {
600                        OverlayCacheEntry::Computing(existing) => Arc::ptr_eq(existing, &waiter),
601                        OverlayCacheEntry::Ready(_) => false,
602                    };
603                    if should_publish {
604                        entry.insert(OverlayCacheEntry::Ready(Arc::clone(&input)));
605                    }
606                }
607
608                Ok(Some(input))
609            }
610        }
611    }
612
613    fn blocks_from_parent_state(
614        parent_state: &BlockState<N>,
615        anchor_hash: B256,
616    ) -> Result<Vec<ExecutedBlock<N>>, StateTrieOverlayError> {
617        let tip_hash = parent_state.hash();
618        let mut hash = tip_hash;
619        let mut blocks = Vec::new();
620        for state in parent_state.chain() {
621            let block = state.block();
622            if block.recovered_block().hash() != hash {
623                return Err(StateTrieOverlayError { tip_hash, anchor_hash })
624            }
625            hash = block.recovered_block().parent_hash();
626            blocks.push(block);
627            if hash == anchor_hash {
628                return Ok(blocks)
629            }
630        }
631        Err(StateTrieOverlayError { tip_hash, anchor_hash })
632    }
633
634    /// Returns every in-memory block in the chain whose tip is `parent_hash`.
635    pub(crate) fn parent_chain(
636        &self,
637        parent_hash: B256,
638    ) -> impl Iterator<Item = ExecutedBlock<N>> + '_ {
639        let mut hash = parent_hash;
640        std::iter::from_fn(move || {
641            let block = self.blocks.get(&hash)?;
642            hash = block.recovered_block().parent_hash();
643            Some(block.clone())
644        })
645    }
646
647    /// Returns true if `hash` is in the parent chain segment from `anchor_hash` inclusive to
648    /// `parent_hash` inclusive.
649    fn contains_hash(&self, parent_hash: B256, anchor_hash: B256, hash: B256) -> bool {
650        let mut current_hash = parent_hash;
651
652        loop {
653            if current_hash == hash {
654                return true
655            }
656            if current_hash == anchor_hash {
657                return false
658            }
659
660            let Some(block) = self.blocks.get(&current_hash) else { return false };
661            current_hash = block.recovered_block().parent_hash();
662        }
663    }
664
665    fn compute_state_trie_overlay(
666        &self,
667        compute_input: ComputeOverlayInput<N, TrieInputSorted>,
668        anchor_hash: B256,
669        _span: tracing::Span,
670        write_to_cache: bool,
671    ) -> TrieInputSorted {
672        if !write_to_cache {
673            return compute_overlay(compute_input, anchor_hash, &self.metrics)
674        }
675
676        #[cfg(feature = "rayon")]
677        {
678            if let Some(worker_pool) = &self.worker_pool {
679                let compute_span = _span;
680                let metrics = self.metrics.clone();
681                return worker_pool.spawn_and_wait(move || {
682                    let _guard = compute_span.enter();
683                    compute_overlay(compute_input, anchor_hash, &metrics)
684                })
685            }
686        }
687
688        compute_overlay(compute_input, anchor_hash, &self.metrics)
689    }
690
691    fn compute_execution_overlay(
692        &self,
693        compute_input: ComputeOverlayInput<N, ExecutionOverlay>,
694        anchor_hash: B256,
695        _span: tracing::Span,
696        write_to_cache: bool,
697    ) -> ExecutionOverlay {
698        if !write_to_cache {
699            return compute_execution_overlay_inner(
700                compute_input,
701                anchor_hash,
702                &self.execution_metrics,
703            )
704        }
705
706        #[cfg(feature = "rayon")]
707        {
708            if let Some(worker_pool) = &self.worker_pool {
709                let compute_span = _span;
710                let metrics = self.execution_metrics.clone();
711                return worker_pool.spawn_and_wait(move || {
712                    let _guard = compute_span.enter();
713                    compute_execution_overlay_inner(compute_input, anchor_hash, &metrics)
714                })
715            }
716        }
717
718        compute_execution_overlay_inner(compute_input, anchor_hash, &self.execution_metrics)
719    }
720}
721
722/// Controls how an overlay computation interacts with the manager cache.
723#[derive(Clone, Copy, Debug)]
724pub(crate) struct OverlayCacheConfig {
725    /// Whether this is a best-effort cache fill that must not wait for an existing computation.
726    pub(crate) precompute: bool,
727    /// Whether to retain the computed overlay in the manager cache.
728    pub(crate) write_to_cache: bool,
729}
730
731impl Default for OverlayCacheConfig {
732    fn default() -> Self {
733        Self { precompute: false, write_to_cache: true }
734    }
735}
736
737/// Error returned when a state trie overlay cannot be built from the manager's current block set.
738#[derive(Debug)]
739pub(crate) struct StateTrieOverlayError {
740    /// Requested in-memory tip hash.
741    pub(crate) tip_hash: B256,
742    /// Requested anchor hash.
743    pub(crate) anchor_hash: B256,
744}
745
746impl fmt::Display for StateTrieOverlayError {
747    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
748        write!(
749            f,
750            "state trie overlay for tip {} cannot be anchored to {} with current blocks",
751            self.tip_hash, self.anchor_hash
752        )
753    }
754}
755
756impl std::error::Error for StateTrieOverlayError {}
757
758#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
759struct OverlayCacheKey {
760    anchor_hash: B256,
761    tip_hash: B256,
762}
763
764struct OverlayCache<T> {
765    entries: Arc<DashMap<OverlayCacheKey, OverlayCacheEntry<T>>>,
766}
767
768impl<T> Default for OverlayCache<T> {
769    fn default() -> Self {
770        Self { entries: Default::default() }
771    }
772}
773
774impl<T> Clone for OverlayCache<T> {
775    fn clone(&self) -> Self {
776        Self { entries: Arc::clone(&self.entries) }
777    }
778}
779
780impl<T> OverlayCache<T> {
781    fn len(&self) -> usize {
782        self.entries.len()
783    }
784
785    fn retain(&self, mut keep: impl FnMut(&OverlayCacheKey, &OverlayCacheEntry<T>) -> bool) {
786        self.entries.retain(|key, entry| keep(key, entry))
787    }
788
789    /// Returns a ready entry without removing it from the cache.
790    fn ready(&self, key: &OverlayCacheKey) -> Option<Arc<T>> {
791        self.entries.get(key).and_then(|entry| match entry.value() {
792            OverlayCacheEntry::Ready(input) => Some(Arc::clone(input)),
793            OverlayCacheEntry::Computing(_) => None,
794        })
795    }
796
797    /// Removes and returns a ready entry.
798    ///
799    /// Transferring a parent entry lets `Arc::make_mut` extend it in place when no caller retains
800    /// it. Keeping the cache entry would otherwise guarantee a clone.
801    fn take_ready(&self, key: &OverlayCacheKey) -> Option<Arc<T>> {
802        let (_, entry) =
803            self.entries.remove_if(key, |_, entry| matches!(entry, OverlayCacheEntry::Ready(_)))?;
804        let OverlayCacheEntry::Ready(input) = entry else { unreachable!() };
805        Some(input)
806    }
807}
808
809enum OverlayCacheEntry<T> {
810    Ready(Arc<T>),
811    Computing(Arc<OverlayWaiter<T>>),
812}
813
814impl<T> Clone for OverlayCacheEntry<T> {
815    fn clone(&self) -> Self {
816        match self {
817            Self::Ready(input) => Self::Ready(Arc::clone(input)),
818            Self::Computing(waiter) => Self::Computing(Arc::clone(waiter)),
819        }
820    }
821}
822
823struct OverlayWaiter<T> {
824    input: OnceLock<Arc<T>>,
825}
826
827impl<T> OverlayWaiter<T> {
828    const fn new() -> Self {
829        Self { input: OnceLock::new() }
830    }
831
832    fn wait(&self, metrics: &impl OverlayCacheMetrics) -> Arc<T> {
833        let start = Instant::now();
834        let input = self.input.wait();
835        metrics.record_wait_duration(start.elapsed());
836        Arc::clone(input)
837    }
838
839    fn finish(&self, computed: Arc<T>) {
840        let _ = self.input.set(computed);
841    }
842}
843
844enum ComputeOverlayInput<N: NodePrimitives, T> {
845    ExtendCached { block: ExecutedBlock<N>, parent_input: Arc<T> },
846    MergeBlocks(Vec<ExecutedBlock<N>>),
847}
848
849#[tracing::instrument(
850    level = "trace",
851    target = "storage::overlay::manager",
852    skip_all,
853    fields(
854        anchor_hash = %anchor_hash,
855        block_count = tracing::field::Empty,
856        parent_overlay = tracing::field::Empty,
857        elapsed_us = tracing::field::Empty,
858    )
859)]
860fn compute_overlay<N: NodePrimitives>(
861    input: ComputeOverlayInput<N, TrieInputSorted>,
862    anchor_hash: B256,
863    metrics: &StateTrieOverlayMetrics,
864) -> TrieInputSorted {
865    let started_at = Instant::now();
866    let block_count = match &input {
867        ComputeOverlayInput::ExtendCached { .. } => 1,
868        ComputeOverlayInput::MergeBlocks(blocks) => blocks.len(),
869    };
870    let parent_overlay = matches!(&input, ComputeOverlayInput::ExtendCached { .. });
871    tracing::Span::current().record("block_count", block_count);
872    tracing::Span::current().record("parent_overlay", parent_overlay);
873
874    let overlay = match input {
875        ComputeOverlayInput::ExtendCached { block, parent_input } => {
876            let trie_data = block.trie_data();
877
878            trace!(
879                target: "storage::overlay::manager",
880                %anchor_hash,
881                head = %block.recovered_block().hash(),
882                "extending cached parent state trie overlay"
883            );
884
885            let mut parent_input = parent_input;
886            extend_overlay(
887                Arc::make_mut(&mut parent_input),
888                &trie_data.hashed_state,
889                &trie_data.trie_updates,
890            );
891            Arc::try_unwrap(parent_input).expect("Arc::make_mut leaves the child overlay unique")
892        }
893        ComputeOverlayInput::MergeBlocks(blocks) => merge_blocks(blocks),
894    };
895
896    let elapsed = started_at.elapsed();
897    metrics.overlay_computation_duration_seconds.record(elapsed.as_secs_f64());
898    tracing::Span::current().record("elapsed_us", elapsed.as_micros() as u64);
899    debug!(
900        target: "storage::overlay::manager",
901        %anchor_hash,
902        block_count,
903        parent_overlay,
904        ?elapsed,
905        "computed state trie overlay"
906    );
907
908    overlay
909}
910
911fn merge_blocks<N: NodePrimitives>(blocks: Vec<ExecutedBlock<N>>) -> TrieInputSorted {
912    let trie_data = blocks.iter().map(ExecutedBlock::trie_data).collect::<Vec<_>>();
913
914    #[cfg(feature = "rayon")]
915    let (nodes, state) = rayon::join(
916        || {
917            TrieUpdatesSorted::merge_batch(
918                trie_data.iter().map(|data| Arc::clone(&data.trie_updates)),
919            )
920        },
921        || {
922            HashedPostStateSorted::merge_batch(
923                trie_data.iter().map(|data| Arc::clone(&data.hashed_state)),
924            )
925        },
926    );
927
928    #[cfg(not(feature = "rayon"))]
929    let (nodes, state) = (
930        TrieUpdatesSorted::merge_batch(trie_data.iter().map(|data| Arc::clone(&data.trie_updates))),
931        HashedPostStateSorted::merge_batch(
932            trie_data.iter().map(|data| Arc::clone(&data.hashed_state)),
933        ),
934    );
935
936    TrieInputSorted::new(nodes, state, Default::default())
937}
938
939fn extend_overlay(
940    overlay: &mut TrieInputSorted,
941    hashed_state: &HashedPostStateSorted,
942    trie_updates: &TrieUpdatesSorted,
943) {
944    #[cfg(feature = "rayon")]
945    {
946        rayon::join(
947            || {
948                if !hashed_state.is_empty() {
949                    Arc::make_mut(&mut overlay.state).extend_ref_and_sort(hashed_state);
950                }
951            },
952            || {
953                if !trie_updates.is_empty() {
954                    Arc::make_mut(&mut overlay.nodes).extend_ref_and_sort(trie_updates);
955                }
956            },
957        );
958    }
959
960    #[cfg(not(feature = "rayon"))]
961    {
962        if !hashed_state.is_empty() {
963            Arc::make_mut(&mut overlay.state).extend_ref_and_sort(hashed_state);
964        }
965        if !trie_updates.is_empty() {
966            Arc::make_mut(&mut overlay.nodes).extend_ref_and_sort(trie_updates);
967        }
968    }
969}
970
971fn compute_execution_overlay_inner<N: NodePrimitives>(
972    input: ComputeOverlayInput<N, ExecutionOverlay>,
973    anchor_hash: B256,
974    metrics: &ExecutionOverlayMetrics,
975) -> ExecutionOverlay {
976    let started_at = Instant::now();
977    let block_count = match &input {
978        ComputeOverlayInput::ExtendCached { .. } => 1,
979        ComputeOverlayInput::MergeBlocks(blocks) => blocks.len(),
980    };
981    let parent_overlay = matches!(&input, ComputeOverlayInput::ExtendCached { .. });
982    tracing::Span::current().record("block_count", block_count);
983    tracing::Span::current().record("parent_overlay", parent_overlay);
984
985    let overlay = match input {
986        ComputeOverlayInput::ExtendCached { block, parent_input } => {
987            let mut parent_input = parent_input;
988            Arc::make_mut(&mut parent_input).extend_block(&block);
989            Arc::try_unwrap(parent_input).expect("Arc::make_mut leaves the child overlay unique")
990        }
991        ComputeOverlayInput::MergeBlocks(blocks) => {
992            let mut overlay = ExecutionOverlay::default();
993            for block in blocks.iter().rev() {
994                overlay.extend_block(block);
995            }
996            overlay
997        }
998    };
999
1000    let elapsed = started_at.elapsed();
1001    metrics.overlay_computation_duration_seconds.record(elapsed.as_secs_f64());
1002    tracing::Span::current().record("elapsed_us", elapsed.as_micros() as u64);
1003    debug!(
1004        target: "storage::overlay::manager",
1005        %anchor_hash,
1006        block_count,
1007        parent_overlay,
1008        ?elapsed,
1009        "computed execution overlay"
1010    );
1011
1012    overlay
1013}
1014
1015#[cfg(test)]
1016mod tests {
1017    use super::*;
1018    use alloy_primitives::{map::HashMap, Address, U256};
1019    use reth_chain_state::{test_utils::TestBlockBuilder, ExecutedBlock};
1020    use reth_ethereum_primitives::EthPrimitives;
1021    use reth_primitives_traits::Account;
1022    #[cfg(feature = "rayon")]
1023    use reth_tasks::WorkerPool;
1024    use reth_trie::{updates::TrieUpdatesSorted, HashedPostState, HashedStorage, SortedTrieData};
1025    use revm::{
1026        bytecode::Bytecode,
1027        database::BundleState,
1028        state::{AccountId, AccountInfo},
1029    };
1030    use std::{
1031        sync::{mpsc, Arc},
1032        thread,
1033        time::Duration,
1034    };
1035
1036    fn with_unique_state(
1037        block: &ExecutedBlock<EthPrimitives>,
1038        id: u8,
1039    ) -> ExecutedBlock<EthPrimitives> {
1040        let hashed_address = B256::with_last_byte(id);
1041        let hashed_slot = B256::with_last_byte(id.saturating_add(32));
1042        let hashed_state = HashedPostState::default()
1043            .with_accounts([(hashed_address, Some(Account::default()))])
1044            .with_storages([(
1045                hashed_address,
1046                HashedStorage::from_iter([(hashed_slot, U256::from(id))]),
1047            )])
1048            .into_sorted();
1049        let address = Address::with_last_byte(id);
1050        let slot = U256::from(id);
1051        let code_hash = B256::with_last_byte(id.saturating_add(64));
1052        let state = BundleState::builder(block.block_number()..=block.block_number())
1053            .state_present_account_info(
1054                address,
1055                AccountInfo {
1056                    nonce: id as u64,
1057                    balance: U256::from(id),
1058                    account_id: AccountId::new(id as usize),
1059                    ..Default::default()
1060                },
1061            )
1062            .state_storage(address, HashMap::from_iter([(slot, (U256::ZERO, U256::from(id)))]))
1063            .contract(code_hash, Bytecode::new_raw(vec![id].into()))
1064            .build();
1065        let mut execution_output = (*block.execution_output).clone();
1066        execution_output.state = state;
1067
1068        ExecutedBlock::new(
1069            Arc::clone(&block.recovered_block),
1070            Arc::new(execution_output),
1071            SortedTrieData::new(Arc::new(hashed_state), Arc::new(TrieUpdatesSorted::default())),
1072        )
1073    }
1074
1075    fn test_blocks() -> Vec<ExecutedBlock<EthPrimitives>> {
1076        TestBlockBuilder::eth()
1077            .get_executed_blocks(1..4)
1078            .enumerate()
1079            .map(|(index, block)| with_unique_state(&block, index as u8 + 1))
1080            .collect()
1081    }
1082
1083    impl OverlayManager {
1084        fn execution_overlay_for_parent(
1085            &self,
1086            parent_hash: B256,
1087            anchor_hash: B256,
1088        ) -> Result<Arc<ExecutionOverlay>, StateTrieOverlayError> {
1089            if parent_hash == anchor_hash {
1090                return Ok(Arc::new(ExecutionOverlay::default()))
1091            }
1092            let parent_state = self
1093                .block_state(parent_hash)
1094                .ok_or(StateTrieOverlayError { tip_hash: parent_hash, anchor_hash })?;
1095            self.execution_overlay_for_block_state(
1096                &parent_state,
1097                anchor_hash,
1098                OverlayCacheConfig::default(),
1099            )
1100        }
1101    }
1102
1103    fn overlay_for_parent(
1104        manager: &OverlayManager,
1105        parent_hash: B256,
1106        anchor_hash: B256,
1107    ) -> Result<(Arc<TrieUpdatesSorted>, Arc<HashedPostStateSorted>), StateTrieOverlayError> {
1108        let parent_state = manager
1109            .block_state(parent_hash)
1110            .ok_or(StateTrieOverlayError { tip_hash: parent_hash, anchor_hash })?;
1111        manager.overlay_for_parent(&parent_state, anchor_hash, OverlayCacheConfig::default())
1112    }
1113
1114    #[test]
1115    fn builds_managed_overlay_for_inserted_blocks() {
1116        let manager = OverlayManager::default();
1117        let blocks = test_blocks();
1118        for block in &blocks {
1119            manager.insert_block(block.clone());
1120        }
1121
1122        let anchor_hash = blocks[0].recovered_block().parent_hash();
1123
1124        let (_, state) =
1125            overlay_for_parent(&manager, blocks[2].recovered_block().hash(), anchor_hash).unwrap();
1126        assert_eq!(state.accounts.len(), 3);
1127
1128        let short_anchor = blocks[1].recovered_block().hash();
1129        let (_, short) =
1130            overlay_for_parent(&manager, blocks[2].recovered_block().hash(), short_anchor).unwrap();
1131        assert_eq!(short.accounts.len(), 1);
1132        let (_, cached_short) =
1133            overlay_for_parent(&manager, blocks[2].recovered_block().hash(), short_anchor).unwrap();
1134        assert!(Arc::ptr_eq(&short, &cached_short));
1135    }
1136
1137    #[test]
1138    fn builds_execution_overlay_for_inserted_blocks() {
1139        let manager = OverlayManager::default();
1140        let blocks = test_blocks();
1141        for block in &blocks {
1142            manager.insert_block(block.clone());
1143        }
1144
1145        let anchor_hash = blocks[0].recovered_block().parent_hash();
1146        let overlay = manager
1147            .execution_overlay_for_parent(blocks[2].recovered_block().hash(), anchor_hash)
1148            .unwrap();
1149
1150        for id in 1..=3 {
1151            let address = Address::with_last_byte(id);
1152            let code_hash = B256::with_last_byte(id + 64);
1153            assert_eq!(overlay.accounts()[&address].as_ref().unwrap().nonce, id as u64);
1154            assert_eq!(overlay.accounts()[&address].as_ref().unwrap().account_id, None);
1155            assert_eq!(overlay.storage()[&address][&U256::from(id)], U256::from(id));
1156            assert_eq!(overlay.code_hashes()[&code_hash], Bytecode::new_raw(vec![id].into()));
1157        }
1158        assert_eq!(
1159            overlay.block_hashes(),
1160            blocks[..=2].iter().map(|block| block.recovered_block().num_hash()).collect::<Vec<_>>(),
1161        );
1162
1163        let cached = manager
1164            .execution_overlay_for_parent(blocks[2].recovered_block().hash(), anchor_hash)
1165            .unwrap();
1166        assert!(Arc::ptr_eq(&overlay, &cached));
1167
1168        let short_anchor = blocks[1].recovered_block().hash();
1169        let short = manager
1170            .execution_overlay_for_parent(blocks[2].recovered_block().hash(), short_anchor)
1171            .unwrap();
1172        assert_eq!(short.accounts().len(), 1);
1173    }
1174
1175    #[test]
1176    fn computing_sibling_evicts_ready_cached_overlays() {
1177        let manager = OverlayManager::default();
1178        let mut builder = TestBlockBuilder::eth();
1179        let anchor_hash = B256::random();
1180        let parent = builder.get_executed_block_with_number(1, anchor_hash);
1181        let parent_hash = parent.recovered_block().hash();
1182        let first = builder.get_executed_block_with_number(2, parent_hash);
1183        let sibling = builder.get_executed_block_with_number(2, parent_hash);
1184        let sibling_hash = sibling.recovered_block().hash();
1185        let first_key = OverlayCacheKey { anchor_hash, tip_hash: first.recovered_block().hash() };
1186
1187        manager.insert_block(parent);
1188        manager.insert_block(first);
1189        manager.insert_block(sibling);
1190        manager
1191            .state_trie_overlays
1192            .entries
1193            .insert(first_key, OverlayCacheEntry::Ready(Arc::new(TrieInputSorted::default())));
1194        manager
1195            .execution_overlays
1196            .entries
1197            .insert(first_key, OverlayCacheEntry::Ready(Arc::new(ExecutionOverlay::default())));
1198
1199        overlay_for_parent(&manager, sibling_hash, anchor_hash).unwrap();
1200        manager.execution_overlay_for_parent(sibling_hash, anchor_hash).unwrap();
1201
1202        assert!(!manager.state_trie_overlays.entries.contains_key(&first_key));
1203        assert!(!manager.execution_overlays.entries.contains_key(&first_key));
1204    }
1205
1206    #[test]
1207    fn promotes_ready_parent_overlays_to_the_child() {
1208        let manager = OverlayManager::default();
1209        let blocks = test_blocks();
1210        for block in &blocks {
1211            manager.insert_block(block.clone());
1212        }
1213
1214        let anchor_hash = blocks[0].recovered_block().parent_hash();
1215        let parent_hash = blocks[1].recovered_block().hash();
1216        let child_hash = blocks[2].recovered_block().hash();
1217        let parent_key = OverlayCacheKey { anchor_hash, tip_hash: parent_hash };
1218        let child_key = OverlayCacheKey { anchor_hash, tip_hash: child_hash };
1219
1220        overlay_for_parent(&manager, parent_hash, anchor_hash).unwrap();
1221        manager.execution_overlay_for_parent(parent_hash, anchor_hash).unwrap();
1222
1223        overlay_for_parent(&manager, child_hash, anchor_hash).unwrap();
1224        manager.execution_overlay_for_parent(child_hash, anchor_hash).unwrap();
1225
1226        assert!(!manager.state_trie_overlays.entries.contains_key(&parent_key));
1227        assert!(manager.state_trie_overlays.entries.contains_key(&child_key));
1228        assert!(!manager.execution_overlays.entries.contains_key(&parent_key));
1229        assert!(manager.execution_overlays.entries.contains_key(&child_key));
1230    }
1231
1232    #[test]
1233    fn promotes_parent_overlays_held_by_callers() {
1234        let manager = OverlayManager::default();
1235        let blocks = test_blocks();
1236        for block in &blocks {
1237            manager.insert_block(block.clone());
1238        }
1239
1240        let anchor_hash = blocks[0].recovered_block().parent_hash();
1241        let parent_hash = blocks[1].recovered_block().hash();
1242        let child_hash = blocks[2].recovered_block().hash();
1243        let parent_key = OverlayCacheKey { anchor_hash, tip_hash: parent_hash };
1244
1245        overlay_for_parent(&manager, parent_hash, anchor_hash).unwrap();
1246        let state_parent = manager
1247            .state_trie_overlays
1248            .entries
1249            .get(&parent_key)
1250            .and_then(|entry| match entry.value() {
1251                OverlayCacheEntry::Ready(input) => Some(Arc::clone(input)),
1252                OverlayCacheEntry::Computing(_) => None,
1253            })
1254            .unwrap();
1255        let execution_parent =
1256            manager.execution_overlay_for_parent(parent_hash, anchor_hash).unwrap();
1257
1258        let (_, child_state) = overlay_for_parent(&manager, child_hash, anchor_hash).unwrap();
1259        let child_execution =
1260            manager.execution_overlay_for_parent(child_hash, anchor_hash).unwrap();
1261
1262        assert!(!manager.state_trie_overlays.entries.contains_key(&parent_key));
1263        assert!(!manager.execution_overlays.entries.contains_key(&parent_key));
1264        assert_eq!(state_parent.state.accounts.len(), 2);
1265        assert_eq!(execution_parent.accounts().len(), 2);
1266        assert_eq!(child_state.accounts.len(), 3);
1267        assert_eq!(child_execution.accounts().len(), 3);
1268        assert!(child_execution
1269            .accounts()
1270            .values()
1271            .flatten()
1272            .all(|account| account.account_id.is_none()));
1273    }
1274
1275    #[test]
1276    fn does_not_cache_or_take_parent_overlays_for_unmanaged_blocks() {
1277        let manager = OverlayManager::default();
1278        let blocks = test_blocks();
1279        for block in &blocks[..2] {
1280            manager.insert_block(block.clone());
1281        }
1282
1283        let anchor_hash = blocks[0].recovered_block().parent_hash();
1284        let parent_hash = blocks[1].recovered_block().hash();
1285        let child_hash = blocks[2].recovered_block().hash();
1286        let parent_key = OverlayCacheKey { anchor_hash, tip_hash: parent_hash };
1287        let child_key = OverlayCacheKey { anchor_hash, tip_hash: child_hash };
1288        let parent_state = manager.block_state(parent_hash).unwrap();
1289        let child_state = BlockState::with_parent(blocks[2].clone(), Some(Arc::new(parent_state)));
1290        let cache_config = OverlayCacheConfig { precompute: false, write_to_cache: false };
1291
1292        overlay_for_parent(&manager, parent_hash, anchor_hash).unwrap();
1293        manager.execution_overlay_for_parent(parent_hash, anchor_hash).unwrap();
1294
1295        let (_, state) =
1296            manager.overlay_for_parent(&child_state, anchor_hash, cache_config).unwrap();
1297        let execution = manager
1298            .execution_overlay_for_block_state(&child_state, anchor_hash, cache_config)
1299            .unwrap();
1300
1301        assert_eq!(state.accounts.len(), 3);
1302        assert_eq!(execution.accounts().len(), 3);
1303        assert!(manager.state_trie_overlays.entries.contains_key(&parent_key));
1304        assert!(!manager.state_trie_overlays.entries.contains_key(&child_key));
1305        assert!(manager.execution_overlays.entries.contains_key(&parent_key));
1306        assert!(!manager.execution_overlays.entries.contains_key(&child_key));
1307    }
1308
1309    #[cfg(feature = "rayon")]
1310    #[test]
1311    fn uncached_overlays_do_not_use_worker_pool() {
1312        let worker_pool = Arc::new(WorkerPool::new(1, "uncached-overlay-test"));
1313        let manager = OverlayManager::new(Arc::clone(&worker_pool));
1314        let block = test_blocks().remove(0);
1315        let anchor_hash = block.recovered_block().parent_hash();
1316        let parent_state = BlockState::new(block);
1317        let cache_config = OverlayCacheConfig { precompute: false, write_to_cache: false };
1318
1319        let (started_tx, started_rx) = mpsc::channel();
1320        let (release_tx, release_rx) = mpsc::channel();
1321        worker_pool.spawn(move || {
1322            started_tx.send(()).unwrap();
1323            release_rx.recv().unwrap();
1324        });
1325        started_rx.recv().unwrap();
1326
1327        let (completed_tx, completed_rx) = mpsc::channel();
1328        let task = thread::spawn(move || {
1329            let execution =
1330                manager.execution_overlay_for_block_state(&parent_state, anchor_hash, cache_config);
1331            let state = manager.overlay_for_parent(&parent_state, anchor_hash, cache_config);
1332            completed_tx.send((execution, state)).unwrap();
1333        });
1334
1335        let completed = completed_rx.recv_timeout(Duration::from_millis(100));
1336        release_tx.send(()).unwrap();
1337        task.join().unwrap();
1338
1339        assert!(completed.is_ok(), "uncached overlay used the worker pool");
1340        let (execution, state) = completed.unwrap();
1341        assert!(execution.is_ok());
1342        assert!(state.is_ok());
1343    }
1344
1345    #[cfg(feature = "rayon")]
1346    #[test]
1347    fn precomputes_execution_overlay_for_cached_parent() {
1348        let manager = OverlayManager::new(Arc::new(WorkerPool::new(1, "execution-overlay-test")));
1349        let blocks = test_blocks();
1350        let anchor_hash = blocks[0].recovered_block().parent_hash();
1351
1352        manager.insert_block(blocks[0].clone());
1353        manager
1354            .execution_overlay_for_parent(blocks[0].recovered_block().hash(), anchor_hash)
1355            .unwrap();
1356
1357        manager.insert_block(blocks[1].clone());
1358        let key = OverlayCacheKey { anchor_hash, tip_hash: blocks[1].recovered_block().hash() };
1359        let deadline = std::time::Instant::now() + Duration::from_secs(1);
1360        while !manager
1361            .execution_overlays
1362            .entries
1363            .get(&key)
1364            .is_some_and(|entry| matches!(entry.value(), OverlayCacheEntry::Ready(_)))
1365        {
1366            assert!(std::time::Instant::now() < deadline, "execution overlay was not precomputed");
1367            thread::sleep(Duration::from_millis(10));
1368        }
1369        assert!(!manager.execution_overlays.entries.contains_key(&OverlayCacheKey {
1370            anchor_hash,
1371            tip_hash: blocks[0].recovered_block().hash(),
1372        }));
1373    }
1374
1375    #[cfg(feature = "rayon")]
1376    #[test]
1377    fn precomputes_execution_overlay_after_anchor_advances() {
1378        let manager = OverlayManager::new(Arc::new(WorkerPool::new(1, "execution-overlay-test")));
1379        let blocks = test_blocks();
1380        for block in &blocks {
1381            manager.insert_block(block.clone());
1382        }
1383        let tip_hash = blocks[2].recovered_block().hash();
1384        let old_anchor = blocks[0].recovered_block().parent_hash();
1385        manager.execution_overlay_for_parent(tip_hash, old_anchor).unwrap();
1386
1387        // The state/trie frontier advances to block 1. Block 2 may already be persisted at the
1388        // Finish frontier, but its execution state must remain in the overlay until trie
1389        // persistence.
1390        let new_anchor = blocks[0].recovered_block().hash();
1391        manager.remove_blocks([new_anchor]);
1392        assert!(!manager
1393            .execution_overlays
1394            .entries
1395            .contains_key(&OverlayCacheKey { anchor_hash: old_anchor, tip_hash }));
1396
1397        manager.precompute_execution_overlay(tip_hash, new_anchor);
1398
1399        let key = OverlayCacheKey { anchor_hash: new_anchor, tip_hash };
1400        let deadline = std::time::Instant::now() + Duration::from_secs(1);
1401        let overlay = loop {
1402            if let Some(overlay) = manager.execution_overlays.ready(&key) {
1403                break overlay
1404            }
1405            assert!(std::time::Instant::now() < deadline, "new anchor overlay was not precomputed");
1406            thread::sleep(Duration::from_millis(10));
1407        };
1408        assert_eq!(
1409            overlay.block_hashes(),
1410            blocks[1..].iter().map(|block| block.recovered_block().num_hash()).collect::<Vec<_>>(),
1411        );
1412        assert!(!overlay.accounts().contains_key(&Address::with_last_byte(1)));
1413        for id in 2..=3 {
1414            let address = Address::with_last_byte(id);
1415            assert_eq!(overlay.accounts()[&address].as_ref().unwrap().nonce, id as u64);
1416            assert_eq!(overlay.storage()[&address][&U256::from(id)], U256::from(id));
1417        }
1418        let cached = manager.execution_overlay_for_parent(tip_hash, new_anchor).unwrap();
1419        assert!(Arc::ptr_eq(&overlay, &cached));
1420    }
1421
1422    #[cfg(feature = "rayon")]
1423    #[test]
1424    fn execution_overlay_precompute_does_not_wait_for_pending_entry() {
1425        let worker_pool = Arc::new(WorkerPool::new(1, "execution-overlay-pending-test"));
1426        let manager = OverlayManager::new(Arc::clone(&worker_pool));
1427        let block = test_blocks().remove(0);
1428        let anchor_hash = block.recovered_block().parent_hash();
1429        let tip_hash = block.recovered_block().hash();
1430        manager.insert_block(block);
1431
1432        let waiter = Arc::new(OverlayWaiter::new());
1433        manager.execution_overlays.entries.insert(
1434            OverlayCacheKey { anchor_hash, tip_hash },
1435            OverlayCacheEntry::Computing(Arc::clone(&waiter)),
1436        );
1437
1438        let (tx, rx) = mpsc::channel();
1439        worker_pool.spawn(move || {
1440            manager.precompute_execution_overlay_for_parent(tip_hash, anchor_hash).unwrap();
1441            tx.send(()).unwrap();
1442        });
1443
1444        let completed = rx.recv_timeout(Duration::from_millis(100));
1445        waiter.finish(Arc::new(ExecutionOverlay::default()));
1446        assert!(completed.is_ok(), "execution overlay precompute waited for pending entry");
1447    }
1448
1449    #[test]
1450    fn contains_hash_detects_hashes_from_anchor_to_parent() {
1451        let manager = OverlayManager::default();
1452        let blocks = test_blocks();
1453        for block in &blocks {
1454            manager.insert_block(block.clone());
1455        }
1456
1457        let anchor_hash = blocks[0].recovered_block().parent_hash();
1458        let parent_hash = blocks[2].recovered_block().hash();
1459
1460        assert!(manager.contains_hash(parent_hash, anchor_hash, anchor_hash));
1461        for block in &blocks {
1462            assert!(manager.contains_hash(
1463                parent_hash,
1464                anchor_hash,
1465                block.recovered_block().hash()
1466            ));
1467        }
1468        assert!(!manager.contains_hash(parent_hash, anchor_hash, B256::random()));
1469    }
1470
1471    #[test]
1472    fn contains_hash_rejects_hash_before_anchor() {
1473        let manager = OverlayManager::default();
1474        let blocks = test_blocks();
1475        for block in &blocks {
1476            manager.insert_block(block.clone());
1477        }
1478
1479        let parent_hash = blocks[2].recovered_block().hash();
1480        let anchor_hash = blocks[1].recovered_block().hash();
1481        let before_anchor_hash = blocks[0].recovered_block().hash();
1482
1483        assert!(manager.contains_hash(parent_hash, anchor_hash, parent_hash));
1484        assert!(manager.contains_hash(parent_hash, anchor_hash, anchor_hash));
1485        assert!(!manager.contains_hash(parent_hash, anchor_hash, before_anchor_hash));
1486    }
1487
1488    #[test]
1489    fn contains_hash_rejects_unknown_anchor() {
1490        let manager = OverlayManager::default();
1491        let blocks = test_blocks();
1492        for block in &blocks {
1493            manager.insert_block(block.clone());
1494        }
1495
1496        let parent_hash = blocks[2].recovered_block().hash();
1497        let anchor_hash = B256::random();
1498
1499        assert!(!manager.contains_hash(parent_hash, anchor_hash, anchor_hash));
1500    }
1501
1502    #[test]
1503    fn required_lookup_waits_for_in_progress_overlay() {
1504        let manager = OverlayManager::<EthPrimitives>::default();
1505        let block = test_blocks().remove(0);
1506        let parent_state = BlockState::new(block);
1507        let key = OverlayCacheKey {
1508            anchor_hash: parent_state.block_ref().recovered_block().parent_hash(),
1509            tip_hash: parent_state.hash(),
1510        };
1511        let waiter = Arc::new(OverlayWaiter::new());
1512        manager
1513            .state_trie_overlays
1514            .entries
1515            .insert(key, OverlayCacheEntry::Computing(Arc::clone(&waiter)));
1516
1517        let (tx, rx) = mpsc::channel();
1518        thread::spawn(move || {
1519            let res = manager
1520                .overlay_for_parent(&parent_state, key.anchor_hash, OverlayCacheConfig::default())
1521                .map(|(_, state)| state);
1522            tx.send(res).unwrap();
1523        });
1524
1525        assert!(matches!(
1526            rx.recv_timeout(Duration::from_millis(50)),
1527            Err(mpsc::RecvTimeoutError::Timeout)
1528        ));
1529
1530        waiter.finish(Arc::new(TrieInputSorted::default()));
1531
1532        let state = rx.recv_timeout(Duration::from_secs(1)).unwrap().unwrap();
1533        assert!(state.is_empty());
1534    }
1535
1536    #[test]
1537    fn prunes_cached_overlays_after_removing_blocks() {
1538        let manager = OverlayManager::default();
1539        let blocks = test_blocks();
1540        for block in &blocks {
1541            manager.insert_block(block.clone());
1542        }
1543
1544        let original_anchor = blocks[0].recovered_block().parent_hash();
1545        overlay_for_parent(&manager, blocks[2].recovered_block().hash(), original_anchor).unwrap();
1546        manager
1547            .execution_overlay_for_parent(blocks[2].recovered_block().hash(), original_anchor)
1548            .unwrap();
1549
1550        manager.remove_blocks([
1551            blocks[0].recovered_block().hash(),
1552            blocks[1].recovered_block().hash(),
1553        ]);
1554
1555        let anchor_hash = blocks[1].recovered_block().hash();
1556        assert!(overlay_for_parent(&manager, blocks[2].recovered_block().hash(), original_anchor)
1557            .is_err());
1558        assert!(manager
1559            .execution_overlay_for_parent(blocks[2].recovered_block().hash(), original_anchor)
1560            .is_err());
1561
1562        let (_, state) =
1563            overlay_for_parent(&manager, blocks[2].recovered_block().hash(), anchor_hash).unwrap();
1564        assert_eq!(state.accounts.len(), 1);
1565        let execution = manager
1566            .execution_overlay_for_parent(blocks[2].recovered_block().hash(), anchor_hash)
1567            .unwrap();
1568        assert_eq!(execution.accounts().len(), 1);
1569    }
1570}