feat(anvil): add load/dump state options (#3730)

anvil/src/cmd.rs

@@ -1,22 +1,30 @@
 use crate::{
-    config::DEFAULT_MNEMONIC, eth::pool::transactions::TransactionOrder, genesis::Genesis,
+    config::DEFAULT_MNEMONIC,
+    eth::{backend::db::SerializableState, pool::transactions::TransactionOrder, EthApi},
+    genesis::Genesis,
     AccountGenerator, Hardfork, NodeConfig, CHAIN_ID,
 };
 use anvil_server::ServerConfig;
 use clap::Parser;
 use core::fmt;
 use ethers::utils::WEI_IN_ETHER;
 use foundry_config::Chain;
+use futures::FutureExt;
 use std::{
+    future::Future,
     net::IpAddr,
+    path::PathBuf,
+    pin::Pin,
     str::FromStr,
     sync::{
         atomic::{AtomicUsize, Ordering},
         Arc,
     },
+    task::{Context, Poll},
     time::Duration,
 };
-use tracing::log::trace;
+use tokio::time::{Instant, Interval};
+use tracing::{error, trace};
 
 #[derive(Clone, Debug, Parser)]
 pub struct NodeArgs {
@@ -126,6 +134,40 @@ pub struct NodeArgs {
     )]
     pub init: Option<Genesis>,
 
+    #[clap(
+        long,
+        help = "This is an alias for bot --load-state and --dump-state. It initializes the chain with the state stored at the file, if it exists, and dumps the chain's state on exit",
+        value_name = "PATH",
+        value_parser = StateFile::parse,
+        conflicts_with_all = &["init", "dump_state", "load_state"]
+    )]
+    pub state: Option<StateFile>,
+
+    #[clap(
+        short,
+        long,
+        help = "Interval in seconds at which the status is to be dumped to disk. See --state and --dump-state",
+        value_name = "SECONDS"
+    )]
+    pub state_interval: Option<u64>,
+
+    #[clap(
+        long,
+        help = "Dump the state of chain on exit to the given file. If the value is a directory, the state will be written to `<VALUE>/state.json`.",
+        value_name = "PATH",
+        conflicts_with = "init"
+    )]
+    pub dump_state: Option<PathBuf>,
+
+    #[clap(
+        long,
+        help = "Initialize the chain from a previously saved state snapshot.",
+        value_name = "PATH",
+        value_parser = SerializableState::parse,
+        conflicts_with = "init"
+    )]
+    pub load_state: Option<SerializableState>,
+
     #[clap(
         long,
         help = IPC_HELP,
@@ -146,6 +188,9 @@ const IPC_HELP: &str =
 #[cfg(not(windows))]
 const IPC_HELP: &str = "Launch an ipc server at the given path or default path = `/tmp/anvil.ipc`";
 
+/// Default interval for periodically dumping the state.
+const DEFAULT_DUMP_INTERVAL: Duration = Duration::from_secs(60);
+
 impl NodeArgs {
     pub fn into_node_config(self) -> NodeConfig {
         let genesis_balance = WEI_IN_ETHER.saturating_mul(self.balance.into());
@@ -159,7 +204,7 @@ impl NodeArgs {
             .with_gas_limit(self.evm_opts.gas_limit)
             .with_gas_price(self.evm_opts.gas_price)
             .with_hardfork(self.hardfork)
-            .with_blocktime(self.block_time.map(std::time::Duration::from_secs))
+            .with_blocktime(self.block_time.map(Duration::from_secs))
             .with_no_mining(self.no_mining)
             .with_account_generator(self.account_generator())
             .with_genesis_balance(genesis_balance)
@@ -188,6 +233,8 @@ impl NodeArgs {
             .with_ipc(self.ipc)
             .with_code_size_limit(self.evm_opts.code_size_limit)
             .set_pruned_history(self.prune_history)
+            .with_init_state(self.load_state)
+            .with_init_state(self.state.and_then(|s| s.state))
     }
 
     fn account_generator(&self) -> AccountGenerator {
@@ -203,10 +250,19 @@ impl NodeArgs {
         gen
     }
 
+    /// Returns the location where to dump the state to.
+    fn dump_state_path(&self) -> Option<PathBuf> {
+        self.dump_state.as_ref().or_else(|| self.state.as_ref().map(|s| &s.path)).cloned()
+    }
+
     /// Starts the node
     ///
     /// See also [crate::spawn()]
     pub async fn run(self) -> Result<(), Box<dyn std::error::Error>> {
+        let dump_state = self.dump_state_path();
+        let dump_interval =
+            self.state_interval.map(Duration::from_secs).unwrap_or(DEFAULT_DUMP_INTERVAL);
+
         let (api, mut handle) = crate::spawn(self.into_node_config()).await;
 
         // sets the signal handler to gracefully shutdown.
@@ -218,10 +274,20 @@ impl NodeArgs {
         let mut signal = handle.shutdown_signal_mut().take();
 
         let task_manager = handle.task_manager();
-        let on_shutdown = task_manager.on_shutdown();
+        let mut on_shutdown = task_manager.on_shutdown();
+
+        let mut state_dumper = PeriodicStateDumper::new(api, dump_state, dump_interval);
 
         task_manager.spawn(async move {
-            on_shutdown.await;
+            // await shutdown signal but also periodically flush state
+            tokio::select! {
+                _ = &mut on_shutdown =>{}
+                _ = &mut state_dumper =>{}
+            }
+
+            // shutdown received
+            state_dumper.dump().await;
+
             // cleaning up and shutting down
             // this will make sure that the fork RPC cache is flushed if caching is configured
             if let Some(fork) = fork.take() {
@@ -369,6 +435,115 @@ pub struct AnvilEvmArgs {
     pub steps_tracing: bool,
 }
 
+/// Helper type to periodically dump the state of the chain to disk
+struct PeriodicStateDumper {
+    in_progress_dump: Option<Pin<Box<dyn Future<Output = ()> + Send + Sync + 'static>>>,
+    api: EthApi,
+    dump_state: Option<PathBuf>,
+    interval: Interval,
+}
+
+impl PeriodicStateDumper {
+    fn new(api: EthApi, dump_state: Option<PathBuf>, interval: Duration) -> Self {
+        let dump_state = dump_state.map(|mut dump_state| {
+            if dump_state.is_dir() {
+                dump_state = dump_state.join("state.json");
+            }
+            dump_state
+        });
+
+        // periodically flush the state
+        let interval = tokio::time::interval_at(Instant::now() + interval, interval);
+        Self { in_progress_dump: None, api, dump_state, interval }
+    }
+
+    async fn dump(&self) {
+        if let Some(state) = self.dump_state.clone() {
+            Self::dump_state(self.api.clone(), state).await
+        }
+    }
+
+    /// Infallible state dump
+    async fn dump_state(api: EthApi, dump_state: PathBuf) {
+        trace!(path=?dump_state, "Dumping state on shutdown");
+        match api.serialized_state().await {
+            Ok(state) => {
+                if let Err(err) = foundry_common::fs::write_json_file(&dump_state, &state) {
+                    error!(?err, "Failed to dump state");
+                } else {
+                    trace!(path=?dump_state, "Dumped state on shutdown");
+                }
+            }
+            Err(err) => {
+                error!(?err, "Failed to extract state");
+            }
+        }
+    }
+}
+
+// An endless future that periodically dumps the state to disk if configured.
+impl Future for PeriodicStateDumper {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.get_mut();
+        if this.dump_state.is_none() {
+            return Poll::Pending
+        }
+
+        loop {
+            if let Some(mut flush) = this.in_progress_dump.take() {
+                match flush.poll_unpin(cx) {
+                    Poll::Ready(_) => {
+                        this.interval.reset();
+                    }
+                    Poll::Pending => {
+                        this.in_progress_dump = Some(flush);
+                        return Poll::Pending
+                    }
+                }
+            }
+
+            if this.interval.poll_tick(cx).is_ready() {
+                let api = this.api.clone();
+                let path = this.dump_state.clone().expect("exists; see above");
+                this.in_progress_dump =
+                    Some(Box::pin(async move { PeriodicStateDumper::dump_state(api, path).await }));
+            } else {
+                break
+            }
+        }
+
+        Poll::Pending
+    }
+}
+
+/// Represents the --state flag and where to load from, or dump the state to
+#[derive(Debug, Clone)]
+pub struct StateFile {
+    pub path: PathBuf,
+    pub state: Option<SerializableState>,
+}
+
+impl StateFile {
+    /// This is used as the clap `value_parser` implementation to parse from file but only if it
+    /// exists
+    fn parse(path: &str) -> Result<Self, String> {
+        let mut path = PathBuf::from(path);
+        if path.is_dir() {
+            path = path.join("state.json");
+        }
+        let mut state = Self { path, state: None };
+        if !state.path.exists() {
+            return Ok(state)
+        }
+
+        state.state = Some(SerializableState::load(&state.path).map_err(|err| err.to_string())?);
+
+        Ok(state)
+    }
+}
+
 /// Represents the input URL for a fork with an optional trailing block number:
 /// `http://localhost:8545@1000000`
 #[derive(Debug, Clone, PartialEq, Eq)]

anvil/src/config.rs

@@ -1,7 +1,7 @@
 use crate::{
     eth::{
         backend::{
-            db::Db,
+            db::{Db, SerializableState},
             fork::{ClientFork, ClientForkConfig},
             genesis::GenesisConfig,
             mem::fork_db::ForkedDatabase,
@@ -143,6 +143,8 @@ pub struct NodeConfig {
     pub code_size_limit: Option<usize>,
     /// If set to true, remove historic state entirely
     pub prune_history: bool,
+    /// The file where to load the state from
+    pub init_state: Option<SerializableState>,
 }
 
 impl NodeConfig {
@@ -359,6 +361,7 @@ impl Default for NodeConfig {
             ipc_path: None,
             code_size_limit: None,
             prune_history: false,
+            init_state: None,
         }
     }
 }
@@ -388,6 +391,13 @@ impl NodeConfig {
         self
     }
 
+    /// Sets a custom code size limit
+    #[must_use]
+    pub fn with_init_state(mut self, init_state: Option<SerializableState>) -> Self {
+        self.init_state = init_state;
+        self
+    }
+
     /// Sets the chain ID
     #[must_use]
     pub fn with_chain_id<U: Into<u64>>(mut self, chain_id: Option<U>) -> Self {
@@ -891,7 +901,7 @@ impl NodeConfig {
         };
 
         // only memory based backend for now
-        mem::Backend::with_genesis(
+        let backend = mem::Backend::with_genesis(
             db,
             Arc::new(RwLock::new(env)),
             genesis,
@@ -900,7 +910,18 @@ impl NodeConfig {
             self.enable_steps_tracing,
             self.prune_history,
         )
-        .await
+        .await;
+
+        if let Some(ref state) = self.init_state {
+            backend
+                .get_db()
+                .write()
+                .await
+                .load_state(state.clone())
+                .expect("Failed to load init state");
+        }
+
+        backend
     }
 }
 

anvil/src/eth/api.rs

@@ -2,7 +2,7 @@ use crate::{
     eth::{
         backend,
         backend::{
-            mem::MIN_TRANSACTION_GAS, notifications::NewBlockNotifications,
+            db::SerializableState, mem::MIN_TRANSACTION_GAS, notifications::NewBlockNotifications,
             validate::TransactionValidator,
         },
         error::{
@@ -1468,6 +1468,11 @@ impl EthApi {
         self.backend.dump_state().await
     }
 
+    /// Returns the current state
+    pub async fn serialized_state(&self) -> Result<SerializableState> {
+        self.backend.serialized_state().await
+    }
+
     /// Append chain state buffer to current chain. Will overwrite any conflicting addresses or
     /// storage.
     ///

anvil/src/eth/backend/db.rs

@@ -8,6 +8,7 @@ use ethers::{
     utils::keccak256,
 };
 use forge::revm::KECCAK_EMPTY;
+use foundry_common::errors::FsPathError;
 use foundry_evm::{
     executor::{
         backend::{snapshot::StateSnapshot, DatabaseError, DatabaseResult, MemDb},
@@ -21,7 +22,7 @@ use foundry_evm::{
 };
 use hash_db::HashDB;
 use serde::{Deserialize, Serialize};
-use std::fmt;
+use std::{fmt, path::Path};
 
 /// Type alias for the `HashDB` representation of the Database
 pub type AsHashDB = Box<dyn HashDB<KeccakHasher, Vec<u8>>>;
@@ -278,6 +279,25 @@ pub struct SerializableState {
     pub accounts: HashMap<Address, SerializableAccountRecord>,
 }
 
+// === impl SerializableState ===
+
+impl SerializableState {
+    /// Loads the `Genesis` object from the given json file path
+    pub fn load(path: impl AsRef<Path>) -> Result<Self, FsPathError> {
+        let path = path.as_ref();
+        if path.is_dir() {
+            foundry_common::fs::read_json_file(&path.join("state.json"))
+        } else {
+            foundry_common::fs::read_json_file(path)
+        }
+    }
+
+    /// This is used as the clap `value_parser` implementation
+    pub(crate) fn parse(path: &str) -> Result<Self, String> {
+        Self::load(path).map_err(|err| err.to_string())
+    }
+}
+
 #[derive(Serialize, Deserialize, Clone, Debug)]
 pub struct SerializableAccountRecord {
     pub nonce: u64,