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