reth_bench/bench/
context.rs

1//! This contains the [`BenchContext`], which is information that all replay-based benchmarks need.
2//! The initialization code is also the same, so this can be shared across benchmark commands.
3
4use crate::{authenticated_transport::AuthenticatedTransportConnect, bench_mode::BenchMode};
5use alloy_eips::BlockNumberOrTag;
6use alloy_primitives::address;
7use alloy_provider::{network::AnyNetwork, Provider, RootProvider};
8use alloy_rpc_client::ClientBuilder;
9use alloy_rpc_types_engine::JwtSecret;
10use alloy_transport::layers::RetryBackoffLayer;
11use reqwest::Url;
12use reth_node_core::args::BenchmarkArgs;
13use tracing::info;
14
15/// This is intended to be used by benchmarks that replay blocks from an RPC.
16///
17/// It contains an authenticated provider for engine API queries, a block provider for block
18/// queries, a [`BenchMode`] to determine whether the benchmark should run for a closed or open
19/// range of blocks, and the next block to fetch.
20pub(crate) struct BenchContext {
21    /// The auth provider is used for engine API queries.
22    pub(crate) auth_provider: RootProvider<AnyNetwork>,
23    /// The block provider is used for block queries.
24    pub(crate) block_provider: RootProvider<AnyNetwork>,
25    /// The benchmark mode, which defines whether the benchmark should run for a closed or open
26    /// range of blocks.
27    pub(crate) benchmark_mode: BenchMode,
28    /// The next block to fetch.
29    pub(crate) next_block: u64,
30    /// Whether the chain is an OP rollup.
31    pub(crate) is_optimism: bool,
32}
33
34impl BenchContext {
35    /// This is the initialization code for most benchmarks, taking in a [`BenchmarkArgs`] and
36    /// returning the providers needed to run a benchmark.
37    pub(crate) async fn new(bench_args: &BenchmarkArgs, rpc_url: String) -> eyre::Result<Self> {
38        info!("Running benchmark using data from RPC URL: {}", rpc_url);
39
40        // Ensure that output directory exists and is a directory
41        if let Some(output) = &bench_args.output {
42            if output.is_file() {
43                return Err(eyre::eyre!("Output path must be a directory"));
44            }
45            // Create the directory if it doesn't exist
46            if !output.exists() {
47                std::fs::create_dir_all(output)?;
48                info!("Created output directory: {:?}", output);
49            }
50        }
51
52        // set up alloy client for blocks
53        let client = ClientBuilder::default()
54            .layer(RetryBackoffLayer::new(10, 800, u64::MAX))
55            .http(rpc_url.parse()?);
56        let block_provider = RootProvider::<AnyNetwork>::new(client);
57
58        // Check if this is an OP chain by checking code at a predeploy address.
59        let is_optimism = !block_provider
60            .get_code_at(address!("0x420000000000000000000000000000000000000F"))
61            .await?
62            .is_empty();
63
64        // construct the authenticated provider
65        let auth_jwt = bench_args
66            .auth_jwtsecret
67            .clone()
68            .ok_or_else(|| eyre::eyre!("--jwt-secret must be provided for authenticated RPC"))?;
69
70        // fetch jwt from file
71        //
72        // the jwt is hex encoded so we will decode it after
73        let jwt = std::fs::read_to_string(auth_jwt)?;
74        let jwt = JwtSecret::from_hex(jwt)?;
75
76        // get engine url
77        let auth_url = Url::parse(&bench_args.engine_rpc_url)?;
78
79        // construct the authed transport
80        info!("Connecting to Engine RPC at {} for replay", auth_url);
81        let auth_transport = AuthenticatedTransportConnect::new(auth_url, jwt);
82        let client = ClientBuilder::default().connect_with(auth_transport).await?;
83        let auth_provider = RootProvider::<AnyNetwork>::new(client);
84
85        // Computes the block range for the benchmark.
86        //
87        // - If `--advance` is provided, fetches the latest block and sets:
88        //     - `from = head + 1`
89        //     - `to = head + advance`
90        // - Otherwise, uses the values from `--from` and `--to`.
91        let (from, to) = if let Some(advance) = bench_args.advance {
92            if advance == 0 {
93                return Err(eyre::eyre!("--advance must be greater than 0"));
94            }
95
96            let head_block = auth_provider
97                .get_block_by_number(BlockNumberOrTag::Latest)
98                .await?
99                .ok_or_else(|| eyre::eyre!("Failed to fetch latest block for --advance"))?;
100            let head_number = head_block.header.number;
101            (Some(head_number), Some(head_number + advance))
102        } else {
103            (bench_args.from, bench_args.to)
104        };
105
106        // If neither `--from` nor `--to` are provided, we will run the benchmark continuously,
107        // starting at the latest block.
108        let mut benchmark_mode = BenchMode::new(from, to)?;
109
110        let first_block = match benchmark_mode {
111            BenchMode::Continuous => {
112                // fetch Latest block
113                block_provider.get_block_by_number(BlockNumberOrTag::Latest).full().await?.unwrap()
114            }
115            BenchMode::Range(ref mut range) => {
116                match range.next() {
117                    Some(block_number) => {
118                        // fetch first block in range
119                        block_provider
120                            .get_block_by_number(block_number.into())
121                            .full()
122                            .await?
123                            .unwrap()
124                    }
125                    None => {
126                        return Err(eyre::eyre!(
127                            "Benchmark mode range is empty, please provide a larger range"
128                        ));
129                    }
130                }
131            }
132        };
133
134        let next_block = first_block.header.number + 1;
135        Ok(Self { auth_provider, block_provider, benchmark_mode, next_block, is_optimism })
136    }
137}