reth_exex/wal/
storage.rs

1use std::{
2    fs::File,
3    io::{BufReader, BufWriter, Write},
4    path::{Path, PathBuf},
5};
6
7use crate::wal::{WalError, WalResult};
8use reth_ethereum_primitives::EthPrimitives;
9use reth_exex_types::ExExNotification;
10use reth_node_api::NodePrimitives;
11use reth_tracing::tracing::debug;
12use tracing::instrument;
13
14static FILE_EXTENSION: &str = "wal";
15
16/// The underlying WAL storage backed by a directory of files.
17///
18/// Each notification is represented by a single file that contains a MessagePack-encoded
19/// notification.
20#[derive(Debug, Clone)]
21pub struct Storage<N: NodePrimitives = EthPrimitives> {
22    /// The path to the WAL file.
23    path: PathBuf,
24    _pd: std::marker::PhantomData<N>,
25}
26
27impl<N> Storage<N>
28where
29    N: NodePrimitives,
30{
31    /// Creates a new instance of [`Storage`] backed by the file at the given path and creates
32    /// it doesn't exist.
33    pub(super) fn new(path: impl AsRef<Path>) -> WalResult<Self> {
34        reth_fs_util::create_dir_all(&path)?;
35
36        Ok(Self { path: path.as_ref().to_path_buf(), _pd: std::marker::PhantomData })
37    }
38
39    fn file_path(&self, id: u32) -> PathBuf {
40        self.path.join(format!("{id}.{FILE_EXTENSION}"))
41    }
42
43    fn parse_filename(filename: &str) -> WalResult<u32> {
44        filename
45            .strip_suffix(".wal")
46            .and_then(|s| s.parse().ok())
47            .ok_or_else(|| WalError::Parse(filename.to_string()))
48    }
49
50    /// Removes notification for the given file ID from the storage.
51    ///
52    /// # Returns
53    ///
54    /// The size of the file that was removed in bytes, if any.
55    #[instrument(skip(self))]
56    fn remove_notification(&self, file_id: u32) -> Option<u64> {
57        let path = self.file_path(file_id);
58        let size = path.metadata().ok()?.len();
59
60        match reth_fs_util::remove_file(self.file_path(file_id)) {
61            Ok(()) => {
62                debug!(target: "exex::wal::storage", "Notification was removed from the storage");
63                Some(size)
64            }
65            Err(err) => {
66                debug!(target: "exex::wal::storage", ?err, "Failed to remove notification from the storage");
67                None
68            }
69        }
70    }
71
72    /// Returns the file IDs in the storage in ascending order.
73    pub(super) fn file_ids(&self) -> WalResult<Vec<u32>> {
74        let mut file_ids = Vec::new();
75
76        for entry in reth_fs_util::read_dir(&self.path)? {
77            let entry = entry.map_err(|err| WalError::DirEntry(self.path.clone(), err))?;
78
79            if entry.path().extension() == Some(FILE_EXTENSION.as_ref()) {
80                let file_name = entry.file_name();
81                let file_id = Self::parse_filename(&file_name.to_string_lossy())?;
82                file_ids.push(file_id);
83            }
84        }
85
86        file_ids.sort_unstable();
87        Ok(file_ids)
88    }
89
90    /// Removes notifications from the storage according to the given list of file IDs.
91    ///
92    /// # Returns
93    ///
94    /// Number of removed notifications and the total size of the removed files in bytes.
95    pub(super) fn remove_notifications(
96        &self,
97        file_ids: impl IntoIterator<Item = u32>,
98    ) -> WalResult<(usize, u64)> {
99        let mut deleted_total = 0;
100        let mut deleted_size = 0;
101
102        for id in file_ids {
103            if let Some(size) = self.remove_notification(id) {
104                deleted_total += 1;
105                deleted_size += size;
106            }
107        }
108
109        Ok((deleted_total, deleted_size))
110    }
111
112    pub(super) fn iter_notifications<'a>(
113        &'a self,
114        file_ids: impl IntoIterator<Item = u32> + 'a,
115    ) -> impl Iterator<Item = WalResult<(u32, u64, ExExNotification<N>)>> + 'a {
116        file_ids.into_iter().map(move |id| {
117            let (notification, size) =
118                self.read_notification(id)?.ok_or(WalError::FileNotFound(id))?;
119
120            Ok((id, size, notification))
121        })
122    }
123
124    /// Reads the notification from the file with the given ID.
125    #[instrument(skip(self))]
126    pub(super) fn read_notification(
127        &self,
128        file_id: u32,
129    ) -> WalResult<Option<(ExExNotification<N>, u64)>> {
130        let file_path = self.file_path(file_id);
131        debug!(target: "exex::wal::storage", ?file_path, "Reading notification from WAL");
132
133        let mut file = match File::open(&file_path) {
134            Ok(file) => file,
135            Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
136            Err(err) => return Err(reth_fs_util::FsPathError::open(err, &file_path).into()),
137        };
138        let size = file.metadata().map_err(|err| WalError::FileMetadata(file_id, err))?.len();
139
140        // Deserialize using the bincode- and msgpack-compatible serde wrapper
141        let notification: reth_exex_types::serde_bincode_compat::ExExNotification<'_, N> =
142            rmp_serde::decode::from_read(BufReader::new(&mut file))
143                .map_err(|err| WalError::Decode(file_id, file_path, err))?;
144
145        Ok(Some((notification.into(), size)))
146    }
147
148    /// Writes the notification to the file with the given ID.
149    ///
150    /// # Returns
151    ///
152    /// The size of the file that was written in bytes.
153    #[instrument(skip(self, notification))]
154    pub(super) fn write_notification(
155        &self,
156        file_id: u32,
157        notification: &ExExNotification<N>,
158    ) -> WalResult<u64> {
159        let file_path = self.file_path(file_id);
160        debug!(target: "exex::wal::storage", ?file_path, "Writing notification to WAL");
161
162        // Serialize using the bincode- and msgpack-compatible serde wrapper
163        let notification =
164            reth_exex_types::serde_bincode_compat::ExExNotification::<N>::from(notification);
165
166        reth_fs_util::atomic_write_file(&file_path, |file| {
167            let mut writer = BufWriter::new(file);
168            rmp_serde::encode::write(&mut writer, &notification)?;
169            // a `BufWriter` dropped without an explicit flush discards write errors, and
170            // `atomic_write_file` fsyncs as soon as this returns
171            writer.flush()?;
172            Ok::<_, Box<dyn core::error::Error + Send + Sync>>(())
173        })?;
174
175        Ok(file_path.metadata().map_err(|err| WalError::FileMetadata(file_id, err))?.len())
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::Storage;
182    use alloy_consensus::BlockHeader;
183    use alloy_primitives::{
184        map::{HashMap, HashSet},
185        B256, U256,
186    };
187    use reth_exex_types::ExExNotification;
188    use reth_primitives_traits::Account;
189    use reth_provider::Chain;
190    use reth_testing_utils::generators::{self, random_block};
191    use reth_trie_common::{
192        serde_bincode_compat,
193        updates::{StorageTrieUpdates, StorageTrieUpdatesSorted, TrieUpdates},
194        BranchNodeCompact, ComputedTrieData, HashedPostState, HashedStorage, HashedStorageSorted,
195        LazyTrieData, Nibbles,
196    };
197    use std::{collections::BTreeMap, fs::File, sync::Arc};
198
199    #[test]
200    fn test_roundtrip() -> eyre::Result<()> {
201        let mut rng = generators::rng();
202
203        let temp_dir = tempfile::tempdir()?;
204        let storage: Storage = Storage::new(&temp_dir)?;
205
206        let old_block = random_block(&mut rng, 0, Default::default()).try_recover()?;
207        let new_block = random_block(&mut rng, 0, Default::default()).try_recover()?;
208
209        let notification = ExExNotification::ChainReorged {
210            new: Arc::new(Chain::new(vec![new_block], Default::default(), BTreeMap::new())),
211            old: Arc::new(Chain::new(vec![old_block], Default::default(), BTreeMap::new())),
212        };
213
214        // Do a round trip serialization and deserialization
215        let file_id = 0;
216        storage.write_notification(file_id, &notification)?;
217        let deserialized_notification = storage.read_notification(file_id)?;
218        assert_eq!(
219            deserialized_notification.map(|(notification, _)| notification),
220            Some(notification)
221        );
222
223        Ok(())
224    }
225
226    #[test]
227    fn test_decode_legacy_sorted_trie_data() -> eyre::Result<()> {
228        let storage_nodes =
229            vec![(Nibbles::from_nibbles_unchecked([0x01]), Some(BranchNodeCompact::default()))];
230        let encoded = rmp_serde::encode::to_vec(&(false, &storage_nodes))?;
231        let decoded: serde_bincode_compat::updates::StorageTrieUpdatesSorted<'_> =
232            rmp_serde::decode::from_slice(&encoded)?;
233        let decoded: StorageTrieUpdatesSorted = decoded.into();
234        assert_eq!(decoded.storage_nodes, storage_nodes);
235
236        let storage_slots = vec![(B256::from([1; 32]), U256::from(1))];
237        let encoded = rmp_serde::encode::to_vec(&(&storage_slots, false))?;
238        let decoded: serde_bincode_compat::hashed_state::HashedStorageSorted<'_> =
239            rmp_serde::decode::from_slice(&encoded)?;
240        let decoded: HashedStorageSorted = decoded.into();
241        assert_eq!(decoded.storage_slots, storage_slots);
242
243        Ok(())
244    }
245
246    /// Generate a new WAL file for testing.
247    ///
248    /// Run this test with `--ignored` to generate a new test WAL file:
249    /// ```sh
250    /// cargo test -p reth-exex generate_test_wal -- --ignored --nocapture
251    /// ```
252    #[test]
253    #[ignore]
254    fn generate_test_wal() -> eyre::Result<()> {
255        use std::io::Write;
256
257        let notification = get_test_notification_data()?;
258
259        // Serialize the notification
260        let notification_compat =
261            reth_exex_types::serde_bincode_compat::ExExNotification::from(&notification);
262        let encoded = rmp_serde::encode::to_vec(&notification_compat)?;
263
264        // Write to test-data directory
265        let test_data_dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("test-data");
266        std::fs::create_dir_all(&test_data_dir)?;
267
268        let output_path = test_data_dir.join("new_format.wal");
269        let mut file = File::create(&output_path)?;
270        file.write_all(&encoded)?;
271
272        println!("Generated WAL file at: {}", output_path.display());
273        println!("File size: {} bytes", encoded.len());
274        println!("✓ WAL file created successfully!");
275
276        Ok(())
277    }
278
279    /// Helper function to generate deterministic test data for WAL tests
280    fn get_test_notification_data(
281    ) -> eyre::Result<ExExNotification<reth_ethereum_primitives::EthPrimitives>> {
282        use reth_ethereum_primitives::Block;
283        use reth_primitives_traits::Block as _;
284
285        // Create a block with a transaction
286        let block = Block::default().seal_slow().try_recover()?;
287        let block_number = block.header().number();
288
289        let hashed_address = B256::from([1; 32]);
290        let storage_key = B256::from([2; 32]);
291
292        let trie_updates = TrieUpdates {
293            account_nodes: HashMap::from_iter([
294                (Nibbles::from_nibbles_unchecked([0x01]), BranchNodeCompact::default()),
295                (Nibbles::from_nibbles_unchecked([0x02]), BranchNodeCompact::default()),
296            ]),
297            removed_nodes: HashSet::from_iter([Nibbles::from_nibbles_unchecked([0x03])]),
298            storage_tries: HashMap::from_iter([(
299                hashed_address,
300                StorageTrieUpdates {
301                    storage_nodes: HashMap::from_iter([(
302                        Nibbles::from_nibbles_unchecked([0x04]),
303                        BranchNodeCompact::default(),
304                    )]),
305                    removed_nodes: Default::default(),
306                },
307            )]),
308        };
309
310        let hashed_state = HashedPostState {
311            accounts: HashMap::from_iter([(
312                hashed_address,
313                Some(Account { nonce: 1, ..Default::default() }),
314            )]),
315            storages: HashMap::from_iter([(
316                hashed_address,
317                HashedStorage { storage: HashMap::from_iter([(storage_key, U256::from(101))]) },
318            )]),
319        };
320
321        let trie_data = LazyTrieData::ready(ComputedTrieData::new(
322            Arc::new(hashed_state.into_sorted()),
323            Arc::new(trie_updates.into_sorted()),
324        ));
325
326        let notification: ExExNotification<reth_ethereum_primitives::EthPrimitives> =
327            ExExNotification::ChainCommitted {
328                new: Arc::new(Chain::new(
329                    vec![block],
330                    Default::default(),
331                    BTreeMap::from([(block_number, trie_data)]),
332                )),
333            };
334        Ok(notification)
335    }
336
337    #[test]
338    fn test_file_ids() -> eyre::Result<()> {
339        let temp_dir = tempfile::tempdir()?;
340        let storage: Storage = Storage::new(&temp_dir)?;
341
342        // Create WAL files
343        File::create(storage.file_path(1))?;
344        File::create(storage.file_path(3))?;
345
346        // Create non-WAL files that should be ignored
347        File::create(temp_dir.path().join("0.tmp"))?;
348        File::create(temp_dir.path().join("4.tmp"))?;
349
350        // Check existing file IDs are returned in order without filling the gap
351        assert_eq!(storage.file_ids()?, vec![1, 3]);
352
353        Ok(())
354    }
355}
reth_exex/wal/storage.rs

reth_exex/wal/
storage.rs
