feat(anvil): support max mem history value for --prune-history (#4263)

anvil/src/cmd.rs

@@ -176,8 +176,11 @@ pub struct NodeArgs {
     )]
     pub ipc: Option<Option<String>>,
 
-    #[clap(long, help = "Don't keep full chain history.")]
-    pub prune_history: bool,
+    #[clap(
+        long,
+        help = "Don't keep full chain history. If a number argument is specified, at most this number of states is kept in memory."
+    )]
+    pub prune_history: Option<Option<usize>>,
 
     #[clap(long, help = "Number of blocks with transactions to keep in memory.")]
     pub transaction_block_keeper: Option<usize>,
@@ -642,4 +645,13 @@ mod tests {
         let args: NodeArgs = NodeArgs::parse_from(["anvil", "--hardfork", "berlin"]);
         assert_eq!(args.hardfork, Some(Hardfork::Berlin));
     }
+
+    #[test]
+    fn can_parse_prune_config() {
+        let args: NodeArgs = NodeArgs::parse_from(["anvil", "--prune-history"]);
+        assert!(args.prune_history.is_some());
+
+        let args: NodeArgs = NodeArgs::parse_from(["anvil", "--prune-history", "100"]);
+        assert_eq!(args.prune_history, Some(Some(100)));
+    }
 }

anvil/src/config.rs

@@ -144,8 +144,10 @@ pub struct NodeConfig {
     pub enable_steps_tracing: bool,
     /// Configure the code size limit
     pub code_size_limit: Option<usize>,
-    /// If set to true, remove historic state entirely
-    pub prune_history: bool,
+    /// Configures how to remove historic state.
+    ///
+    /// If set to `Some(num)` keep latest num state in memory only.
+    pub prune_history: PruneStateHistoryConfig,
     /// The file where to load the state from
     pub init_state: Option<SerializableState>,
     /// max number of blocks with transactions in memory
@@ -370,7 +372,7 @@ impl Default for NodeConfig {
             compute_units_per_second: ALCHEMY_FREE_TIER_CUPS,
             ipc_path: None,
             code_size_limit: None,
-            prune_history: false,
+            prune_history: Default::default(),
             init_state: None,
             transaction_block_keeper: None,
         }
@@ -453,8 +455,8 @@ impl NodeConfig {
 
     /// Sets prune history status.
     #[must_use]
-    pub fn set_pruned_history(mut self, prune_history: bool) -> Self {
-        self.prune_history = prune_history;
+    pub fn set_pruned_history(mut self, prune_history: Option<Option<usize>>) -> Self {
+        self.prune_history = PruneStateHistoryConfig::from_args(prune_history);
         self
     }
 
@@ -967,6 +969,30 @@ impl NodeConfig {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub struct PruneStateHistoryConfig {
+    pub enabled: bool,
+    pub max_memory_history: Option<usize>,
+}
+
+// === impl PruneStateHistoryConfig ===
+
+impl PruneStateHistoryConfig {
+    /// Returns `true` if writing state history is supported
+    pub fn is_state_history_supported(&self) -> bool {
+        !self.enabled || self.max_memory_history.is_some()
+    }
+
+    /// Returns tru if this setting was enabled.
+    pub fn is_config_enabled(&self) -> bool {
+        self.enabled
+    }
+
+    pub fn from_args(val: Option<Option<usize>>) -> Self {
+        val.map(|max_memory_history| Self { enabled: true, max_memory_history }).unwrap_or_default()
+    }
+}
+
 /// Can create dev accounts
 #[derive(Debug, Clone)]
 pub struct AccountGenerator {
@@ -1069,3 +1095,18 @@ async fn find_latest_fork_block<M: Middleware>(provider: M) -> Result<u64, M::Er
 
     Ok(num)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_prune_history() {
+        let config = PruneStateHistoryConfig::default();
+        assert!(config.is_state_history_supported());
+        let config = PruneStateHistoryConfig::from_args(Some(None));
+        assert!(!config.is_state_history_supported());
+        let config = PruneStateHistoryConfig::from_args(Some(Some(10)));
+        assert!(config.is_state_history_supported());
+    }
+}

anvil/src/eth/backend/mem/mod.rs

@@ -1,5 +1,6 @@
 //! In memory blockchain backend
 use crate::{
+    config::PruneStateHistoryConfig,
     eth::{
         backend::{
             cheats::CheatsManager,
@@ -143,8 +144,8 @@ pub struct Backend {
     /// keeps track of active snapshots at a specific block
     active_snapshots: Arc<Mutex<HashMap<U256, (u64, H256)>>>,
     enable_steps_tracing: bool,
-    /// Whether to keep history state
-    prune_history: bool,
+    /// How to keep history state
+    prune_state_history_config: PruneStateHistoryConfig,
     /// max number of blocks with transactions in memory
     transaction_block_keeper: Option<usize>,
 }
@@ -159,7 +160,7 @@ impl Backend {
         fees: FeeManager,
         fork: Option<ClientFork>,
         enable_steps_tracing: bool,
-        prune_history: bool,
+        prune_state_history_config: PruneStateHistoryConfig,
         transaction_block_keeper: Option<usize>,
         automine_block_time: Option<Duration>,
     ) -> Self {
@@ -177,10 +178,22 @@ impl Backend {
 
         let start_timestamp =
             if let Some(fork) = fork.as_ref() { fork.timestamp() } else { genesis.timestamp };
+
+        let states = if prune_state_history_config.is_config_enabled() {
+            // if prune state history is enabled, configure the state cache only for memory
+            prune_state_history_config
+                .max_memory_history
+                .map(InMemoryBlockStates::new)
+                .unwrap_or_default()
+                .memory_only()
+        } else {
+            Default::default()
+        };
+
         let backend = Self {
             db,
             blockchain,
-            states: Arc::new(RwLock::new(Default::default())),
+            states: Arc::new(RwLock::new(states)),
             env,
             fork,
             time: TimeManager::new(start_timestamp),
@@ -190,7 +203,7 @@ impl Backend {
             genesis,
             active_snapshots: Arc::new(Mutex::new(Default::default())),
             enable_steps_tracing,
-            prune_history,
+            prune_state_history_config,
             transaction_block_keeper,
         };
 
@@ -669,7 +682,7 @@ impl Backend {
 
             let best_hash = self.blockchain.storage.read().best_hash;
 
-            if !self.prune_history {
+            if self.prune_state_history_config.is_state_history_supported() {
                 let db = self.db.read().await.current_state();
                 // store current state before executing all transactions
                 self.states.write().insert(best_hash, db);

anvil/src/eth/backend/mem/storage.rs

@@ -72,6 +72,12 @@ impl InMemoryBlockStates {
         }
     }
 
+    /// Configures no disk caching
+    pub fn memory_only(mut self) -> Self {
+        self.max_on_disk_limit = 0;
+        self
+    }
+
     /// This modifies the `limit` what to keep stored in memory.
     ///
     /// This will ensure the new limit adjusts based on the block time.
@@ -87,6 +93,11 @@ impl InMemoryBlockStates {
         }
     }
 
+    /// Returns true if only memory caching is supported.
+    fn is_memory_only(&self) -> bool {
+        self.max_on_disk_limit == 0
+    }
+
     /// Inserts a new (hash -> state) pair
     ///
     /// When the configured limit for the number of states that can be stored in memory is reached,
@@ -98,7 +109,7 @@ impl InMemoryBlockStates {
     ///
     /// When a state that was previously written to disk is requested, it is simply read from disk.
     pub fn insert(&mut self, hash: H256, state: StateDb) {
-        if self.present.len() >= self.in_memory_limit {
+        if !self.is_memory_only() && self.present.len() >= self.in_memory_limit {
             // once we hit the max limit we gradually decrease it
             self.in_memory_limit =
                 self.in_memory_limit.saturating_sub(1).max(self.min_in_memory_limit);
@@ -120,10 +131,13 @@ impl InMemoryBlockStates {
                 .pop_front()
                 .and_then(|hash| self.states.remove(&hash).map(|state| (hash, state)))
             {
-                let snapshot = state.0.clear_into_snapshot();
-                self.disk_cache.write(hash, snapshot);
-                self.on_disk_states.insert(hash, state);
-                self.oldest_on_disk.push_back(hash);
+                // only write to disk if supported
+                if !self.is_memory_only() {
+                    let snapshot = state.0.clear_into_snapshot();
+                    self.disk_cache.write(hash, snapshot);
+                    self.on_disk_states.insert(hash, state);
+                    self.oldest_on_disk.push_back(hash);
+                }
             }
         }