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
28const REFRESH_INTERVAL: Duration = Duration::from_millis(100);
30
31const HEAD_POLL_INTERVAL: Duration = Duration::from_millis(10);
33
34pub(super) struct Worker<N, P, Evm>
46where
47 N: NodePrimitives,
48 Evm: ConfigureEvm<Primitives = N>,
49{
50 commands: Receiver<Command<Job<N, P, Evm>>>,
52 publication: Publication,
54 source: Arc<dyn Source<N>>,
56 evm_config: Evm,
58 job: Option<(B256, Job<N, P, Evm>)>,
60 pauses: u64,
62 cache: CachedReads,
64 cache_parent: Option<B256>,
66 published_entries: (usize, usize, usize),
69 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 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 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 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 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 fn open_transactions(&mut self, parent_hash: B256) -> bool {
169 if self.transactions.as_ref().is_none_or(|(parent, _)| *parent != parent_hash) {
170 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 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 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 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 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 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 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 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 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
340enum BatchEnd {
341 GoAgain,
344 Rest,
347}
348
349#[derive(Debug, Clone, Copy, PartialEq, Eq)]
351struct ChannelDisconnected;
352
353fn 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 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 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 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 fn start(&self, parent_hash: B256) {
421 self.start_with(parent_hash, None);
422 }
423
424 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 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 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 let (disconnected, _) = unbounded();
489 self.commands = disconnected;
490 if let Some(worker) = self.worker.take() {
491 let _ = worker.join();
492 }
493 }
494 }
495
496 #[derive(Debug, Default)]
499 struct ScriptedPool {
500 queues: Arc<Mutex<HashMap<B256, VecDeque<PoolTransaction<EthPrimitives>>>>>,
501 opened: AtomicUsize,
503 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 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 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}