1use 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#[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 #[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 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 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 pub fn evict_cached_changesets(&self, up_to_block: BlockNumber) {
160 self.changeset_cache.evict(up_to_block);
161 }
162
163 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 pub fn take_sparse_trie(&self) -> Option<PreservedSparseTrie> {
183 self.preserved_sparse_trie.lock().take()
184 }
185
186 pub fn store_sparse_trie(&self, trie: PreservedSparseTrie) {
188 *self.preserved_sparse_trie.lock() = Some(trie);
189 }
190
191 pub fn clear_sparse_trie(&self) {
193 *self.preserved_sparse_trie.lock() = None;
194 }
195
196 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 #[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 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 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 #[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 #[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 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 #[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 #[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 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 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 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 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(¤t_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#[derive(Clone, Copy, Debug)]
661pub(crate) struct OverlayCacheConfig {
662 pub(crate) precompute: bool,
664 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#[derive(Debug)]
676pub(crate) struct StateTrieOverlayError {
677 tip_hash: B256,
679 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 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 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}