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