reth_engine_tree/tree/txpool_prewarm/
worker.rs

1use super::{
2    control::{Command, Publication},
3    Job, Source, Transactions,
4};
5use crate::tree::{StateProviderDatabase, TxPoolPrewarmCacheSnapshot as Snapshot};
6use alloy_evm::Evm;
7use alloy_primitives::B256;
8use crossbeam_channel::{Receiver, RecvTimeoutError, TryRecvError};
9use reth_evm::ConfigureEvm;
10use reth_primitives_traits::NodePrimitives;
11use reth_provider::{
12    BlockNumReader, DatabaseProviderFactory, PruneCheckpointReader, StageCheckpointReader,
13    StorageSettingsCache, TryIntoHistoricalStateProvider,
14};
15use reth_revm::{cached::CachedReads, db::State};
16use std::{
17    sync::Arc,
18    time::{Duration, Instant},
19};
20use tracing::{debug, trace};
21
22/// Maximum interval between snapshot publications and delay when no transaction is ready.
23const REFRESH_INTERVAL: Duration = Duration::from_millis(100);
24
25/// Delay while waiting for pool maintenance to advance to the state being warmed.
26const HEAD_POLL_INTERVAL: Duration = Duration::from_millis(10);
27
28/// The txpool prewarming worker.
29///
30/// A long-lived loop that speculatively executes the pool's best transactions on top of the
31/// current canonical state, recording every state read in a [`CachedReads`]. Roughly every
32/// [`REFRESH_INTERVAL`] it publishes an immutable snapshot of that cache, which block validation
33/// and payload building use to seed their own caches.
34///
35/// The worker is driven by [`Command`]s: `Start` points it at a new parent state, and
36/// `Pause`/`Resume` bracket cache-sensitive work elsewhere. Commands are only applied between
37/// batches, never while an EVM or state provider is alive.
38pub(super) struct Worker<N, P, Evm>
39where
40    N: NodePrimitives,
41    Evm: ConfigureEvm<Primitives = N>,
42{
43    /// Control commands from the [`Handle`](super::Handle).
44    commands: Receiver<Command<Job<N, P, Evm>>>,
45    /// Shared slot the latest snapshot is published into.
46    publication: Publication,
47    /// The txpool view transactions are drawn from.
48    source: Arc<dyn Source<N>>,
49    /// Configures the EVM used for speculative execution.
50    evm_config: Evm,
51    /// The parent state to warm, from the most recent `Start` command.
52    job: Option<(B256, Job<N, P, Evm>)>,
53    /// Outstanding pauses; the worker only warms while this is zero.
54    pauses: u64,
55    /// Read-through cache filled by execution; replaced whenever the warmed parent changes.
56    cache: CachedReads,
57    /// Parent whose state the `cache` reads were collected against.
58    cache_parent: Option<B256>,
59    /// Cache entry counts as of the last publication. The cache only ever grows, so a change
60    /// means it holds unpublished reads.
61    published_entries: (usize, usize, usize),
62    /// Live best-transactions iterator, tagged with the parent it was opened for.
63    transactions: Option<(B256, Transactions<N>)>,
64}
65
66impl<N, P, Evm> Worker<N, P, Evm>
67where
68    N: NodePrimitives,
69    P: DatabaseProviderFactory,
70    P::Provider: BlockNumReader
71        + PruneCheckpointReader
72        + StageCheckpointReader
73        + StorageSettingsCache
74        + TryIntoHistoricalStateProvider
75        + 'static,
76    Evm: ConfigureEvm<Primitives = N>,
77{
78    pub(super) fn new(
79        commands: Receiver<Command<Job<N, P, Evm>>>,
80        publication: Publication,
81        source: Arc<dyn Source<N>>,
82        evm_config: Evm,
83    ) -> Self {
84        Self {
85            commands,
86            publication,
87            source,
88            evm_config,
89            job: None,
90            pauses: 0,
91            cache: CachedReads::default(),
92            cache_parent: None,
93            published_entries: (0, 0, 0),
94            transactions: None,
95        }
96    }
97
98    /// Runs until the control side is dropped, which is the worker's shutdown signal.
99    pub(super) fn run(mut self) {
100        let _ = self.run_until_disconnected();
101    }
102
103    fn run_until_disconnected(&mut self) -> Result<(), ChannelDisconnected> {
104        loop {
105            let parent_hash = self.wait_until_runnable()?;
106
107            // The pool tracks canonical heads on its own schedule; check back shortly if it is
108            // not tracking this parent yet.
109            if !self.open_transactions(parent_hash) {
110                self.idle(HEAD_POLL_INTERVAL)?;
111                continue
112            }
113
114            if self.cache_parent != Some(parent_hash) {
115                self.cache = CachedReads::default();
116                self.cache_parent = Some(parent_hash);
117                self.published_entries = (0, 0, 0);
118                debug!(
119                    target: "engine::tree::txpool_prewarm",
120                    ?parent_hash,
121                    "started txpool prewarming"
122                );
123            }
124
125            let batch = self.warm_one_batch();
126
127            // A pending command may pause the worker or point it at a new parent: apply it (at
128            // the top of the loop) before spending time on publication.
129            if !self.commands.is_empty() {
130                continue
131            }
132            self.publish_snapshot_if_dirty();
133            if batch == BatchEnd::Rest {
134                self.idle(REFRESH_INTERVAL)?;
135            }
136        }
137    }
138
139    /// Blocks until the worker holds a job and no pauses are outstanding, applying every command
140    /// that arrives in the meantime. Returns the parent hash to warm.
141    fn wait_until_runnable(&mut self) -> Result<B256, ChannelDisconnected> {
142        loop {
143            self.apply_pending_commands()?;
144
145            if self.pauses == 0 &&
146                let Some((parent_hash, _)) = self.job.as_ref()
147            {
148                return Ok(*parent_hash)
149            }
150
151            let command = self.commands.recv().map_err(|_| ChannelDisconnected)?;
152            self.apply(command);
153        }
154    }
155
156    /// Ensures the transaction iterator matches `parent_hash`, opening a fresh one when the head
157    /// changed. Returns `false` while the pool is not yet tracking that parent.
158    fn open_transactions(&mut self, parent_hash: B256) -> bool {
159        if self.transactions.as_ref().is_none_or(|(parent, _)| *parent != parent_hash) {
160            // Release the stale iterator before asking the pool for a new one.
161            self.transactions = None;
162            self.transactions = self
163                .source
164                .best_transactions(parent_hash)
165                .map(|transactions| (parent_hash, transactions));
166        }
167        self.transactions.is_some()
168    }
169
170    /// Speculatively executes pool transactions against the parent state for at most
171    /// [`REFRESH_INTERVAL`], filling the cache with every state read.
172    ///
173    /// Stops early once the pool has no transaction ready or a command arrives. Commands are
174    /// never consumed here: a pending command merely ends the batch and is applied by the main
175    /// loop after the EVM and state provider built here are dropped.
176    fn warm_one_batch(&mut self) -> BatchEnd {
177        let (_, job) = self.job.as_ref().expect("wait_until_runnable installed a job");
178        let (parent_hash, transactions) =
179            self.transactions.as_mut().expect("open_transactions installed an iterator");
180
181        // Building a state provider opens a database transaction; don't bother under a pending
182        // command.
183        if !self.commands.is_empty() {
184            return BatchEnd::GoAgain
185        }
186
187        let state_provider = match job.provider_builder.build() {
188            Ok(provider) => provider,
189            Err(err) => {
190                trace!(
191                    target: "engine::tree::txpool_prewarm",
192                    %err,
193                    ?parent_hash,
194                    "failed to build txpool prewarming state provider"
195                );
196                return BatchEnd::Rest
197            }
198        };
199        let mut state = State::builder()
200            .with_database(self.cache.as_db_mut(StateProviderDatabase::new(state_provider)))
201            .build();
202        // The environment is the head block's own, not a predicted next-block one, and execution
203        // is out of context by design: transaction viability is the pool's business, so nonce,
204        // balance and (one-block-stale) basefee checks must not gate which state gets warmed.
205        let mut evm_env = job.evm_env.clone();
206        evm_env.cfg_env.disable_nonce_check = true;
207        evm_env.cfg_env.disable_balance_check = true;
208        evm_env.cfg_env.disable_base_fee = true;
209        let mut evm = self.evm_config.evm_with_env(&mut state, evm_env);
210
211        let deadline = Instant::now() + REFRESH_INTERVAL;
212        while self.commands.is_empty() && Instant::now() < deadline {
213            let Some(transaction) = transactions.next() else { return BatchEnd::Rest };
214            if let Err(err) = evm.transact(transaction.transaction) {
215                trace!(
216                    target: "engine::tree::txpool_prewarm",
217                    %err,
218                    tx_hash = ?transaction.hash,
219                    sender = %transaction.sender,
220                    "speculative txpool transaction execution failed"
221                );
222            }
223        }
224        BatchEnd::GoAgain
225    }
226
227    /// Publishes a fresh snapshot if the cache gained reads since the last publication.
228    fn publish_snapshot_if_dirty(&mut self) {
229        let entries = entry_counts(&self.cache);
230        if entries == self.published_entries {
231            return
232        }
233
234        let parent_hash = self.cache_parent.expect("reads only accumulate after a cache reset");
235        *self.publication.write() = Some(Snapshot::new(parent_hash, Arc::new(self.cache.clone())));
236        self.published_entries = entries;
237        let (accounts, storage, bytecodes) = entries;
238        debug!(
239            target: "engine::tree::txpool_prewarm",
240            ?parent_hash,
241            accounts,
242            storage,
243            bytecodes,
244            "published txpool prewarming snapshot"
245        );
246    }
247
248    /// Rests until a command arrives (applying it) or `timeout` elapses, whichever comes first.
249    fn idle(&mut self, timeout: Duration) -> Result<(), ChannelDisconnected> {
250        match self.commands.recv_timeout(timeout) {
251            Ok(command) => {
252                self.apply(command);
253                Ok(())
254            }
255            Err(RecvTimeoutError::Timeout) => Ok(()),
256            Err(RecvTimeoutError::Disconnected) => Err(ChannelDisconnected),
257        }
258    }
259
260    /// Applies every command already sitting in the channel, without blocking.
261    fn apply_pending_commands(&mut self) -> Result<(), ChannelDisconnected> {
262        loop {
263            match self.commands.try_recv() {
264                Ok(command) => self.apply(command),
265                Err(TryRecvError::Empty) => return Ok(()),
266                Err(TryRecvError::Disconnected) => return Err(ChannelDisconnected),
267            }
268        }
269    }
270
271    /// Applies a single command.
272    ///
273    /// Only called while no EVM or state provider is alive, so a paused worker holds no
274    /// execution resources.
275    fn apply(&mut self, command: Command<Job<N, P, Evm>>) {
276        match command {
277            Command::Start { parent_hash, job } => self.job = Some((parent_hash, job)),
278            Command::Pause => {
279                self.pauses =
280                    self.pauses.checked_add(1).expect("txpool prewarm pause count overflow");
281            }
282            Command::Resume => {
283                self.pauses = self
284                    .pauses
285                    .checked_sub(1)
286                    .expect("txpool prewarm resumed without a matching pause");
287            }
288        }
289    }
290}
291
292/// What to do after a warming batch.
293#[derive(Debug, Clone, Copy, PartialEq, Eq)]
294enum BatchEnd {
295    /// Batch again right away: the refresh deadline passed or a command interrupted the batch
296    /// while transactions may still be ready.
297    GoAgain,
298    /// Idle until something changes: the pool had no transaction ready, or no state provider
299    /// could be built.
300    Rest,
301}
302
303/// The control channel closed: every sender is dropped and the worker shuts down.
304#[derive(Debug, Clone, Copy, PartialEq, Eq)]
305struct ChannelDisconnected;
306
307/// Returns `(accounts, storage slots, bytecodes)` cached in `reads`.
308fn entry_counts(reads: &CachedReads) -> (usize, usize, usize) {
309    (
310        reads.accounts.len(),
311        reads.accounts.values().map(|account| account.storage.len()).sum(),
312        reads.contracts.len(),
313    )
314}
315
316#[cfg(test)]
317mod tests {
318    use super::{super::Transaction as PoolTransaction, *};
319    use crate::tree::StateProviderBuilder;
320    use alloy_consensus::{transaction::Recovered, Signed, TxLegacy};
321    use alloy_primitives::{Address, Signature, TxKind, U256};
322    use crossbeam_channel::{unbounded, Sender};
323    use parking_lot::{Mutex, RwLock};
324    use reth_ethereum_primitives::{EthPrimitives, TransactionSigned};
325    use reth_evm_ethereum::EthEvmConfig;
326    use reth_provider::test_utils::MockEthProvider;
327    use reth_stages_api::{StageCheckpoint, StageId};
328    use std::{
329        collections::{HashMap, VecDeque},
330        sync::atomic::{AtomicUsize, Ordering},
331        thread::{self, JoinHandle},
332    };
333
334    /// Upper bound on any single wait; failures surface as panics well before CI timeouts.
335    const WAIT_LIMIT: Duration = Duration::from_secs(5);
336    const POLL_INTERVAL: Duration = Duration::from_millis(2);
337
338    type TestJob = Job<EthPrimitives, MockEthProvider, EthEvmConfig>;
339
340    /// Drives a live worker thread through its public seams only: commands in, the publication
341    /// slot and the scripted pool out.
342    struct Harness {
343        commands: Sender<Command<TestJob>>,
344        publication: Publication,
345        pool: Arc<ScriptedPool>,
346        worker: Option<JoinHandle<()>>,
347    }
348
349    impl Harness {
350        fn spawn() -> Self {
351            let (commands, receiver) = unbounded();
352            let publication: Publication = Arc::new(RwLock::new(None));
353            let pool = Arc::new(ScriptedPool::default());
354            let worker = thread::spawn({
355                let publication = Arc::clone(&publication);
356                let source: Arc<dyn Source<EthPrimitives>> = pool.clone();
357                move || Worker::new(receiver, publication, source, EthEvmConfig::mainnet()).run()
358            });
359            Self { commands, publication, pool, worker: Some(worker) }
360        }
361
362        /// Drops the control channel and waits for the worker thread to exit.
363        fn shutdown(mut self) {
364            let (disconnected, _) = unbounded();
365            self.commands = disconnected;
366            let worker = self.worker.take().expect("worker already joined");
367            wait_until("the worker thread exits", || worker.is_finished());
368            worker.join().unwrap();
369        }
370
371        /// Points the worker at `parent_hash`, as [`Handle::start`](super::super::Handle) does.
372        fn start(&self, parent_hash: B256) {
373            let provider = MockEthProvider::default();
374            provider.enable_database_provider();
375            provider.add_header(parent_hash, Default::default());
376            provider.add_stage_checkpoint(StageId::Finish, StageCheckpoint::new(0));
377            let job = Job {
378                evm_env: Default::default(),
379                provider_builder: StateProviderBuilder::new(
380                    provider,
381                    parent_hash,
382                    reth_storage_overlay::OverlayManager::default(),
383                ),
384            };
385            self.commands.send(Command::Start { parent_hash, job }).unwrap();
386        }
387
388        /// Pauses the worker, as [`Handle::pause`](super::super::Handle) does. Fire-and-forget
389        /// is still deterministic: the worker never publishes work it performs after the pause
390        /// is queued until [`Self::resume`], because publication is skipped while a command is
391        /// pending and a consumed pause blocks the loop.
392        fn pause(&self) {
393            self.commands.send(Command::Pause).unwrap();
394        }
395
396        fn resume(&self) {
397            self.commands.send(Command::Resume).unwrap();
398        }
399
400        /// Blocks until a published snapshot satisfies `accept`.
401        fn published(&self, accept: impl Fn(&Snapshot) -> bool) -> Snapshot {
402            let deadline = Instant::now() + WAIT_LIMIT;
403            loop {
404                let snapshot = self.publication.read().as_ref().cloned();
405                if let Some(snapshot) = snapshot &&
406                    accept(&snapshot)
407                {
408                    return snapshot
409                }
410                assert!(Instant::now() < deadline, "timed out waiting for a matching snapshot");
411                thread::sleep(POLL_INTERVAL);
412            }
413        }
414
415        fn published_for(&self, parent_hash: B256) -> Snapshot {
416            self.published(|snapshot| snapshot.parent_hash() == parent_hash)
417        }
418
419        fn published_entry_counts(&self) -> Option<(usize, usize, usize)> {
420            self.publication.read().as_ref().map(|snapshot| snapshot.entry_counts())
421        }
422    }
423
424    impl Drop for Harness {
425        fn drop(&mut self) {
426            // Disconnect and reap the worker thread so it cannot outlive the test.
427            let (disconnected, _) = unbounded();
428            self.commands = disconnected;
429            if let Some(worker) = self.worker.take() {
430                let _ = worker.join();
431            }
432        }
433    }
434
435    /// A pool the tests script per parent: [`Self::push`] hands a transaction to the iterator
436    /// opened for that parent, and unknown parents read as "not tracking this head yet".
437    #[derive(Debug, Default)]
438    struct ScriptedPool {
439        queues: Arc<Mutex<HashMap<B256, VecDeque<PoolTransaction<EthPrimitives>>>>>,
440        /// Iterators handed out; the worker is expected to open exactly one per parent.
441        opened: AtomicUsize,
442        /// Polls answered with "not tracking this head yet".
443        not_ready: AtomicUsize,
444    }
445
446    impl ScriptedPool {
447        fn push(&self, parent_hash: B256, transaction: PoolTransaction<EthPrimitives>) {
448            self.queues.lock().entry(parent_hash).or_default().push_back(transaction);
449        }
450    }
451
452    impl Source<EthPrimitives> for ScriptedPool {
453        fn best_transactions(&self, parent_hash: B256) -> Option<Transactions<EthPrimitives>> {
454            if !self.queues.lock().contains_key(&parent_hash) {
455                self.not_ready.fetch_add(1, Ordering::Relaxed);
456                return None
457            }
458            self.opened.fetch_add(1, Ordering::Relaxed);
459            let queues = Arc::clone(&self.queues);
460            Some(Box::new(std::iter::from_fn(move || {
461                queues.lock().get_mut(&parent_hash)?.pop_front()
462            })))
463        }
464    }
465
466    /// A signed transfer to `recipient`. The signature is a dummy: the worker executes with the
467    /// attached sender and disabled nonce/balance checks, so it is never recovered or validated.
468    fn transfer(recipient: u8) -> PoolTransaction<EthPrimitives> {
469        let transaction = TxLegacy {
470            gas_limit: 21_000,
471            to: TxKind::Call(Address::repeat_byte(recipient)),
472            value: U256::from(1),
473            ..Default::default()
474        };
475        let hash = B256::repeat_byte(recipient);
476        let signed = TransactionSigned::Legacy(Signed::new_unchecked(
477            transaction,
478            Signature::test_signature(),
479            hash,
480        ));
481        let sender = Address::repeat_byte(0xAA);
482        PoolTransaction { hash, sender, transaction: Recovered::new_unchecked(signed, sender) }
483    }
484
485    fn wait_until(what: &str, condition: impl Fn() -> bool) {
486        let deadline = Instant::now() + WAIT_LIMIT;
487        while !condition() {
488            assert!(Instant::now() < deadline, "timed out waiting until {what}");
489            thread::sleep(POLL_INTERVAL);
490        }
491    }
492
493    #[test]
494    fn warms_pool_transactions_into_a_published_snapshot() {
495        let harness = Harness::spawn();
496        let parent_hash = B256::repeat_byte(0x01);
497
498        // Start before the pool tracks the head, covering the poll-and-retry path.
499        harness.start(parent_hash);
500        wait_until("the untracked head is polled", || {
501            harness.pool.not_ready.load(Ordering::Relaxed) >= 1
502        });
503        harness.pool.push(parent_hash, transfer(0xB0));
504
505        let snapshot = harness.published_for(parent_hash);
506        let (accounts, _, _) = snapshot.entry_counts();
507        assert!(accounts >= 1, "speculative execution should cache account reads");
508    }
509
510    #[test]
511    fn pause_quiesces_the_worker_until_resume() {
512        let harness = Harness::spawn();
513        let parent_hash = B256::repeat_byte(0x01);
514        harness.start(parent_hash);
515        harness.pool.push(parent_hash, transfer(0xB0));
516        let before = harness.published_for(parent_hash).entry_counts();
517
518        harness.pause();
519        harness.pool.push(parent_hash, transfer(0xB1));
520        thread::sleep(REFRESH_INTERVAL * 2);
521        assert_eq!(
522            harness.published_entry_counts(),
523            Some(before),
524            "a paused worker must not publish"
525        );
526
527        harness.resume();
528        harness.published(|snapshot| snapshot.entry_counts() != before);
529    }
530
531    #[test]
532    fn overlapping_pauses_require_matching_resumes() {
533        let harness = Harness::spawn();
534        let parent_hash = B256::repeat_byte(0x01);
535        harness.start(parent_hash);
536        harness.pool.push(parent_hash, transfer(0xB0));
537        let before = harness.published_for(parent_hash).entry_counts();
538
539        harness.pause();
540        harness.pause();
541        harness.pool.push(parent_hash, transfer(0xB1));
542        harness.resume();
543        thread::sleep(REFRESH_INTERVAL * 2);
544        assert_eq!(
545            harness.published_entry_counts(),
546            Some(before),
547            "one resume must not release two pauses"
548        );
549
550        harness.resume();
551        harness.published(|snapshot| snapshot.entry_counts() != before);
552    }
553
554    #[test]
555    fn reuses_the_iterator_per_head_and_reopens_on_switch() {
556        let harness = Harness::spawn();
557        let first = B256::repeat_byte(0x01);
558        let second = B256::repeat_byte(0x02);
559
560        harness.start(first);
561        harness.pool.push(first, transfer(0xB0));
562        let before = harness.published_for(first).entry_counts();
563        harness.pool.push(first, transfer(0xB1));
564        harness.published(|snapshot| snapshot.entry_counts() != before);
565        assert_eq!(harness.pool.opened.load(Ordering::Relaxed), 1, "one iterator per head");
566
567        harness.start(second);
568        harness.pool.push(second, transfer(0xB2));
569        harness.published_for(second);
570        assert_eq!(harness.pool.opened.load(Ordering::Relaxed), 2);
571    }
572
573    #[test]
574    fn newest_start_wins() {
575        let harness = Harness::spawn();
576        let stale = B256::repeat_byte(0x01);
577        let newest = B256::repeat_byte(0x02);
578
579        harness.start(stale);
580        harness.start(newest);
581        harness.pool.push(newest, transfer(0xB0));
582
583        harness.published_for(newest);
584    }
585
586    #[test]
587    fn shuts_down_when_control_is_dropped() {
588        let harness = Harness::spawn();
589        let parent_hash = B256::repeat_byte(0x01);
590        harness.start(parent_hash);
591        harness.pool.push(parent_hash, transfer(0xB0));
592        harness.published_for(parent_hash);
593
594        harness.shutdown();
595    }
596}
reth_engine_tree/tree/txpool_prewarm/worker.rs

reth_engine_tree/tree/txpool_prewarm/
worker.rs
