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

reth_exex/wal/
storage.rs
