perf(vm): Improve snapshot management in batch executor

core/bin/external_node/src/main.rs

@@ -41,7 +41,7 @@ use zksync_reorg_detector::ReorgDetector;
 use zksync_shared_metrics::rustc::RUST_METRICS;
 use zksync_state::{PostgresStorageCaches, RocksdbStorageOptions};
 use zksync_state_keeper::{
-    seal_criteria::NoopSealer, AsyncRocksdbCache, BatchExecutor, MainBatchExecutor, OutputHandler,
+    seal_criteria::NoopSealer, AsyncRocksdbCache, BatchExecutor, OutputHandler,
     StateKeeperPersistence, TreeWritesPersistence, ZkSyncStateKeeper,
 };
 use zksync_storage::RocksDB;
@@ -94,8 +94,7 @@ async fn build_state_keeper(
         stop_receiver_clone.changed().await?;
         result
     }));
-    let batch_executor_base: Box<dyn BatchExecutor> =
-        Box::new(MainBatchExecutor::new(save_call_traces, true));
+    let batch_executor_base = BatchExecutor::new(save_call_traces, true);
 
     let io = ExternalIO::new(
         connection_pool,

core/bin/external_node/src/node_builder.rs

@@ -37,7 +37,7 @@ use zksync_node_framework::{
         reorg_detector::ReorgDetectorLayer,
         sigint::SigintHandlerLayer,
         state_keeper::{
-            external_io::ExternalIOLayer, main_batch_executor::MainBatchExecutorLayer,
+            batch_executor::BatchExecutorLayer, external_io::ExternalIOLayer,
             output_handler::OutputHandlerLayer, StateKeeperLayer,
         },
         sync_state_updater::SyncStateUpdaterLayer,
@@ -201,7 +201,7 @@ impl ExternalNodeBuilder {
             .api_namespaces()
             .contains(&Namespace::Debug);
         let main_node_batch_executor_builder_layer =
-            MainBatchExecutorLayer::new(save_call_traces, OPTIONAL_BYTECODE_COMPRESSION);
+            BatchExecutorLayer::new(save_call_traces, OPTIONAL_BYTECODE_COMPRESSION);
 
         let rocksdb_options = RocksdbStorageOptions {
             block_cache_capacity: self

core/bin/zksync_server/src/node_builder.rs

@@ -48,7 +48,7 @@ use zksync_node_framework::{
         query_eth_client::QueryEthClientLayer,
         sigint::SigintHandlerLayer,
         state_keeper::{
-            main_batch_executor::MainBatchExecutorLayer, mempool_io::MempoolIOLayer,
+            batch_executor::BatchExecutorLayer, mempool_io::MempoolIOLayer,
             output_handler::OutputHandlerLayer, RocksdbStorageOptions, StateKeeperLayer,
         },
         tee_verifier_input_producer::TeeVerifierInputProducerLayer,
@@ -232,7 +232,7 @@ impl MainNodeBuilder {
         );
         let db_config = try_load_config!(self.configs.db_config);
         let main_node_batch_executor_builder_layer =
-            MainBatchExecutorLayer::new(sk_config.save_call_traces, OPTIONAL_BYTECODE_COMPRESSION);
+            BatchExecutorLayer::new(sk_config.save_call_traces, OPTIONAL_BYTECODE_COMPRESSION);
 
         let rocksdb_options = RocksdbStorageOptions {
             block_cache_capacity: db_config

core/lib/multivm/src/interface/traits/vm.rs

@@ -147,13 +147,25 @@ pub trait VmFactory<S>: VmInterface {
 }
 
 /// Methods of VM requiring history manipulations.
+///
+/// # Snapshot workflow
+///
+/// External callers must follow the following snapshot workflow:
+///
+/// - Each new snapshot created using `make_snapshot()` must be either popped or rolled back before creating the following snapshot.
+///   OTOH, it's not required to call either of these methods by the end of VM execution.
+/// - `pop_snapshot_no_rollback()` may be called spuriously, when no snapshot was created. It is a no-op in this case.
+///
+/// These rules guarantee that at each given moment, a VM instance has <=1 snapshot (unless the VM makes snapshots internally),
+/// which may allow additional VM optimizations.
 pub trait VmInterfaceHistoryEnabled: VmInterface {
     /// Create a snapshot of the current VM state and push it into memory.
     fn make_snapshot(&mut self);
 
     /// Roll back VM state to the latest snapshot and destroy the snapshot.
     fn rollback_to_the_latest_snapshot(&mut self);
 
-    /// Pop the latest snapshot from memory and destroy it.
+    /// Pop the latest snapshot from memory and destroy it. If there are no snapshots, this should be a no-op
+    /// (i.e., the VM must not panic in this case).
     fn pop_snapshot_no_rollback(&mut self);
 }

core/lib/multivm/src/versions/shadow.rs

@@ -23,6 +23,17 @@ pub struct ShadowVm<S, T> {
     shadow: vm_fast::Vm<ImmutableStorageView<S>>,
 }
 
+impl<S, T> ShadowVm<S, T>
+where
+    S: ReadStorage,
+    T: VmFactory<StorageView<S>>,
+{
+    #[doc(hidden)] // used in batch executor tests
+    pub fn snapshot_count(&self) -> usize {
+        self.shadow.snapshot_count()
+    }
+}
+
 impl<S, T> VmFactory<StorageView<S>> for ShadowVm<S, T>
 where
     S: ReadStorage,

core/lib/multivm/src/versions/vm_1_3_2/vm.rs

@@ -297,6 +297,6 @@ impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, crate::vm_latest::Hist
     }
 
     fn pop_snapshot_no_rollback(&mut self) {
-        self.vm.pop_snapshot_no_rollback()
+        self.vm.pop_snapshot_no_rollback();
     }
 }

core/lib/multivm/src/versions/vm_1_3_2/vm_instance.rs

@@ -381,7 +381,7 @@ impl<H: HistoryMode, S: WriteStorage> VmInstance<S, H> {
     /// Removes the latest snapshot without rolling it back.
     /// This function expects that there is at least one snapshot present.
     pub fn pop_snapshot_no_rollback(&mut self) {
-        self.snapshots.pop().unwrap();
+        self.snapshots.pop();
     }
 
     /// Returns the amount of gas remaining to the VM.

core/lib/multivm/src/versions/vm_1_4_1/vm.rs

@@ -180,7 +180,6 @@ impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, crate::vm_latest::Hist
         self.make_snapshot_inner()
     }
 
-    /// Rollback vm state to the latest snapshot and destroy the snapshot
     fn rollback_to_the_latest_snapshot(&mut self) {
         let snapshot = self
             .snapshots
@@ -189,10 +188,7 @@ impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, crate::vm_latest::Hist
         self.rollback_to_snapshot(snapshot);
     }
 
-    /// Pop the latest snapshot from the memory and destroy it
     fn pop_snapshot_no_rollback(&mut self) {
-        self.snapshots
-            .pop()
-            .expect("Snapshot should be created before rolling it back");
+        self.snapshots.pop();
     }
 }

core/lib/multivm/src/versions/vm_1_4_2/vm.rs

@@ -178,14 +178,11 @@ impl<S: WriteStorage, H: HistoryMode> VmFactory<S> for Vm<S, H> {
     }
 }
 
-/// Methods of vm, which required some history manipulations
 impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, crate::vm_latest::HistoryEnabled> {
-    /// Create snapshot of current vm state and push it into the memory
     fn make_snapshot(&mut self) {
-        self.make_snapshot_inner()
+        self.make_snapshot_inner();
     }
 
-    /// Rollback vm state to the latest snapshot and destroy the snapshot
     fn rollback_to_the_latest_snapshot(&mut self) {
         let snapshot = self
             .snapshots
@@ -194,10 +191,7 @@ impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, crate::vm_latest::Hist
         self.rollback_to_snapshot(snapshot);
     }
 
-    /// Pop the latest snapshot from the memory and destroy it
     fn pop_snapshot_no_rollback(&mut self) {
-        self.snapshots
-            .pop()
-            .expect("Snapshot should be created before rolling it back");
+        self.snapshots.pop();
     }
 }

core/lib/multivm/src/versions/vm_boojum_integration/vm.rs

@@ -173,14 +173,11 @@ impl<S: WriteStorage, H: HistoryMode> VmFactory<S> for Vm<S, H> {
     }
 }
 
-/// Methods of vm, which required some history manipulations
 impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, crate::vm_latest::HistoryEnabled> {
-    /// Create snapshot of current vm state and push it into the memory
     fn make_snapshot(&mut self) {
-        self.make_snapshot_inner()
+        self.make_snapshot_inner();
     }
 
-    /// Rollback vm state to the latest snapshot and destroy the snapshot
     fn rollback_to_the_latest_snapshot(&mut self) {
         let snapshot = self
             .snapshots
@@ -189,10 +186,7 @@ impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, crate::vm_latest::Hist
         self.rollback_to_snapshot(snapshot);
     }
 
-    /// Pop the latest snapshot from the memory and destroy it
     fn pop_snapshot_no_rollback(&mut self) {
-        self.snapshots
-            .pop()
-            .expect("Snapshot should be created before rolling it back");
+        self.snapshots.pop();
     }
 }

core/lib/multivm/src/versions/vm_fast/vm.rs

@@ -63,16 +63,17 @@ pub struct Vm<S> {
     pub(crate) inner: VirtualMachine,
     suspended_at: u16,
     gas_for_account_validation: u32,
-
     pub(crate) bootloader_state: BootloaderState,
-
     pub(crate) batch_env: L1BatchEnv,
     pub(crate) system_env: SystemEnv,
-
     snapshots: Vec<VmSnapshot>,
 }
 
 impl<S: ReadStorage> Vm<S> {
+    pub(crate) fn snapshot_count(&self) -> usize {
+        self.snapshots.len()
+    }
+
     fn run(
         &mut self,
         execution_mode: VmExecutionMode,

core/lib/multivm/src/versions/vm_latest/vm.rs

@@ -229,14 +229,11 @@ impl<S: WriteStorage, H: HistoryMode> Vm<S, H> {
     }
 }
 
-/// Methods of vm, which required some history manipulations
 impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, HistoryEnabled> {
-    /// Create snapshot of current vm state and push it into the memory
     fn make_snapshot(&mut self) {
         self.make_snapshot_inner()
     }
 
-    /// Rollback vm state to the latest snapshot and destroy the snapshot
     fn rollback_to_the_latest_snapshot(&mut self) {
         let snapshot = self
             .snapshots
@@ -245,10 +242,7 @@ impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, HistoryEnabled> {
         self.rollback_to_snapshot(snapshot);
     }
 
-    /// Pop the latest snapshot from the memory and destroy it
     fn pop_snapshot_no_rollback(&mut self) {
-        self.snapshots
-            .pop()
-            .expect("Snapshot should be created before rolling it back");
+        self.snapshots.pop();
     }
 }

core/lib/multivm/src/versions/vm_m6/vm.rs

@@ -325,6 +325,6 @@ impl<S: Storage> VmInterfaceHistoryEnabled for Vm<S, crate::vm_latest::HistoryEn
     }
 
     fn pop_snapshot_no_rollback(&mut self) {
-        self.vm.pop_snapshot_no_rollback()
+        self.vm.pop_snapshot_no_rollback();
     }
 }

core/lib/multivm/src/versions/vm_m6/vm_instance.rs

@@ -396,9 +396,8 @@ impl<H: HistoryMode, S: Storage> VmInstance<S, H> {
     }
 
     /// Removes the latest snapshot without rolling back to it.
-    /// This function expects that there is at least one snapshot present.
     pub fn pop_snapshot_no_rollback(&mut self) {
-        self.snapshots.pop().unwrap();
+        self.snapshots.pop();
     }
 
     /// Returns the amount of gas remaining to the VM.

core/lib/multivm/src/versions/vm_refunds_enhancement/vm.rs

@@ -146,14 +146,12 @@ impl<S: WriteStorage, H: HistoryMode> VmFactory<S> for Vm<S, H> {
     }
 }
 
-/// Methods of vm, which required some history manipulations
 impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, HistoryEnabled> {
     /// Create snapshot of current vm state and push it into the memory
     fn make_snapshot(&mut self) {
-        self.make_snapshot_inner()
+        self.make_snapshot_inner();
     }
 
-    /// Rollback vm state to the latest snapshot and destroy the snapshot
     fn rollback_to_the_latest_snapshot(&mut self) {
         let snapshot = self
             .snapshots
@@ -162,10 +160,7 @@ impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, HistoryEnabled> {
         self.rollback_to_snapshot(snapshot);
     }
 
-    /// Pop the latest snapshot from the memory and destroy it
     fn pop_snapshot_no_rollback(&mut self) {
-        self.snapshots
-            .pop()
-            .expect("Snapshot should be created before rolling it back");
+        self.snapshots.pop();
     }
 }

core/lib/multivm/src/versions/vm_virtual_blocks/vm.rs

@@ -146,14 +146,11 @@ impl<S: WriteStorage, H: HistoryMode> VmFactory<S> for Vm<S, H> {
     }
 }
 
-/// Methods of vm, which required some history manipulations
 impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, HistoryEnabled> {
-    /// Create snapshot of current vm state and push it into the memory
     fn make_snapshot(&mut self) {
         self.make_snapshot_inner()
     }
 
-    /// Rollback vm state to the latest snapshot and destroy the snapshot
     fn rollback_to_the_latest_snapshot(&mut self) {
         let snapshot = self
             .snapshots
@@ -162,10 +159,7 @@ impl<S: WriteStorage> VmInterfaceHistoryEnabled for Vm<S, HistoryEnabled> {
         self.rollback_to_snapshot(snapshot);
     }
 
-    /// Pop the latest snapshot from the memory and destroy it
     fn pop_snapshot_no_rollback(&mut self) {
-        self.snapshots
-            .pop()
-            .expect("Snapshot should be created before rolling it back");
+        self.snapshots.pop();
     }
 }

core/lib/state/src/rocksdb/mod.rs

@@ -257,6 +257,12 @@ impl RocksdbStorageBuilder {
     ) -> anyhow::Result<()> {
         self.0.revert(storage, last_l1_batch_to_keep).await
     }
+
+    /// Returns the underlying storage without synchronizing it or checking that it's up-to-date.
+    /// Should not be used other than for tests.
+    pub fn build_unchecked(self) -> RocksdbStorage {
+        self.0
+    }
 }
 
 impl RocksdbStorage {

core/node/consensus/src/testonly.rs

@@ -34,11 +34,10 @@ use zksync_state_keeper::{
     io::{IoCursor, L1BatchParams, L2BlockParams},
     seal_criteria::NoopSealer,
     testonly::{
-        fund, l1_transaction, l2_transaction, test_batch_executor::MockReadStorageFactory,
-        MockBatchExecutor,
+        fund, l1_transaction, l2_transaction, test_batch_vm::MockReadStorageFactory, MockBatchVm,
     },
-    AsyncRocksdbCache, MainBatchExecutor, OutputHandler, StateKeeperPersistence,
-    TreeWritesPersistence, ZkSyncStateKeeper,
+    AsyncRocksdbCache, BatchExecutor, OutputHandler, StateKeeperPersistence, TreeWritesPersistence,
+    ZkSyncStateKeeper,
 };
 use zksync_test_account::Account;
 use zksync_types::{
@@ -547,7 +546,7 @@ impl StateKeeperRunner {
                     ZkSyncStateKeeper::new(
                         stop_recv,
                         Box::new(io),
-                        Box::new(MainBatchExecutor::new(false, false)),
+                        BatchExecutor::new(false, false),
                         OutputHandler::new(Box::new(persistence.with_tx_insertion()))
                             .with_handler(Box::new(self.sync_state.clone())),
                         Arc::new(NoopSealer),
@@ -628,12 +627,13 @@ impl StateKeeperRunner {
                     ZkSyncStateKeeper::new(
                         stop_recv,
                         Box::new(io),
-                        Box::new(MockBatchExecutor),
+                        BatchExecutor::new(false, false)
+                            .with_vm_factory(Arc::<MockBatchVm>::default()),
                         OutputHandler::new(Box::new(persistence.with_tx_insertion()))
                             .with_handler(Box::new(tree_writes_persistence))
                             .with_handler(Box::new(self.sync_state.clone())),
                         Arc::new(NoopSealer),
-                        Arc::new(MockReadStorageFactory),
+                        Arc::<MockReadStorageFactory>::default(),
                     )
                     .run()
                     .await

core/node/node_framework/examples/main_node.rs

@@ -42,7 +42,7 @@ use zksync_node_framework::{
         query_eth_client::QueryEthClientLayer,
         sigint::SigintHandlerLayer,
         state_keeper::{
-            main_batch_executor::MainBatchExecutorLayer, mempool_io::MempoolIOLayer,
+            batch_executor::BatchExecutorLayer, mempool_io::MempoolIOLayer,
             output_handler::OutputHandlerLayer, StateKeeperLayer,
         },
         web3_api::{
@@ -161,7 +161,7 @@ impl MainNodeBuilder {
             wallets.state_keeper.context("State keeper wallets")?,
         );
         let main_node_batch_executor_builder_layer =
-            MainBatchExecutorLayer::new(StateKeeperConfig::from_env()?.save_call_traces, true);
+            BatchExecutorLayer::new(StateKeeperConfig::from_env()?.save_call_traces, true);
         let db_config = DBConfig::from_env()?;
 
         let rocksdb_options = RocksdbStorageOptions {