reth_e2e_test_utils/
snap.rs

1//! Observes and pauses the snap requests of a node's production backfill in e2e tests.
2//!
3//! [`SnapControl`] records every request and can pause one at a [`SnapGate`], while responses
4//! still come from the real serving peer and go through production validation.
5
6use crate::{wait::poll_until, E2ETestSetupBuilder};
7use parking_lot::Mutex;
8use reth_eth_wire_types::snap::{
9    GetAccountRangeMessage, GetBlockAccessListsMessage, GetByteCodesMessage,
10    GetStorageRangesMessage, SnapProtocolMessage,
11};
12use reth_network_p2p::{
13    download::DownloadClient,
14    error::PeerRequestResult,
15    headers::client::{HeadersClient, HeadersRequest},
16    priority::Priority,
17    snap::client::{SnapClient, SnapResponse},
18};
19use reth_network_peers::PeerId;
20use reth_node_builder::sync::{BackfillContext, BackfillSyncBuilder};
21use reth_node_ethereum::{
22    snap::{EthereumBackfill, EthereumBackfillSync},
23    EthereumNode,
24};
25use reth_provider::{providers::ProviderNodeTypes, ProviderFactory};
26use std::{future::Future, mem, pin::Pin, sync::Arc};
27use tokio::sync::watch;
28
29/// Controls the requests of a node's production snap backfill, including across restarts.
30#[derive(Clone, Debug, Default)]
31pub struct SnapControl {
32    // Byte limit set on account and storage range requests, if any.
33    response_bytes: Option<u64>,
34    // Recorded requests and registered gates, shared by every clone.
35    state: Arc<Mutex<ControlState>>,
36}
37
38impl SnapControl {
39    /// Creates a control that limits account and storage responses to force real peer pagination.
40    pub fn with_response_bytes(bytes: u64) -> Self {
41        Self { response_bytes: Some(bytes), state: Default::default() }
42    }
43
44    /// Pauses the next matching request before it reaches the network.
45    ///
46    /// Only requests sent after this call can match. The gate fires once, so a retry of the paused
47    /// request is not paused again. Dropping the gate releases the request.
48    pub fn pause_on(
49        &self,
50        matches: impl Fn(&SnapProtocolMessage) -> bool + Send + Sync + 'static,
51    ) -> SnapGate {
52        let (state, _) = watch::channel(GateState::Waiting);
53        self.state
54            .lock()
55            .gates
56            .push(RequestGate { matches: Box::new(matches), state: state.clone() });
57        SnapGate { state }
58    }
59
60    /// Returns the requests seen since the last call and clears them, including a request
61    /// currently paused at a gate.
62    pub fn take_requests(&self) -> Vec<SnapProtocolMessage> {
63        mem::take(&mut self.state.lock().requests)
64    }
65}
66
67impl E2ETestSetupBuilder<EthereumNode> {
68    /// Wraps the backfill chosen by the node configuration with `control`.
69    pub fn with_snap_control(self, control: &SnapControl) -> Self {
70        let control = control.clone();
71        self.with_backfill(move |config| ControlledBackfill {
72            inner: EthereumBackfill::new(config),
73            control: control.clone(),
74        })
75    }
76}
77
78/// Builds the production backfill with an observed network client.
79#[derive(Debug)]
80pub struct ControlledBackfill {
81    // Backfill chosen by the node configuration.
82    inner: EthereumBackfill,
83    // Control handed to the network client the backfill is built with.
84    control: SnapControl,
85}
86
87impl<N, C> BackfillSyncBuilder<N, C> for ControlledBackfill
88where
89    N: ProviderNodeTypes,
90    C: SnapClient + HeadersClient + Clone + Unpin + 'static,
91{
92    type Backfill = EthereumBackfillSync<N, ControlledSnapClient<C>>;
93
94    fn build(self, ctx: BackfillContext<N, C>) -> eyre::Result<Self::Backfill> {
95        let client = ControlledSnapClient { inner: ctx.client().clone(), control: self.control };
96        let factory = ctx.provider_factory().clone();
97        let runtime = ctx.runtime().clone();
98        self.inner.build(BackfillContext::new(ctx.into_pipeline(), client, factory, runtime))
99    }
100
101    fn recover(&mut self, factory: &ProviderFactory<N>) -> eyre::Result<()> {
102        <EthereumBackfill as BackfillSyncBuilder<N, C>>::recover(&mut self.inner, factory)
103    }
104}
105
106/// A real network client with request observation and deterministic gates.
107#[derive(Clone, Debug)]
108pub struct ControlledSnapClient<C> {
109    // Network client that sends the requests to peers.
110    inner: C,
111    // Control that records and gates the requests.
112    control: SnapControl,
113}
114
115impl<C: DownloadClient> DownloadClient for ControlledSnapClient<C> {
116    fn report_bad_message(&self, peer_id: PeerId) {
117        self.inner.report_bad_message(peer_id);
118    }
119
120    fn num_connected_peers(&self) -> usize {
121        self.inner.num_connected_peers()
122    }
123}
124
125impl<C: HeadersClient> HeadersClient for ControlledSnapClient<C> {
126    type Header = C::Header;
127    type Output = C::Output;
128
129    fn get_headers_with_priority(
130        &self,
131        request: HeadersRequest,
132        priority: Priority,
133    ) -> Self::Output {
134        self.inner.get_headers_with_priority(request, priority)
135    }
136}
137
138impl<C: SnapClient + Clone + 'static> SnapClient for ControlledSnapClient<C> {
139    type Output = Pin<Box<dyn Future<Output = PeerRequestResult<SnapResponse>> + Send + Sync>>;
140
141    fn get_account_range_with_priority(
142        &self,
143        mut request: GetAccountRangeMessage,
144        priority: Priority,
145    ) -> Self::Output {
146        if let Some(bytes) = self.control.response_bytes {
147            request.response_bytes = bytes;
148        }
149        self.gated(SnapProtocolMessage::GetAccountRange(request.clone()), move |client| {
150            client.get_account_range_with_priority(request, priority)
151        })
152    }
153
154    fn get_storage_ranges(&self, request: GetStorageRangesMessage) -> Self::Output {
155        self.get_storage_ranges_with_priority(request, Priority::Normal)
156    }
157
158    fn get_storage_ranges_with_priority(
159        &self,
160        mut request: GetStorageRangesMessage,
161        priority: Priority,
162    ) -> Self::Output {
163        if let Some(bytes) = self.control.response_bytes {
164            request.response_bytes = bytes;
165        }
166        self.gated(SnapProtocolMessage::GetStorageRanges(request.clone()), move |client| {
167            client.get_storage_ranges_with_priority(request, priority)
168        })
169    }
170
171    fn get_byte_codes(&self, request: GetByteCodesMessage) -> Self::Output {
172        self.get_byte_codes_with_priority(request, Priority::Normal)
173    }
174
175    fn get_byte_codes_with_priority(
176        &self,
177        request: GetByteCodesMessage,
178        priority: Priority,
179    ) -> Self::Output {
180        self.gated(SnapProtocolMessage::GetByteCodes(request.clone()), move |client| {
181            client.get_byte_codes_with_priority(request, priority)
182        })
183    }
184
185    fn get_block_access_lists_with_priority(
186        &self,
187        request: GetBlockAccessListsMessage,
188        priority: Priority,
189    ) -> Self::Output {
190        self.gated(SnapProtocolMessage::GetBlockAccessLists(request.clone()), move |client| {
191            client.get_block_access_lists_with_priority(request, priority)
192        })
193    }
194}
195
196impl<C: SnapClient + Clone + 'static> ControlledSnapClient<C> {
197    // Record and gate before dispatching; all responses still come from the real network.
198    fn gated(
199        &self,
200        request: SnapProtocolMessage,
201        send: impl FnOnce(&C) -> C::Output + Send + Sync + 'static,
202    ) -> <Self as SnapClient>::Output {
203        let released = self.control.state.lock().observe(request);
204        let client = self.inner.clone();
205        Box::pin(async move {
206            if let Some(mut released) = released {
207                // An error means the gate was dropped, which releases the request too.
208                let _ = released.wait_for(|state| matches!(state, GateState::Released)).await;
209            }
210            send(&client).await
211        })
212    }
213}
214
215/// A single request boundary a test can await and release.
216#[derive(Debug)]
217#[must_use = "dropping the gate releases the paused request"]
218pub struct SnapGate {
219    // Progress of the gate, shared with its `RequestGate`.
220    state: watch::Sender<GateState>,
221}
222
223impl SnapGate {
224    /// Waits until the downloader reaches this request, with the harness timeout, and returns it.
225    pub async fn reached(&self) -> eyre::Result<SnapProtocolMessage> {
226        poll_until("a paused snap request", || async {
227            Ok(match &*self.state.borrow() {
228                GateState::Reached(request) => Some(request.clone()),
229                GateState::Waiting | GateState::Released => None,
230            })
231        })
232        .await
233    }
234
235    /// Lets the paused request proceed to the real peer.
236    pub fn release(&self) {
237        self.state.send_replace(GateState::Released);
238    }
239}
240
241// State shared by the clones of a `SnapControl`.
242#[derive(Debug, Default)]
243struct ControlState {
244    // Requests seen since the last `take_requests`.
245    requests: Vec<SnapProtocolMessage>,
246    // Gates still waiting for a matching request, in registration order.
247    gates: Vec<RequestGate>,
248}
249
250impl ControlState {
251    // Signals a matching gate right away: the request future may be dropped before it is polled.
252    fn observe(&mut self, request: SnapProtocolMessage) -> Option<watch::Receiver<GateState>> {
253        let released = self.gates.iter().position(|gate| (gate.matches)(&request)).map(|index| {
254            let gate = self.gates.remove(index);
255            gate.state.send_replace(GateState::Reached(request.clone()));
256            gate.state.subscribe()
257        });
258        self.requests.push(request);
259        released
260    }
261}
262
263// A registered gate waiting for its first matching request.
264#[derive(derive_more::Debug)]
265struct RequestGate {
266    // Selects the request to pause.
267    #[debug(skip)]
268    matches: Box<dyn Fn(&SnapProtocolMessage) -> bool + Send + Sync>,
269    // Progress shared with the test's `SnapGate`.
270    state: watch::Sender<GateState>,
271}
272
273// Progress of a gate from registration to release.
274#[derive(Debug)]
275enum GateState {
276    // No request matched yet.
277    Waiting,
278    // The matched request is paused before it reaches the network.
279    Reached(SnapProtocolMessage),
280    // The test released the request.
281    Released,
282}
reth_e2e_test_utils/snap.rs

reth_e2e_test_utils/
snap.rs
