reth_engine_util/
engine_store.rs

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, ExecutionPayload};
6use reth_fs_util as fs;
7use reth_payload_primitives::PayloadTypes;
8use serde::{Deserialize, Serialize};
9use std::{
10    collections::BTreeMap,
11    path::PathBuf,
12    pin::Pin,
13    task::{ready, Context, Poll},
14    time::SystemTime,
15};
16use tracing::*;
17
18/// A message from the engine API that has been stored to disk.
19#[derive(Debug, Serialize, Deserialize)]
20#[serde(rename_all = "camelCase")]
21pub enum StoredEngineApiMessage<T: PayloadTypes> {
22    /// The on-disk representation of an `engine_forkchoiceUpdated` method call.
23    ForkchoiceUpdated {
24        /// The [`ForkchoiceState`] sent in the persisted call.
25        state: ForkchoiceState,
26        /// The payload attributes sent in the persisted call, if any.
27        payload_attrs: Option<T::PayloadAttributes>,
28    },
29    /// The on-disk representation of an `engine_newPayload` method call.
30    NewPayload {
31        /// The [`PayloadTypes::ExecutionData`] sent in the persisted call.
32        #[serde(flatten)]
33        payload: T::ExecutionData,
34    },
35}
36
37/// This can read and write engine API messages in a specific directory.
38#[derive(Debug)]
39pub struct EngineMessageStore {
40    /// The path to the directory that stores the engine API messages.
41    path: PathBuf,
42}
43
44impl EngineMessageStore {
45    /// Creates a new [`EngineMessageStore`] at the given path.
46    ///
47    /// The path is expected to be a directory, where individual message JSON files will be stored.
48    pub const fn new(path: PathBuf) -> Self {
49        Self { path }
50    }
51
52    /// Stores the received [`BeaconEngineMessage`] to disk, appending the `received_at` time to the
53    /// path.
54    pub fn on_message<T>(
55        &self,
56        msg: &BeaconEngineMessage<T>,
57        received_at: SystemTime,
58    ) -> eyre::Result<()>
59    where
60        T: PayloadTypes,
61    {
62        fs::create_dir_all(&self.path)?; // ensure that store path had been created
63        let timestamp = received_at.duration_since(SystemTime::UNIX_EPOCH).unwrap().as_millis();
64        match msg {
65            BeaconEngineMessage::ForkchoiceUpdated {
66                state,
67                payload_attrs,
68                tx: _tx,
69                version: _version,
70            } => {
71                let filename = format!("{}-fcu-{}.json", timestamp, state.head_block_hash);
72                fs::write(
73                    self.path.join(filename),
74                    serde_json::to_vec(&StoredEngineApiMessage::<T>::ForkchoiceUpdated {
75                        state: *state,
76                        payload_attrs: payload_attrs.clone(),
77                    })?,
78                )?;
79            }
80            BeaconEngineMessage::NewPayload { payload, .. } |
81            BeaconEngineMessage::RethNewPayload { payload, .. } => {
82                let filename = format!("{}-new_payload-{}.json", timestamp, payload.block_hash());
83                fs::write(
84                    self.path.join(filename),
85                    serde_json::to_vec(&StoredEngineApiMessage::<T>::NewPayload {
86                        payload: payload.clone(),
87                    })?,
88                )?;
89            }
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_values().flatten())
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, T> Stream for EngineStoreStream<S>
135where
136    S: Stream<Item = BeaconEngineMessage<T>>,
137    T: PayloadTypes,
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            let Err(error) = this.store.on_message(msg, SystemTime::now())
146        {
147            error!(target: "engine::stream::store", ?msg, %error, "Error handling Engine API message");
148        }
149        Poll::Ready(next)
150    }
151}
reth_engine_util/engine_store.rs

reth_engine_util/
engine_store.rs
