Implement forking support with lazy loading backend (#426)

Author: dimartiro

crates/anvil-polkadot/src/cmd.rs

@@ -1,15 +1,16 @@
 use crate::config::{
-    AccountGenerator, AnvilNodeConfig, CHAIN_ID, DEFAULT_MNEMONIC, SubstrateNodeConfig,
+    AccountGenerator, AnvilNodeConfig, CHAIN_ID, DEFAULT_MNEMONIC, ForkChoice, SubstrateNodeConfig,
 };
 use alloy_genesis::Genesis;
 use alloy_primitives::{U256, utils::Unit};
 use alloy_signer_local::coins_bip39::{English, Mnemonic};
 use anvil_server::ServerConfig;
 use clap::Parser;
+use core::fmt;
 use foundry_common::shell;
 use foundry_config::Chain;
 use rand_08::{SeedableRng, rngs::StdRng};
-use std::{net::IpAddr, path::PathBuf, time::Duration};
+use std::{net::IpAddr, path::PathBuf, str::FromStr, time::Duration};
 
 #[derive(Clone, Debug, Parser)]
 pub struct NodeArgs {
@@ -138,7 +139,16 @@ impl NodeArgs {
             .disable_code_size_limit(self.evm.disable_code_size_limit)
             .with_disable_default_create2_deployer(self.evm.disable_default_create2_deployer)
             .with_memory_limit(self.evm.memory_limit)
-            .with_revive_rpc_block_limit(self.revive_rpc_block_limit);
+            .with_fork_choice(self.evm.fork_block_number.map(ForkChoice::Block).or_else(|| {
+                self.evm
+                    .fork_url
+                    .as_ref()
+                    .and_then(|f| f.block)
+                    .map(|num| ForkChoice::Block(num as i128))
+            }))
+            .with_eth_rpc_url(self.evm.fork_url.map(|fork| fork.url))
+            .fork_request_timeout(self.evm.fork_request_timeout.map(Duration::from_millis))
+            .fork_request_retries(self.evm.fork_request_retries);
 
         let substrate_node_config = SubstrateNodeConfig::new(&anvil_config);
 
@@ -174,6 +184,44 @@ impl NodeArgs {
 #[derive(Clone, Debug, Parser)]
 #[command(next_help_heading = "EVM options")]
 pub struct AnvilEvmArgs {
+    /// Fetch state over a remote endpoint instead of starting from an empty state.
+    ///
+    /// If you want to fetch state from a specific block number, add a block number like `http://localhost:8545@1400000` or use the `--fork-block-number` argument.
+    #[arg(
+        long,
+        short,
+        visible_alias = "rpc-url",
+        value_name = "URL",
+        help_heading = "Fork config"
+    )]
+    pub fork_url: Option<ForkUrl>,
+
+    /// Fetch state from a specific block number over a remote endpoint.
+    ///
+    /// If negative, the given value is subtracted from the `latest` block number.
+    ///
+    /// See --fork-url.
+    #[arg(
+        long,
+        requires = "fork_url",
+        value_name = "BLOCK",
+        help_heading = "Fork config",
+        allow_hyphen_values = true
+    )]
+    pub fork_block_number: Option<i128>,
+
+    /// Timeout in ms for requests sent to remote JSON-RPC server in forking mode.
+    ///
+    /// Default value 45000
+    #[arg(id = "timeout", long = "timeout", help_heading = "Fork config", requires = "fork_url")]
+    pub fork_request_timeout: Option<u64>,
+
+    /// Number of retry requests for spurious networks (timed out requests)
+    ///
+    /// Default value 5
+    #[arg(id = "retries", long = "retries", help_heading = "Fork config", requires = "fork_url")]
+    pub fork_request_retries: Option<u32>,
+
     /// The block gas limit.
     #[arg(long, alias = "block-gas-limit", help_heading = "Environment config")]
     pub gas_limit: Option<u128>,
@@ -245,6 +293,46 @@ pub struct AnvilEvmArgs {
     pub memory_limit: Option<u64>,
 }
 
+/// Represents the input URL for a fork with an optional trailing block number:
+/// `http://localhost:8545@1000000`
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ForkUrl {
+    /// The endpoint url
+    pub url: String,
+    /// Optional trailing block
+    pub block: Option<u64>,
+}
+
+impl fmt::Display for ForkUrl {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.url.fmt(f)?;
+        if let Some(block) = self.block {
+            write!(f, "@{block}")?;
+        }
+        Ok(())
+    }
+}
+
+impl FromStr for ForkUrl {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if let Some((url, block)) = s.rsplit_once('@') {
+            if block == "latest" {
+                return Ok(Self { url: url.to_string(), block: None });
+            }
+            // this will prevent false positives for auths `user:[email protected]`
+            if !block.is_empty() && !block.contains(':') && !block.contains('.') {
+                let block: u64 = block
+                    .parse()
+                    .map_err(|_| format!("Failed to parse block number: `{block}`"))?;
+                return Ok(Self { url: url.to_string(), block: Some(block) });
+            }
+        }
+        Ok(Self { url: s.to_string(), block: None })
+    }
+}
+
 /// Clap's value parser for genesis. Loads a genesis.json file.
 fn read_genesis_file(path: &str) -> Result<Genesis, String> {
     foundry_common::fs::read_json_file(path.as_ref()).map_err(|err| err.to_string())

crates/anvil-polkadot/src/config.rs

@@ -8,7 +8,7 @@ use alloy_signer_local::{
 };
 use anvil_server::ServerConfig;
 use eyre::{Context, Result};
-use foundry_common::{duration_since_unix_epoch, sh_println};
+use foundry_common::{REQUEST_TIMEOUT, duration_since_unix_epoch, sh_println};
 use polkadot_sdk::{
     pallet_revive::evm::Account,
     sc_cli::{
@@ -326,6 +326,14 @@ pub struct AnvilNodeConfig {
     pub memory_limit: Option<u64>,
     /// Do not print log messages.
     pub silent: bool,
+    /// url of the rpc server that should be used for any rpc calls
+    pub eth_rpc_url: Option<String>,
+    /// pins the block number or transaction hash for the state fork
+    pub fork_choice: Option<ForkChoice>,
+    /// Timeout in for requests sent to remote JSON-RPC server in forking mode
+    pub fork_request_timeout: Duration,
+    /// Number of request retries for spurious networks
+    pub fork_request_retries: u32,
 }
 
 impl AnvilNodeConfig {
@@ -548,6 +556,10 @@ impl Default for AnvilNodeConfig {
             disable_default_create2_deployer: false,
             memory_limit: None,
             silent: false,
+            eth_rpc_url: None,
+            fork_choice: None,
+            fork_request_timeout: REQUEST_TIMEOUT,
+            fork_request_retries: 5,
         }
     }
 }
@@ -855,6 +867,69 @@ impl AnvilNodeConfig {
         self.silent = silent;
         self
     }
+
+    /// Sets the `eth_rpc_url` to use when forking
+    #[must_use]
+    pub fn with_eth_rpc_url<U: Into<String>>(mut self, eth_rpc_url: Option<U>) -> Self {
+        self.eth_rpc_url = eth_rpc_url.map(Into::into);
+        self
+    }
+
+    /// Sets the `fork_choice` to use to fork off from based on a block number
+    #[must_use]
+    pub fn with_fork_block_number<U: Into<u64>>(self, fork_block_number: Option<U>) -> Self {
+        self.with_fork_choice(fork_block_number.map(Into::into))
+    }
+
+    /// Sets the `fork_choice` to use to fork off from
+    #[must_use]
+    pub fn with_fork_choice<U: Into<ForkChoice>>(mut self, fork_choice: Option<U>) -> Self {
+        self.fork_choice = fork_choice.map(Into::into);
+        self
+    }
+
+    /// Sets the `fork_request_timeout` to use for requests
+    #[must_use]
+    pub fn fork_request_timeout(mut self, fork_request_timeout: Option<Duration>) -> Self {
+        if let Some(fork_request_timeout) = fork_request_timeout {
+            self.fork_request_timeout = fork_request_timeout;
+        }
+        self
+    }
+
+    /// Sets the `fork_request_retries` to use for spurious networks
+    #[must_use]
+    pub fn fork_request_retries(mut self, fork_request_retries: Option<u32>) -> Self {
+        if let Some(fork_request_retries) = fork_request_retries {
+            self.fork_request_retries = fork_request_retries;
+        }
+        self
+    }
+}
+
+/// Fork delimiter used to specify which block to fork from.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum ForkChoice {
+    /// Block number to fork from.
+    ///
+    /// If negative, the given value is subtracted from the `latest` block number.
+    Block(i128),
+}
+
+impl ForkChoice {
+    /// Returns the block number to fork from
+    pub fn block_number(&self) -> Option<i128> {
+        match self {
+            Self::Block(block_number) => Some(*block_number),
+        }
+    }
+}
+
+/// Convert a decimal block number into a ForkChoice
+impl From<u64> for ForkChoice {
+    fn from(block: u64) -> Self {
+        Self::Block(block as i128)
+    }
 }
 
 /// Can create dev accounts

crates/anvil-polkadot/src/substrate_node/genesis.rs

@@ -171,6 +171,7 @@ pub struct DevelopmentGenesisBlockBuilder<Block: BlockT, B, E> {
     commit_genesis_state: bool,
     backend: Arc<B>,
     executor: E,
+    checkpoint: Option<Block>,
     _phantom: PhantomData<Block>,
 }
 
@@ -217,6 +218,40 @@ impl<Block: BlockT, B: Backend<Block>, E: RuntimeVersionOf>
             commit_genesis_state,
             backend,
             executor,
+            checkpoint: None,
+            _phantom: PhantomData::<Block>,
+        })
+    }
+
+    pub fn new_with_checkpoint(
+        build_genesis_storage: &dyn BuildStorage,
+        commit_genesis_state: bool,
+        backend: Arc<B>,
+        executor: E,
+        checkpoint: Block,
+    ) -> sp_blockchain::Result<Self> {
+        let genesis_storage =
+            build_genesis_storage.build_storage().map_err(sp_blockchain::Error::Storage)?;
+
+        // Extract genesis number from checkpoint header
+        let genesis_number: u32 = (*checkpoint.header().number()).try_into().map_err(|_| {
+            sp_blockchain::Error::Application(
+                format!(
+                    "Checkpoint block number {} is too large for u32 (max: {})",
+                    checkpoint.header().number(),
+                    u32::MAX
+                )
+                .into(),
+            )
+        })?;
+
+        Ok(Self {
+            genesis_number,
+            genesis_storage,
+            commit_genesis_state,
+            backend,
+            executor,
+            checkpoint: Some(checkpoint),
             _phantom: PhantomData::<Block>,
         })
     }
@@ -234,30 +269,47 @@ impl<Block: BlockT, B: Backend<Block>, E: RuntimeVersionOf> BuildGenesisBlock<Bl
             commit_genesis_state,
             backend,
             executor,
+            checkpoint,
             _phantom,
         } = self;
 
-        let genesis_state_version =
-            resolve_state_version_from_wasm::<_, HashingFor<Block>>(&genesis_storage, &executor)?;
-        let mut op = backend.begin_operation()?;
-        let state_root =
-            op.set_genesis_state(genesis_storage, commit_genesis_state, genesis_state_version)?;
-        let extrinsics_root = <<<Block as BlockT>::Header as HeaderT>::Hashing as HashT>::trie_root(
-            Vec::new(),
-            genesis_state_version,
-        );
-        let genesis_block = Block::new(
-            <<Block as BlockT>::Header as HeaderT>::new(
-                genesis_number.into(),
-                extrinsics_root,
-                state_root,
-                Default::default(),
+        // If we have a checkpoint (fork mode), use it as the genesis block
+        if let Some(checkpoint) = checkpoint {
+            tracing::info!(
+                "Using checkpoint block as genesis: number={}, hash={:?}",
+                checkpoint.header().number(),
+                checkpoint.header().hash()
+            );
+
+            let op = backend.begin_operation()?;
+            Ok((checkpoint, op))
+        } else {
+            // Normal mode: create new genesis block
+            let genesis_state_version = resolve_state_version_from_wasm::<_, HashingFor<Block>>(
+                &genesis_storage,
+                &executor,
+            )?;
+            let mut op = backend.begin_operation()?;
+            let state_root =
+                op.set_genesis_state(genesis_storage, commit_genesis_state, genesis_state_version)?;
+            let extrinsics_root =
+                <<<Block as BlockT>::Header as HeaderT>::Hashing as HashT>::trie_root(
+                    Vec::new(),
+                    genesis_state_version,
+                );
+            let genesis_block = Block::new(
+                <<Block as BlockT>::Header as HeaderT>::new(
+                    genesis_number.into(),
+                    extrinsics_root,
+                    state_root,
+                    Default::default(),
+                    Default::default(),
+                ),
                 Default::default(),
-            ),
-            Default::default(),
-        );
+            );
 
-        Ok((genesis_block, op))
+            Ok((genesis_block, op))
+        }
     }
 }
 

crates/anvil-polkadot/src/substrate_node/lazy_loading/backend/forked_lazy_backend.rs

@@ -1,6 +1,5 @@
 use super::make_composite_child_key;
-use crate::substrate_node::lazy_loading::{LAZY_LOADING_LOG_TARGET, rpc_client::RPCClient};
-use alloy_primitives::hex;
+use crate::substrate_node::lazy_loading::rpc_client::RPCClient;
 use parking_lot::RwLock;
 use polkadot_sdk::{
     sc_client_api::StorageKey,
@@ -329,13 +328,6 @@ impl<Block: BlockT + DeserializeOwned> sp_state_machine::Backend<HashingFor<Bloc
         }
         .filter(|next_key| next_key != key);
 
-        tracing::trace!(
-            target: LAZY_LOADING_LOG_TARGET,
-            "next_storage_key: (key: {:?}, next_key: {:?})",
-            hex::encode(key),
-            maybe_next_key.clone().map(hex::encode)
-        );
-
         Ok(maybe_next_key)
     }
 
@@ -377,14 +369,6 @@ impl<Block: BlockT + DeserializeOwned> sp_state_machine::Backend<HashingFor<Bloc
         }
         .filter(|next_key| next_key != key);
 
-        tracing::trace!(
-            target: LAZY_LOADING_LOG_TARGET,
-            "next_child_storage_key: (child_info: {:?}, key: {:?}, next_key: {:?})",
-            child_info,
-            hex::encode(key),
-            maybe_next_key.clone().map(hex::encode)
-        );
-
         Ok(maybe_next_key)
     }