reth_engine_util/
engine_store.rs

Help
1//! Stores engine API messages to disk for later inspection and replay.
2
3use alloy_rpc_types_engine::ForkchoiceState;
4use futures::{Stream, StreamExt};
5use reth_engine_primitives::{BeaconEngineMessage, EngineTypes, ExecutionPayload};
6use reth_fs_util as fs;
7use serde::{Deserialize, Serialize};
8use std::{
9    collections::BTreeMap,
10    path::PathBuf,
11    pin::Pin,
12    task::{ready, Context, Poll},
13    time::SystemTime,
14};
15use tracing::*;
16
17/// A message from the engine API that has been stored to disk.
18#[derive(Debug, Serialize, Deserialize)]
19#[serde(rename_all = "camelCase")]
20pub enum StoredEngineApiMessage<EngineT: EngineTypes> {
21    /// The on-disk representation of an `engine_forkchoiceUpdated` method call.
22    ForkchoiceUpdated {
23        /// The [`ForkchoiceState`] sent in the persisted call.
24        state: ForkchoiceState,
25        /// The payload attributes sent in the persisted call, if any.
26        payload_attrs: Option<EngineT::PayloadAttributes>,
27    },
28    /// The on-disk representation of an `engine_newPayload` method call.
29    NewPayload {
30        /// The [`EngineTypes::ExecutionData`] sent in the persisted call.
31        #[serde(flatten)]
32        payload: EngineT::ExecutionData,
33    },
34}
35
36/// This can read and write engine API messages in a specific directory.
37#[derive(Debug)]
38pub struct EngineMessageStore {
39    /// The path to the directory that stores the engine API messages.
40    path: PathBuf,
41}
42
43impl EngineMessageStore {
44    /// Creates a new [`EngineMessageStore`] at the given path.
45    ///
46    /// The path is expected to be a directory, where individual message JSON files will be stored.
47    pub const fn new(path: PathBuf) -> Self {
48        Self { path }
49    }
50
51    /// Stores the received [`BeaconEngineMessage`] to disk, appending the `received_at` time to the
52    /// path.
53    pub fn on_message<Engine>(
54        &self,
55        msg: &BeaconEngineMessage<Engine>,
56        received_at: SystemTime,
57    ) -> eyre::Result<()>
58    where
59        Engine: EngineTypes,
60    {
61        fs::create_dir_all(&self.path)?; // ensure that store path had been created
62        let timestamp = received_at.duration_since(SystemTime::UNIX_EPOCH).unwrap().as_millis();
63        match msg {
64            BeaconEngineMessage::ForkchoiceUpdated {
65                state,
66                payload_attrs,
67                tx: _tx,
68                version: _version,
69            } => {
70                let filename = format!("{}-fcu-{}.json", timestamp, state.head_block_hash);
71                fs::write(
72                    self.path.join(filename),
73                    serde_json::to_vec(&StoredEngineApiMessage::<Engine>::ForkchoiceUpdated {
74                        state: *state,
75                        payload_attrs: payload_attrs.clone(),
76                    })?,
77                )?;
78            }
79            BeaconEngineMessage::NewPayload { payload, tx: _tx } => {
80                let filename = format!("{}-new_payload-{}.json", timestamp, payload.block_hash());
81                fs::write(
82                    self.path.join(filename),
83                    serde_json::to_vec(&StoredEngineApiMessage::<Engine>::NewPayload {
84                        payload: payload.clone(),
85                    })?,
86                )?;
87            }
88            // noop
89            BeaconEngineMessage::TransitionConfigurationExchanged => (),
90        };
91        Ok(())
92    }
93
94    /// Finds and iterates through any stored engine API message files, ordered by timestamp.
95    pub fn engine_messages_iter(&self) -> eyre::Result<impl Iterator<Item = PathBuf>> {
96        let mut filenames_by_ts = BTreeMap::<u64, Vec<PathBuf>>::default();
97        for entry in fs::read_dir(&self.path)? {
98            let entry = entry?;
99            let filename = entry.file_name();
100            if let Some(filename) = filename.to_str().filter(|n| n.ends_with(".json")) {
101                if let Some(Ok(timestamp)) = filename.split('-').next().map(|n| n.parse::<u64>()) {
102                    filenames_by_ts.entry(timestamp).or_default().push(entry.path());
103                    tracing::debug!(target: "engine::store", timestamp, filename, "Queued engine API message");
104                } else {
105                    tracing::warn!(target: "engine::store", %filename, "Could not parse timestamp from filename")
106                }
107            } else {
108                tracing::warn!(target: "engine::store", ?filename, "Skipping non json file");
109            }
110        }
111        Ok(filenames_by_ts.into_iter().flat_map(|(_, paths)| paths))
112    }
113}
114
115/// A wrapper stream that stores Engine API messages in
116/// the specified directory.
117#[derive(Debug)]
118#[pin_project::pin_project]
119pub struct EngineStoreStream<S> {
120    /// Inner message stream.
121    #[pin]
122    stream: S,
123    /// Engine message store.
124    store: EngineMessageStore,
125}
126
127impl<S> EngineStoreStream<S> {
128    /// Create new engine store stream wrapper.
129    pub const fn new(stream: S, path: PathBuf) -> Self {
130        Self { stream, store: EngineMessageStore::new(path) }
131    }
132}
133
134impl<S, Engine> Stream for EngineStoreStream<S>
135where
136    S: Stream<Item = BeaconEngineMessage<Engine>>,
137    Engine: EngineTypes,
138{
139    type Item = S::Item;
140
141    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
142        let mut this = self.project();
143        let next = ready!(this.stream.poll_next_unpin(cx));
144        if let Some(msg) = &next {
145            if let Err(error) = this.store.on_message(msg, SystemTime::now()) {
146                error!(target: "engine::stream::store", ?msg, %error, "Error handling Engine API message");
147            }
148        }
149        Poll::Ready(next)
150    }
151}
reth_engine_util/engine_store.rs

reth_engine_util/
engine_store.rs
