anvil/eth/backend/mem/
cache.rs

Help
1use crate::config::anvil_tmp_dir;
2use alloy_primitives::B256;
3use foundry_evm::backend::StateSnapshot;
4use std::{
5    io,
6    path::{Path, PathBuf},
7};
8use tempfile::TempDir;
9
10/// On disk state cache
11///
12/// A basic tempdir which stores states on disk
13pub struct DiskStateCache {
14    /// The path where to create the tempdir in
15    pub(crate) temp_path: Option<PathBuf>,
16    /// Holds the temp dir object.
17    pub(crate) temp_dir: Option<TempDir>,
18}
19
20impl DiskStateCache {
21    /// Specify the path where to create the tempdir in
22    pub fn with_path(self, temp_path: PathBuf) -> Self {
23        Self { temp_path: Some(temp_path), temp_dir: None }
24    }
25    /// Returns the cache file for the given hash
26    fn with_cache_file<F, R>(&mut self, hash: B256, f: F) -> Option<R>
27    where
28        F: FnOnce(PathBuf) -> R,
29    {
30        if self.temp_dir.is_none() {
31            let tmp_dir = self
32                .temp_path
33                .as_ref()
34                .map(|p| -> io::Result<TempDir> {
35                    std::fs::create_dir_all(p)?;
36                    build_tmp_dir(Some(p))
37                })
38                .unwrap_or_else(|| build_tmp_dir(None));
39
40            match tmp_dir {
41                Ok(temp_dir) => {
42                    trace!(target: "backend", path=?temp_dir.path(), "created disk state cache dir");
43                    self.temp_dir = Some(temp_dir);
44                }
45                Err(err) => {
46                    error!(target: "backend", %err, "failed to create disk state cache dir");
47                }
48            }
49        }
50        if let Some(temp_dir) = &self.temp_dir {
51            let path = temp_dir.path().join(format!("{hash:?}.json"));
52            Some(f(path))
53        } else {
54            None
55        }
56    }
57
58    /// Stores the snapshot for the given hash
59    ///
60    /// Note: this writes the state on a new spawned task
61    ///
62    /// Caution: this requires a running tokio Runtime.
63    pub fn write(&mut self, hash: B256, state: StateSnapshot) {
64        self.with_cache_file(hash, |file| {
65            tokio::task::spawn(async move {
66                match foundry_common::fs::write_json_file(&file, &state) {
67                    Ok(_) => {
68                        trace!(target: "backend", ?hash, "wrote state json file");
69                    }
70                    Err(err) => {
71                        error!(target: "backend", %err, ?hash, "Failed to load state snapshot");
72                    }
73                };
74            });
75        });
76    }
77
78    /// Loads the snapshot file for the given hash
79    ///
80    /// Returns None if it doesn't exist or deserialization failed
81    pub fn read(&mut self, hash: B256) -> Option<StateSnapshot> {
82        self.with_cache_file(hash, |file| {
83            match foundry_common::fs::read_json_file::<StateSnapshot>(&file) {
84                Ok(state) => {
85                    trace!(target: "backend", ?hash,"loaded cached state");
86                    Some(state)
87                }
88                Err(err) => {
89                    error!(target: "backend", %err, ?hash, "Failed to load state snapshot");
90                    None
91                }
92            }
93        })
94        .flatten()
95    }
96
97    /// Removes the cache file for the given hash, if it exists
98    pub fn remove(&mut self, hash: B256) {
99        self.with_cache_file(hash, |file| {
100            foundry_common::fs::remove_file(file).map_err(|err| {
101                error!(target: "backend", %err, %hash, "Failed to remove state snapshot");
102            })
103        });
104    }
105}
106
107impl Default for DiskStateCache {
108    fn default() -> Self {
109        Self { temp_path: anvil_tmp_dir(), temp_dir: None }
110    }
111}
112
113/// Returns the temporary dir for the cached state
114///
115/// This will create a prefixed temp dir with `anvil-state-06-11-2022-12-50`
116fn build_tmp_dir(p: Option<&Path>) -> io::Result<TempDir> {
117    let mut builder = tempfile::Builder::new();
118    let now = chrono::offset::Utc::now();
119    let prefix = now.format("anvil-state-%d-%m-%Y-%H-%M").to_string();
120    builder.prefix(&prefix);
121
122    if let Some(p) = p {
123        builder.tempdir_in(p)
124    } else {
125        builder.tempdir()
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132    use tempfile::tempdir;
133
134    #[test]
135    fn can_build_temp_dir() {
136        let dir = tempdir().unwrap();
137        let p = dir.path();
138        let cache_dir = build_tmp_dir(Some(p)).unwrap();
139        assert!(cache_dir
140            .path()
141            .file_name()
142            .unwrap()
143            .to_str()
144            .unwrap()
145            .starts_with("anvil-state-"));
146        let cache_dir = build_tmp_dir(None).unwrap();
147        assert!(cache_dir
148            .path()
149            .file_name()
150            .unwrap()
151            .to_str()
152            .unwrap()
153            .starts_with("anvil-state-"));
154    }
155}
anvil/eth/backend/mem/cache.rs

anvil/eth/backend/mem/
cache.rs
