Skip to main content

reth_engine_tree/tree/txpool_prewarm/
worker.rs

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