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

reth_engine_tree/tree/txpool_prewarm/
worker.rs
