1use 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#[derive(Debug)]
18pub enum StateRootMessage {
19 PrefetchProofs(MultiProofTargetsV2),
21 StateUpdate(StateUpdate),
23 HashedStateUpdate(HashedPostState),
25 FinishedStateUpdates,
30}
31
32#[derive(Debug, Clone)]
35pub struct StateRootComputeOutcome {
36 pub state_root: B256,
38 pub trie_updates: Arc<TrieUpdates>,
40 pub hashed_state: Arc<HashedPostState>,
42}
43
44#[derive(Debug)]
49pub struct StateRootHandle {
50 cached_trie_state_root: B256,
52 hint: Option<StateRootHintStream>,
54 authoritative: Option<StateRootUpdateStream>,
61 cancel_guard: StateRootTaskCancelGuard,
63 state_root_rx:
65 Option<std::sync::mpsc::Receiver<Result<StateRootComputeOutcome, StateRootTaskError>>>,
66 hashed_state_rx: Option<std::sync::mpsc::Receiver<Arc<HashedPostState>>>,
68}
69
70impl StateRootHandle {
71 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 pub const fn cached_trie_state_root(&self) -> B256 {
94 self.cached_trie_state_root
95 }
96
97 pub const fn take_hint_stream(&mut self) -> StateRootHintStream {
103 self.hint.take().expect("hint stream already taken")
104 }
105
106 pub fn take_execution_hook(&mut self) -> StateRootUpdateHook {
115 self.take_hashed_update_stream().into_state_hook()
116 }
117
118 pub const fn take_hashed_update_stream(&mut self) -> StateRootUpdateStream {
128 self.authoritative.take().expect("authoritative update capability already taken")
129 }
130
131 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 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 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 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#[derive(Debug)]
189pub struct StateRootTaskCancelGuard(#[allow(dead_code)] crossbeam_channel::Sender<()>);
190
191impl StateRootTaskCancelGuard {
192 pub fn channel() -> (Self, crossbeam_channel::Receiver<()>) {
194 let (tx, rx) = crossbeam_channel::bounded(0);
195 (Self(tx), rx)
196 }
197}
198
199pub struct PayloadStateRootHandle {
201 name: &'static str,
202 hook: Option<StateRootUpdateHook>,
204 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 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 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 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 pub fn take_on_payload_built(&mut self) -> Option<Box<dyn FnOnce(B256, B256) + Send>> {
262 self.on_payload_built.take()
263 }
264
265 pub const fn name(&self) -> &'static str {
267 self.name
268 }
269
270 pub const fn take_state_hook(&mut self) -> StateRootUpdateHook {
276 self.hook.take().expect("payload state root task missing execution hook")
277 }
278
279 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 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 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#[derive(Debug, Clone, Default)]
319pub struct StateAccessHint {
320 pub accounts: Vec<B256>,
322 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
356trait StateRootSink: Send + Sync + 'static {
361 fn on_access_hint(&self, _hint: StateAccessHint) {}
363
364 fn on_state_update(&self, state: StateUpdate);
366
367 fn on_hashed_state_update(&self, state: HashedPostState);
369
370 fn on_updates_finished(&self);
372}
373
374#[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 fn new(inner: Arc<dyn StateRootSink>) -> Self {
389 Self { inner }
390 }
391
392 pub fn on_access_hint(&self, hint: StateAccessHint) {
394 self.inner.on_access_hint(hint);
395 }
396}
397
398pub 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 fn new(inner: Arc<dyn StateRootSink>) -> Self {
422 Self { inner }
423 }
424
425 pub fn on_hashed_state_update(&self, state: HashedPostState) {
427 self.inner.on_hashed_state_update(state);
428 }
429
430 pub fn finish(self) {
432 self.inner.on_updates_finished();
433 }
434
435 pub fn into_state_hook(self) -> StateRootUpdateHook {
440 StateRootUpdateHook { inner: self.inner }
441 }
442}
443
444pub 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 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
494impl 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 #[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 #[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 #[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}