reth_engine_tree/tree/txpool_prewarm/
control.rs

1use crate::tree::TxPoolPrewarmCacheSnapshot as Snapshot;
2use alloy_primitives::B256;
3use crossbeam_channel::{unbounded, Receiver, Sender};
4use parking_lot::RwLock;
5use std::{
6    fmt::Debug,
7    sync::{Arc, Weak},
8};
9
10/// Sending side of the txpool prewarming worker's control channel.
11pub(super) struct Control<J> {
12    commands: Sender<Command<J>>,
13    publication: Publication,
14}
15
16impl<J> Debug for Control<J> {
17    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
18        f.debug_struct("Control")
19            .field(
20                "published",
21                &self.publication.read().as_ref().map(|snapshot| snapshot.parent_hash()),
22            )
23            .finish_non_exhaustive()
24    }
25}
26
27impl<J> Control<J> {
28    pub(super) fn new() -> (Arc<Self>, Receiver<Command<J>>) {
29        let (commands, receiver) = unbounded();
30        let publication = Arc::new(RwLock::new(None));
31        (Arc::new(Self { commands, publication }), receiver)
32    }
33
34    pub(super) fn publication(&self) -> Publication {
35        Arc::clone(&self.publication)
36    }
37
38    pub(super) fn start(&self, parent_hash: B256, job: J) {
39        let _ = self.commands.send(Command::Start { parent_hash, job });
40    }
41
42    pub(super) fn pause(self: &Arc<Self>) -> PauseGuard<J> {
43        // Fire-and-forget: the worker never interrupts a transaction mid-execution, so waiting
44        // for it to observe the pause would only add that execution time to the caller's
45        // latency without reducing the overlap.
46        let _ = self.commands.send(Command::Pause);
47        PauseGuard { control: Arc::downgrade(self) }
48    }
49
50    pub(super) fn snapshot(&self, parent_hash: B256) -> Option<Snapshot> {
51        self.publication
52            .read()
53            .as_ref()
54            .filter(|snapshot| snapshot.parent_hash() == parent_hash)
55            .cloned()
56    }
57}
58
59/// Immutable snapshot publication shared by the handle and worker.
60pub(super) type Publication = Arc<RwLock<Option<Snapshot>>>;
61
62/// Commands sent to the txpool prewarming worker.
63pub(super) enum Command<J> {
64    /// Starts prewarming for a canonical head, replacing any previous job.
65    Start { parent_hash: B256, job: J },
66    /// Pauses prewarming until a matching [`Resume`](Self::Resume).
67    Pause,
68    /// Releases one active pause.
69    Resume,
70}
71
72/// One outstanding pause, released when dropped.
73pub(super) struct PauseGuard<J> {
74    control: Weak<Control<J>>,
75}
76
77impl<J> Drop for PauseGuard<J> {
78    fn drop(&mut self) {
79        // If the pause was never delivered the channel is disconnected, and this send fails
80        // just the same, so the worker can never see an unmatched resume.
81        if let Some(control) = self.control.upgrade() {
82            let _ = control.commands.send(Command::Resume);
83        }
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    type TestControl = Control<u64>;
92
93    fn control() -> (Arc<TestControl>, Receiver<Command<u64>>) {
94        TestControl::new()
95    }
96
97    fn snapshot(parent_hash: B256) -> Snapshot {
98        Snapshot::new(parent_hash, Default::default())
99    }
100
101    #[test]
102    fn snapshot_is_published_for_matching_parent() {
103        let (control, _receiver) = control();
104        let parent_hash = B256::repeat_byte(0x01);
105        *control.publication.write() = Some(snapshot(parent_hash));
106
107        assert_eq!(
108            control.snapshot(parent_hash).map(|snapshot| snapshot.parent_hash()),
109            Some(parent_hash)
110        );
111        assert!(control.snapshot(B256::ZERO).is_none());
112    }
113
114    #[test]
115    fn start_sends_job() {
116        let (control, receiver) = control();
117        let parent_hash = B256::repeat_byte(0x01);
118        control.start(parent_hash, 1);
119
120        assert!(matches!(
121            receiver.try_recv(),
122            Ok(Command::Start { parent_hash: received_parent, job: 1 })
123                if received_parent == parent_hash
124        ));
125    }
126
127    #[test]
128    fn pause_queues_without_waiting_and_resumes_on_drop() {
129        let (control, receiver) = control();
130        // Nothing reads the channel, so this returning at all proves pause does not block.
131        let guard = control.pause();
132        assert!(matches!(receiver.try_recv(), Ok(Command::Pause)));
133
134        drop(guard);
135        assert!(matches!(receiver.try_recv(), Ok(Command::Resume)));
136    }
137
138    #[test]
139    fn pause_guard_does_not_retain_control() {
140        let (control, _receiver) = control();
141        let weak_control = Arc::downgrade(&control);
142        let guard = control.pause();
143
144        drop(control);
145        assert!(weak_control.upgrade().is_none());
146        drop(guard);
147    }
148
149    #[test]
150    fn overlapping_pause_guards_send_matching_resumes() {
151        let (control, receiver) = control();
152        let first = control.pause();
153        let second = control.pause();
154        assert!(matches!(receiver.try_recv(), Ok(Command::Pause)));
155        assert!(matches!(receiver.try_recv(), Ok(Command::Pause)));
156
157        drop(first);
158        assert!(matches!(receiver.try_recv(), Ok(Command::Resume)));
159        drop(second);
160        assert!(matches!(receiver.try_recv(), Ok(Command::Resume)));
161    }
162
163    #[test]
164    fn dropping_control_disconnects_worker() {
165        let (control, receiver) = control();
166        drop(control);
167
168        assert!(receiver.recv().is_err());
169    }
170}
reth_engine_tree/tree/txpool_prewarm/control.rs

reth_engine_tree/tree/txpool_prewarm/
control.rs
