Skip to main content

reth_trie_parallel/
state_root_task.rs

1//! State-root task interface types shared between the engine tree and the payload builder.
2//!
3//! The "state-root task" is the background multiproof and sparse-trie pipeline that computes
4//! state roots incrementally while a block executes. This module holds its boundary types:
5//! the input messages and
6//! stream views that feed it, and the handles
7//! that await its result. The per-block strategy abstraction that decides whether and how the
8//! task runs lives in `reth-engine-tree` under `tree::state_root_strategy`.
9
10use crate::error::StateRootTaskError;
11use alloy_primitives::{map::B256Map, B256};
12use reth_execution_types::{OnStateHook, StateUpdate};
13use reth_trie::{updates::TrieUpdates, HashedPostState, MultiProofTargetsV2, ProofV2Target};
14use std::{fmt, sync::Arc};
15
16/// Messages used internally by the multi proof task.
17#[derive(Debug)]
18pub enum StateRootMessage {
19    /// Prefetch proof targets
20    PrefetchProofs(MultiProofTargetsV2),
21    /// New state update from transaction execution.
22    StateUpdate(StateUpdate),
23    /// Pre-hashed state update from BAL conversion that can be applied directly without proofs.
24    HashedStateUpdate(HashedPostState),
25    /// Signals state update stream end.
26    ///
27    /// This is triggered by block execution, indicating that no additional state updates are
28    /// expected.
29    FinishedStateUpdates,
30}
31
32/// Outcome of the state root computation, including the state root itself with
33/// the trie updates.
34#[derive(Debug, Clone)]
35pub struct StateRootComputeOutcome {
36    /// The state root.
37    pub state_root: B256,
38    /// The trie updates.
39    pub trie_updates: Arc<TrieUpdates>,
40    /// Hashed post state produced while computing the state root.
41    pub hashed_state: Arc<HashedPostState>,
42}
43
44/// Handle to a background sparse trie state root computation.
45///
46/// Provides best-effort access hints, one authoritative update stream, and the final computation
47/// result. Dropping the handle cancels the task if it is still running.
48#[derive(Debug)]
49pub struct StateRootHandle {
50    /// The state root that the cached sparse trie is anchored at (parent block's state root).
51    cached_trie_state_root: B256,
52    /// Best-effort hint capability, taken once by prewarm wiring.
53    hint: Option<StateRootHintStream>,
54    /// The single authoritative update capability.
55    ///
56    /// Taken exactly once, either as an execution hook (serial execution) or as a hashed
57    /// update stream (parallel BAL streaming), so per block exactly one producer can finish
58    /// the update stream. Only producers hold update senders: once the taken capabilities are
59    /// dropped or finished, the update channel closes and the task knows producers are done.
60    authoritative: Option<StateRootUpdateStream>,
61    /// Guard whose drop cancels the state-root task if it is still running.
62    cancel_guard: StateRootTaskCancelGuard,
63    /// Receiver for the final state root result.
64    state_root_rx:
65        Option<std::sync::mpsc::Receiver<Result<StateRootComputeOutcome, StateRootTaskError>>>,
66    /// Receiver for the hashed post state.
67    hashed_state_rx: Option<std::sync::mpsc::Receiver<Arc<HashedPostState>>>,
68}
69
70impl StateRootHandle {
71    /// Creates a new [`StateRootHandle`].
72    pub fn new(
73        cached_trie_state_root: B256,
74        updates_tx: crossbeam_channel::Sender<StateRootMessage>,
75        cancel_guard: StateRootTaskCancelGuard,
76        state_root_rx: std::sync::mpsc::Receiver<
77            Result<StateRootComputeOutcome, StateRootTaskError>,
78        >,
79        hashed_state_rx: std::sync::mpsc::Receiver<Arc<HashedPostState>>,
80    ) -> Self {
81        let sink: Arc<dyn StateRootSink> = Arc::new(SparseTrieStateRootSink::new(updates_tx));
82        Self {
83            cached_trie_state_root,
84            hint: Some(StateRootHintStream::new(Arc::clone(&sink))),
85            authoritative: Some(StateRootUpdateStream::new(sink)),
86            cancel_guard,
87            state_root_rx: Some(state_root_rx),
88            hashed_state_rx: Some(hashed_state_rx),
89        }
90    }
91
92    /// Returns the state root that the cached sparse trie is anchored at.
93    pub const fn cached_trie_state_root(&self) -> B256 {
94        self.cached_trie_state_root
95    }
96
97    /// Takes the best-effort hint capability used by transaction prewarming.
98    ///
99    /// # Panics
100    ///
101    /// If called more than once.
102    pub const fn take_hint_stream(&mut self) -> StateRootHintStream {
103        self.hint.take().expect("hint stream already taken")
104    }
105
106    /// Takes the authoritative update capability as an EVM state hook.
107    ///
108    /// The hook finishes the update stream when dropped. It shares one slot with
109    /// [`Self::take_hashed_update_stream`], so only one of the two can exist per block.
110    ///
111    /// # Panics
112    ///
113    /// If the authoritative capability was already taken in either form.
114    pub fn take_execution_hook(&mut self) -> StateRootUpdateHook {
115        self.take_hashed_update_stream().into_state_hook()
116    }
117
118    /// Takes the authoritative update capability as a pre-hashed update stream.
119    ///
120    /// The stream is finished explicitly with [`StateRootUpdateStream::finish`]. It shares
121    /// one slot with [`Self::take_execution_hook`], so only one of the two can exist per
122    /// block.
123    ///
124    /// # Panics
125    ///
126    /// If the authoritative capability was already taken in either form.
127    pub const fn take_hashed_update_stream(&mut self) -> StateRootUpdateStream {
128        self.authoritative.take().expect("authoritative update capability already taken")
129    }
130
131    /// Awaits the state root computation result.
132    ///
133    /// # Panics
134    ///
135    /// If called more than once.
136    pub fn state_root(&mut self) -> Result<StateRootComputeOutcome, StateRootTaskError> {
137        self.state_root_rx
138            .take()
139            .expect("state_root already taken")
140            .recv()
141            .map_err(|_| StateRootTaskError::Other("sparse trie task dropped".into()))?
142    }
143
144    /// Takes the state root receiver for use with custom waiting logic (e.g., timeouts).
145    ///
146    /// # Panics
147    ///
148    /// If called more than once.
149    pub const fn take_state_root_rx(
150        &mut self,
151    ) -> std::sync::mpsc::Receiver<Result<StateRootComputeOutcome, StateRootTaskError>> {
152        self.state_root_rx.take().expect("state_root already taken")
153    }
154
155    /// Takes the hashed state receiver
156    ///
157    /// # Panics
158    ///
159    /// If called more than once.
160    pub const fn take_hashed_state_rx(
161        &mut self,
162    ) -> std::sync::mpsc::Receiver<Arc<HashedPostState>> {
163        self.hashed_state_rx.take().expect("hashed_state already taken")
164    }
165
166    /// Converts this sparse-trie handle into the opaque handle passed to payload builders.
167    ///
168    /// The payload builder only executes transactions, so the handle carries the execution
169    /// hook; the hint capability is dropped here.
170    pub fn into_payload_state_root_handle(mut self) -> PayloadStateRootHandle {
171        let hook = self.take_execution_hook();
172        PayloadStateRootHandle {
173            name: "sparse-trie",
174            hook: Some(hook),
175            cancel_guard: Some(self.cancel_guard),
176            state_root_rx: self.state_root_rx.take(),
177            hashed_state_rx: self.hashed_state_rx.take(),
178            on_payload_built: None,
179        }
180    }
181}
182
183/// Guard that cancels a state-root task when dropped.
184///
185/// The task watches the paired receiver in its event loop. No message is ever sent: the guard
186/// dropping disconnects the channel, which the task treats as the consumer abandoning the
187/// computation (for example on a timeout fallback or when a payload job is dropped unused).
188#[derive(Debug)]
189pub struct StateRootTaskCancelGuard(#[allow(dead_code)] crossbeam_channel::Sender<()>);
190
191impl StateRootTaskCancelGuard {
192    /// Creates a guard and the receiver a task watches for cancellation.
193    pub fn channel() -> (Self, crossbeam_channel::Receiver<()>) {
194        let (tx, rx) = crossbeam_channel::bounded(0);
195        (Self(tx), rx)
196    }
197}
198
199/// Opaque state-root task handle passed to payload builders.
200pub struct PayloadStateRootHandle {
201    name: &'static str,
202    /// Execution hook that streams per-transaction updates; taken once when building starts.
203    hook: Option<StateRootUpdateHook>,
204    /// Cancels the backing task when the handle is dropped without consuming the result.
205    cancel_guard: Option<StateRootTaskCancelGuard>,
206    state_root_rx:
207        Option<std::sync::mpsc::Receiver<Result<StateRootComputeOutcome, StateRootTaskError>>>,
208    hashed_state_rx: Option<std::sync::mpsc::Receiver<Arc<HashedPostState>>>,
209    /// Returns retained state to the engine once the completed block's identity is known.
210    on_payload_built: Option<Box<dyn FnOnce(B256, B256) + Send>>,
211}
212
213impl fmt::Debug for PayloadStateRootHandle {
214    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
215        f.debug_struct("PayloadStateRootHandle")
216            .field("name", &self.name)
217            .field("has_hook", &self.hook.is_some())
218            .field("has_cancel_guard", &self.cancel_guard.is_some())
219            .field("has_state_root_rx", &self.state_root_rx.is_some())
220            .field("has_hashed_state_rx", &self.hashed_state_rx.is_some())
221            .finish()
222    }
223}
224
225impl PayloadStateRootHandle {
226    /// Creates an opaque payload state-root handle.
227    ///
228    /// Tasks with a drop-to-cancel guard should attach it via the `StateRootHandle`
229    /// conversion; handles created here rely on their own task lifecycle.
230    pub const fn new(
231        name: &'static str,
232        hook: Option<StateRootUpdateHook>,
233        state_root_rx: std::sync::mpsc::Receiver<
234            Result<StateRootComputeOutcome, StateRootTaskError>,
235        >,
236        hashed_state_rx: Option<std::sync::mpsc::Receiver<Arc<HashedPostState>>>,
237    ) -> Self {
238        Self {
239            name,
240            hook,
241            cancel_guard: None,
242            state_root_rx: Some(state_root_rx),
243            hashed_state_rx,
244            on_payload_built: None,
245        }
246    }
247
248    /// Attaches a callback receiving the completed payload's block hash and state root.
249    ///
250    /// The callback must verify that any retained state matches the returned payload.
251    pub fn with_on_payload_built(
252        mut self,
253        callback: impl FnOnce(B256, B256) + Send + 'static,
254    ) -> Self {
255        self.on_payload_built = Some(Box::new(callback));
256        self
257    }
258
259    /// Takes the callback to invoke after a successful build, before returning its payload.
260    /// Aborted, cancelled, or failed builds must drop the callback without invoking it.
261    pub fn take_on_payload_built(&mut self) -> Option<Box<dyn FnOnce(B256, B256) + Send>> {
262        self.on_payload_built.take()
263    }
264
265    /// Returns the task name used in logs.
266    pub const fn name(&self) -> &'static str {
267        self.name
268    }
269
270    /// Takes the state hook that streams execution updates and finishes the stream on drop.
271    ///
272    /// # Panics
273    ///
274    /// If the handle was created without an execution hook, or the hook was already taken.
275    pub const fn take_state_hook(&mut self) -> StateRootUpdateHook {
276        self.hook.take().expect("payload state root task missing execution hook")
277    }
278
279    /// Awaits the state root computation result.
280    ///
281    /// # Panics
282    ///
283    /// If called more than once.
284    pub fn state_root(&mut self) -> Result<StateRootComputeOutcome, StateRootTaskError> {
285        self.state_root_rx
286            .take()
287            .expect("state_root already taken")
288            .recv()
289            .map_err(|_| StateRootTaskError::Other("state root task dropped".into()))?
290    }
291
292    /// Takes the state root receiver for use with custom waiting logic (e.g., timeouts).
293    ///
294    /// Dropping the handle continues to cancel the backing task if it owns a cancellation guard.
295    ///
296    /// # Panics
297    ///
298    /// If called more than once.
299    pub const fn take_state_root_rx(
300        &mut self,
301    ) -> std::sync::mpsc::Receiver<Result<StateRootComputeOutcome, StateRootTaskError>> {
302        self.state_root_rx.take().expect("state_root already taken")
303    }
304
305    /// Takes the hashed state receiver, if the handle was built with one and it was not taken
306    /// yet.
307    pub const fn try_take_hashed_state_rx(
308        &mut self,
309    ) -> Option<std::sync::mpsc::Receiver<Arc<HashedPostState>>> {
310        self.hashed_state_rx.take()
311    }
312}
313
314/// Hashed account and storage keys that a state-root task may want to prefetch.
315///
316/// Hints are not authoritative. They may be missing, duplicated, stale, or ignored by a task.
317/// Conversions to and from proof-target types allocate new collections.
318#[derive(Debug, Clone, Default)]
319pub struct StateAccessHint {
320    /// Hashed account keys that may be touched later in the block.
321    pub accounts: Vec<B256>,
322    /// Hashed storage keys keyed by hashed account.
323    pub storages: B256Map<Vec<B256>>,
324}
325
326impl From<MultiProofTargetsV2> for StateAccessHint {
327    fn from(targets: MultiProofTargetsV2) -> Self {
328        Self {
329            accounts: targets.account_targets.into_iter().map(|target| target.key()).collect(),
330            storages: targets
331                .storage_targets
332                .into_iter()
333                .map(|(account, slots)| {
334                    (account, slots.into_iter().map(|target| target.key()).collect())
335                })
336                .collect(),
337        }
338    }
339}
340
341impl From<StateAccessHint> for MultiProofTargetsV2 {
342    fn from(hint: StateAccessHint) -> Self {
343        Self {
344            account_targets: hint.accounts.into_iter().map(ProofV2Target::from).collect(),
345            storage_targets: hint
346                .storages
347                .into_iter()
348                .map(|(account, slots)| {
349                    (account, slots.into_iter().map(ProofV2Target::from).collect())
350                })
351                .collect(),
352        }
353    }
354}
355
356/// Semantic update stream consumed by state-root tasks.
357///
358/// These callbacks submit messages without reporting task completion or failure. Consumers must
359/// obtain the computation result separately, for example through [`StateRootHandle::state_root`].
360trait StateRootSink: Send + Sync + 'static {
361    /// Best-effort access hint from transaction prewarming.
362    fn on_access_hint(&self, _hint: StateAccessHint) {}
363
364    /// Authoritative state update from normal block execution.
365    fn on_state_update(&self, state: StateUpdate);
366
367    /// Authoritative pre-hashed state update, currently used by BAL streaming.
368    fn on_hashed_state_update(&self, state: HashedPostState);
369
370    /// Signals that no more authoritative state updates are expected.
371    fn on_updates_finished(&self);
372}
373
374/// Hint-only view of a state-root stream.
375#[derive(Clone)]
376pub struct StateRootHintStream {
377    inner: Arc<dyn StateRootSink>,
378}
379
380impl fmt::Debug for StateRootHintStream {
381    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
382        f.debug_struct("StateRootHintStream").finish_non_exhaustive()
383    }
384}
385
386impl StateRootHintStream {
387    /// Creates a new hint stream view.
388    fn new(inner: Arc<dyn StateRootSink>) -> Self {
389        Self { inner }
390    }
391
392    /// Emits a best-effort access hint.
393    pub fn on_access_hint(&self, hint: StateAccessHint) {
394        self.inner.on_access_hint(hint);
395    }
396}
397
398/// Authoritative update capability of a state-root stream.
399///
400/// A [`StateRootHandle`] provides this capability exactly once, so its producer can end the
401/// update stream either through the EVM state hook made with [`Self::into_state_hook`] (finishes
402/// on drop) or through a pre-hashed update stream (calls [`Self::finish`]). The
403/// type is deliberately not `Clone` and finishing consumes it, so a second end-of-stream
404/// signal cannot be produced.
405///
406/// Dropping the stream without calling [`Self::finish`] (for example when a producer dies)
407/// deliberately does not finish it: an unfinished stream means the updates are incomplete,
408/// and the task must not compute a root from them.
409pub struct StateRootUpdateStream {
410    inner: Arc<dyn StateRootSink>,
411}
412
413impl fmt::Debug for StateRootUpdateStream {
414    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
415        f.debug_struct("StateRootUpdateStream").finish_non_exhaustive()
416    }
417}
418
419impl StateRootUpdateStream {
420    /// Creates a new authoritative update stream backed by the given sink.
421    fn new(inner: Arc<dyn StateRootSink>) -> Self {
422        Self { inner }
423    }
424
425    /// Emits an authoritative pre-hashed state update.
426    pub fn on_hashed_state_update(&self, state: HashedPostState) {
427        self.inner.on_hashed_state_update(state);
428    }
429
430    /// Finishes the authoritative update stream.
431    pub fn finish(self) {
432        self.inner.on_updates_finished();
433    }
434
435    /// Converts this capability into an EVM state hook that finishes the stream on drop.
436    ///
437    /// See [`StateRootUpdateHook`] for why the hook finishes on drop while the bare stream
438    /// does not, and how a panic during execution is excluded from that.
439    pub fn into_state_hook(self) -> StateRootUpdateHook {
440        StateRootUpdateHook { inner: self.inner }
441    }
442}
443
444/// EVM hook that forwards state updates to the state-root task.
445///
446/// Dropping the hook signals the end of the update stream, so the hook is deliberately not
447/// `Clone`: a second copy would fire a spurious end-of-stream signal.
448///
449/// Unlike [`StateRootUpdateStream::finish`], the end of the stream is signaled by drop and
450/// not by an explicit call, because the EVM owns the hook until it is dropped and gives it no
451/// other end-of-execution signal. A drop during a panic unwind is excluded: execution died
452/// mid-block, so the stream stays unfinished and the task reports an error instead of
453/// computing a root from incomplete updates. Execution that fails by returning an error still
454/// drops the hook normally and finishes the stream; the caller abandons the result in that
455/// case, and the stored trie is rejected by the anchor check on the next block.
456pub struct StateRootUpdateHook {
457    inner: Arc<dyn StateRootSink>,
458}
459
460impl fmt::Debug for StateRootUpdateHook {
461    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
462        f.debug_struct("StateRootUpdateHook").finish_non_exhaustive()
463    }
464}
465
466impl OnStateHook for StateRootUpdateHook {
467    fn on_state(&mut self, state: StateUpdate) {
468        self.inner.on_state_update(state);
469    }
470}
471
472impl Drop for StateRootUpdateHook {
473    fn drop(&mut self) {
474        // A drop during a panic unwind means execution died mid-block. Leave the stream
475        // unfinished so the task fails instead of computing a root from partial updates.
476        if std::thread::panicking() {
477            return;
478        }
479        self.inner.on_updates_finished();
480    }
481}
482
483#[derive(Debug, Clone)]
484struct SparseTrieStateRootSink {
485    sender: crossbeam_channel::Sender<StateRootMessage>,
486}
487
488impl SparseTrieStateRootSink {
489    const fn new(sender: crossbeam_channel::Sender<StateRootMessage>) -> Self {
490        Self { sender }
491    }
492}
493
494// Send errors mean the receiving pipeline has stopped, so further updates and the finish
495// signal cannot be processed. Task errors are delivered through the separate result channel;
496// if the task drops that channel without a result, StateRootHandle::state_root reports an error.
497// Cancellation abandons the result. Ignoring send errors here does not report a successful root.
498impl StateRootSink for SparseTrieStateRootSink {
499    fn on_access_hint(&self, hint: StateAccessHint) {
500        let _ = self.sender.send(StateRootMessage::PrefetchProofs(hint.into()));
501    }
502
503    fn on_state_update(&self, state: StateUpdate) {
504        let _ = self.sender.send(StateRootMessage::StateUpdate(state));
505    }
506
507    fn on_hashed_state_update(&self, state: HashedPostState) {
508        let _ = self.sender.send(StateRootMessage::HashedStateUpdate(state));
509    }
510
511    fn on_updates_finished(&self) {
512        let _ = self.sender.send(StateRootMessage::FinishedStateUpdates);
513    }
514}
515
516#[cfg(test)]
517mod tests {
518    use super::*;
519    use std::{
520        sync::atomic::{AtomicUsize, Ordering},
521        time::Duration,
522    };
523
524    #[derive(Default)]
525    struct CountingSink {
526        state_updates: AtomicUsize,
527        finished_updates: AtomicUsize,
528    }
529
530    impl StateRootSink for CountingSink {
531        fn on_state_update(&self, state: StateUpdate) {
532            assert!(state.is_empty());
533            self.state_updates.fetch_add(1, Ordering::Relaxed);
534        }
535
536        fn on_hashed_state_update(&self, _state: HashedPostState) {}
537
538        fn on_updates_finished(&self) {
539            self.finished_updates.fetch_add(1, Ordering::Relaxed);
540        }
541    }
542
543    #[test]
544    fn state_access_hint_converts_to_sparse_targets() {
545        let account = B256::repeat_byte(0x01);
546        let storage_account = B256::repeat_byte(0x02);
547        let storage_slot = B256::repeat_byte(0x03);
548
549        let mut storages = B256Map::default();
550        storages.insert(storage_account, vec![storage_slot]);
551        let hint = StateAccessHint { accounts: vec![account], storages };
552
553        let targets = MultiProofTargetsV2::from(hint);
554        assert_eq!(targets.account_targets.len(), 1);
555        assert_eq!(targets.account_targets[0].key(), account);
556        assert_eq!(targets.storage_targets.len(), 1);
557        assert_eq!(targets.storage_targets[&storage_account].len(), 1);
558        assert_eq!(targets.storage_targets[&storage_account][0].key(), storage_slot);
559
560        let hint = StateAccessHint::from(targets);
561        assert_eq!(hint.accounts, vec![account]);
562        assert_eq!(hint.storages.len(), 1);
563        assert_eq!(hint.storages[&storage_account], vec![storage_slot]);
564    }
565
566    /// A hook dropped by a panic unwind must not finish the stream: the updates are
567    /// incomplete, and a finish marker would make the task compute a root from them.
568    #[test]
569    fn hook_dropped_during_panic_does_not_finish_stream() {
570        let sink = Arc::new(CountingSink::default());
571        let hook = StateRootUpdateStream::new(sink.clone()).into_state_hook();
572
573        let result = std::thread::spawn(move || {
574            let _hook = hook;
575            panic!("execution died mid-block");
576        })
577        .join();
578
579        assert!(result.is_err());
580        assert_eq!(sink.finished_updates.load(Ordering::Relaxed), 0);
581    }
582
583    /// The authoritative capability is a single slot: taking it as a hook and then again as
584    /// a hashed update stream (or in any other combination) must panic.
585    #[test]
586    #[should_panic(expected = "authoritative update capability already taken")]
587    fn authoritative_capability_can_only_be_taken_once() {
588        let (updates_tx, _updates_rx) = crossbeam_channel::unbounded();
589        let (cancel_guard, _cancel_rx) = StateRootTaskCancelGuard::channel();
590        let (_state_root_tx, state_root_rx) = std::sync::mpsc::channel();
591        let (_hashed_state_tx, hashed_state_rx) = std::sync::mpsc::channel();
592        let mut handle = StateRootHandle::new(
593            B256::ZERO,
594            updates_tx,
595            cancel_guard,
596            state_root_rx,
597            hashed_state_rx,
598        );
599
600        let _hook = handle.take_execution_hook();
601        let _ = handle.take_hashed_update_stream();
602    }
603
604    /// Lifecycle of the opaque handle a strategy hands to the payload builder: the execution
605    /// hook streams updates into the sink and signals completion on drop, the hashed-state
606    /// receiver can be taken exactly once, and the outcome arrives through the state-root
607    /// channel.
608    #[test]
609    fn payload_state_root_handle_lifecycle() {
610        let sink = Arc::new(CountingSink::default());
611        let hook = StateRootUpdateStream::new(sink.clone()).into_state_hook();
612
613        let (state_root_tx, state_root_rx) = std::sync::mpsc::channel();
614        let (hashed_state_tx, hashed_state_rx) = std::sync::mpsc::channel();
615        let mut handle =
616            PayloadStateRootHandle::new("test", Some(hook), state_root_rx, Some(hashed_state_rx));
617
618        assert_eq!(handle.name(), "test");
619
620        {
621            let mut hook = handle.take_state_hook();
622            hook.on_state(StateUpdate::default());
623        }
624        assert_eq!(sink.state_updates.load(Ordering::Relaxed), 1);
625        assert_eq!(sink.finished_updates.load(Ordering::Relaxed), 1);
626
627        hashed_state_tx.send(Arc::new(HashedPostState::default())).unwrap();
628        let rx = handle.try_take_hashed_state_rx().expect("first take returns the receiver");
629        assert!(rx.recv().is_ok());
630        assert!(handle.try_take_hashed_state_rx().is_none(), "second take returns None");
631
632        state_root_tx
633            .send(Ok(StateRootComputeOutcome {
634                state_root: B256::repeat_byte(0x42),
635                trie_updates: Arc::new(TrieUpdates::default()),
636                hashed_state: Arc::new(HashedPostState::default()),
637            }))
638            .unwrap();
639        let outcome = handle.state_root().expect("outcome is delivered");
640        assert_eq!(outcome.state_root, B256::repeat_byte(0x42));
641    }
642
643    #[test]
644    fn payload_state_root_receiver_retains_cancellation() {
645        let (updates_tx, _updates_rx) = crossbeam_channel::unbounded();
646        let (cancel_guard, cancel_rx) = StateRootTaskCancelGuard::channel();
647        let (_state_root_tx, state_root_rx) = std::sync::mpsc::channel();
648        let (_hashed_state_tx, hashed_state_rx) = std::sync::mpsc::channel();
649        let mut handle = StateRootHandle::new(
650            B256::ZERO,
651            updates_tx,
652            cancel_guard,
653            state_root_rx,
654            hashed_state_rx,
655        )
656        .into_payload_state_root_handle();
657
658        let state_root_rx = handle.take_state_root_rx();
659        assert!(matches!(
660            state_root_rx.recv_timeout(Duration::ZERO),
661            Err(std::sync::mpsc::RecvTimeoutError::Timeout)
662        ));
663        assert!(matches!(cancel_rx.try_recv(), Err(crossbeam_channel::TryRecvError::Empty)));
664
665        drop(handle);
666        assert!(matches!(
667            cancel_rx.recv_timeout(Duration::from_secs(1)),
668            Err(crossbeam_channel::RecvTimeoutError::Disconnected)
669        ));
670    }
671}