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