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