feat(vm): Allow switching between VMs for latest protocol version (#2508

)

## What ❔

- Allows using old (latest) VM by default, new VM only or old + new VM
in a shadow mode in the `MainBatchExecutor`.
- Allows configuring this mode for a new `VmRunner`-powered component,
VM playground.

## Why ❔

A separate component seems a relatively safe place to start integration from.

core/bin/external_node/src/main.rs

@@ -94,8 +94,8 @@ 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 = MainBatchExecutor::new(save_call_traces, true);
+    let batch_executor: Box<dyn BatchExecutor> = Box::new(batch_executor);
 
     let io = ExternalIO::new(
         connection_pool,
@@ -108,7 +108,7 @@ async fn build_state_keeper(
     Ok(ZkSyncStateKeeper::new(
         stop_receiver,
         Box::new(io),
-        batch_executor_base,
+        batch_executor,
         output_handler,
         Arc::new(NoopSealer),
         Arc::new(storage_factory),

core/bin/zksync_server/src/main.rs

@@ -12,10 +12,10 @@ use zksync_config::{
         fri_prover_group::FriProverGroupConfig,
         house_keeper::HouseKeeperConfig,
         BasicWitnessInputProducerConfig, ContractsConfig, DatabaseSecrets,
-        ExternalPriceApiClientConfig, FriProofCompressorConfig, FriProverConfig,
-        FriProverGatewayConfig, FriWitnessGeneratorConfig, FriWitnessVectorGeneratorConfig,
-        L1Secrets, ObservabilityConfig, PrometheusConfig, ProofDataHandlerConfig,
-        ProtectiveReadsWriterConfig, Secrets,
+        ExperimentalVmPlaygroundConfig, ExternalPriceApiClientConfig, FriProofCompressorConfig,
+        FriProverConfig, FriProverGatewayConfig, FriWitnessGeneratorConfig,
+        FriWitnessVectorGeneratorConfig, L1Secrets, ObservabilityConfig, PrometheusConfig,
+        ProofDataHandlerConfig, ProtectiveReadsWriterConfig, Secrets,
     },
     ApiConfig, BaseTokenAdjusterConfig, ContractVerifierConfig, DADispatcherConfig, DBConfig,
     EthConfig, EthWatchConfig, GasAdjusterConfig, GenesisConfig, ObjectStoreConfig, PostgresConfig,
@@ -230,6 +230,7 @@ fn load_env_config() -> anyhow::Result<TempConfigStore> {
         da_dispatcher_config: DADispatcherConfig::from_env().ok(),
         protective_reads_writer_config: ProtectiveReadsWriterConfig::from_env().ok(),
         basic_witness_input_producer_config: BasicWitnessInputProducerConfig::from_env().ok(),
+        vm_playground_config: ExperimentalVmPlaygroundConfig::from_env().ok(),
         core_object_store: ObjectStoreConfig::from_env().ok(),
         base_token_adjuster_config: BaseTokenAdjusterConfig::from_env().ok(),
         commitment_generator: None,

core/bin/zksync_server/src/node_builder.rs

@@ -53,7 +53,8 @@ use zksync_node_framework::{
         },
         tee_verifier_input_producer::TeeVerifierInputProducerLayer,
         vm_runner::{
-            bwip::BasicWitnessInputProducerLayer, protective_reads::ProtectiveReadsWriterLayer,
+            bwip::BasicWitnessInputProducerLayer, playground::VmPlaygroundLayer,
+            protective_reads::ProtectiveReadsWriterLayer,
         },
         web3_api::{
             caches::MempoolCacheLayer,
@@ -550,6 +551,16 @@ impl MainNodeBuilder {
         Ok(self)
     }
 
+    fn add_vm_playground_layer(mut self) -> anyhow::Result<Self> {
+        let vm_playground_config = try_load_config!(self.configs.vm_playground_config);
+        self.node.add_layer(VmPlaygroundLayer::new(
+            vm_playground_config,
+            self.genesis_config.l2_chain_id,
+        ));
+
+        Ok(self)
+    }
+
     fn add_base_token_ratio_persister_layer(mut self) -> anyhow::Result<Self> {
         let config = try_load_config!(self.configs.base_token_adjuster);
         let contracts_config = self.contracts_config.clone();
@@ -700,6 +711,9 @@ impl MainNodeBuilder {
                 Component::VmRunnerBwip => {
                     self = self.add_vm_runner_bwip_layer()?;
                 }
+                Component::VmPlayground => {
+                    self = self.add_vm_playground_layer()?;
+                }
             }
         }
         Ok(self.node.build()?)

core/lib/basic_types/src/lib.rs

@@ -28,7 +28,7 @@ pub mod protocol_version;
 pub mod prover_dal;
 pub mod tee_types;
 pub mod url;
-pub mod vm_version;
+pub mod vm;
 pub mod web3;
 
 /// Account place in the global state tree is uniquely identified by its address.

core/lib/basic_types/src/protocol_version.rs

@@ -12,7 +12,7 @@ use serde_with::{DeserializeFromStr, SerializeDisplay};
 
 use crate::{
     ethabi::Token,
-    vm_version::VmVersion,
+    vm::VmVersion,
     web3::contract::{Detokenize, Error},
     H256, U256,
 };

core/lib/basic_types/src/vm.rs

@@ -0,0 +1,39 @@
+//! Basic VM types that shared widely enough to not put them in the `multivm` crate.
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Copy)]
+pub enum VmVersion {
+    M5WithoutRefunds,
+    M5WithRefunds,
+    M6Initial,
+    M6BugWithCompressionFixed,
+    Vm1_3_2,
+    VmVirtualBlocks,
+    VmVirtualBlocksRefundsEnhancement,
+    VmBoojumIntegration,
+    Vm1_4_1,
+    Vm1_4_2,
+    Vm1_5_0SmallBootloaderMemory,
+    Vm1_5_0IncreasedBootloaderMemory,
+}
+
+impl VmVersion {
+    /// Returns the latest supported VM version.
+    pub const fn latest() -> VmVersion {
+        Self::Vm1_5_0IncreasedBootloaderMemory
+    }
+}
+
+/// Mode in which to run the new fast VM implementation.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum FastVmMode {
+    /// Run only the old VM.
+    #[default]
+    Old,
+    /// Run only the new Vm.
+    New,
+    /// Run both the new and old VM and compare their outputs for each transaction execution.
+    Shadow,
+}

core/lib/config/src/configs/experimental.rs

@@ -3,6 +3,7 @@
 use std::num::NonZeroU32;
 
 use serde::Deserialize;
+use zksync_basic_types::{vm::FastVmMode, L1BatchNumber};
 
 #[derive(Debug, Clone, PartialEq, Deserialize)]
 pub struct ExperimentalDBConfig {
@@ -60,3 +61,37 @@ impl ExperimentalDBConfig {
         100
     }
 }
+
+#[derive(Debug, Deserialize, Clone, PartialEq)]
+pub struct ExperimentalVmPlaygroundConfig {
+    /// Mode in which to run the fast VM implementation. Note that for it to actually be used, L1 batches should have a recent version.
+    #[serde(default)]
+    pub fast_vm_mode: FastVmMode,
+    /// Path to the RocksDB cache directory.
+    #[serde(default = "ExperimentalVmPlaygroundConfig::default_db_path")]
+    pub db_path: String,
+    /// First L1 batch to consider processed. Will not be used if the processing cursor is persisted, unless the `reset` flag is set.
+    #[serde(default)]
+    pub first_processed_batch: L1BatchNumber,
+    /// If set to true, processing cursor will reset `first_processed_batch` regardless of the current progress. Beware that this will likely
+    /// require to drop the RocksDB cache.
+    #[serde(default)]
+    pub reset: bool,
+}
+
+impl Default for ExperimentalVmPlaygroundConfig {
+    fn default() -> Self {
+        Self {
+            fast_vm_mode: FastVmMode::default(),
+            db_path: Self::default_db_path(),
+            first_processed_batch: L1BatchNumber(0),
+            reset: false,
+        }
+    }
+}
+
+impl ExperimentalVmPlaygroundConfig {
+    pub fn default_db_path() -> String {
+        "./db/vm_playground".to_owned()
+    }
+}

core/lib/config/src/configs/general.rs

@@ -9,10 +9,10 @@ use crate::{
         pruning::PruningConfig,
         snapshot_recovery::SnapshotRecoveryConfig,
         vm_runner::{BasicWitnessInputProducerConfig, ProtectiveReadsWriterConfig},
-        CommitmentGeneratorConfig, ExternalPriceApiClientConfig, FriProofCompressorConfig,
-        FriProverConfig, FriProverGatewayConfig, FriWitnessGeneratorConfig,
-        FriWitnessVectorGeneratorConfig, ObservabilityConfig, PrometheusConfig,
-        ProofDataHandlerConfig,
+        CommitmentGeneratorConfig, ExperimentalVmPlaygroundConfig, ExternalPriceApiClientConfig,
+        FriProofCompressorConfig, FriProverConfig, FriProverGatewayConfig,
+        FriWitnessGeneratorConfig, FriWitnessVectorGeneratorConfig, ObservabilityConfig,
+        PrometheusConfig, ProofDataHandlerConfig,
     },
     ApiConfig, ContractVerifierConfig, DBConfig, EthConfig, ObjectStoreConfig, PostgresConfig,
     SnapshotsCreatorConfig,
@@ -43,6 +43,7 @@ pub struct GeneralConfig {
     pub da_dispatcher_config: Option<DADispatcherConfig>,
     pub protective_reads_writer_config: Option<ProtectiveReadsWriterConfig>,
     pub basic_witness_input_producer_config: Option<BasicWitnessInputProducerConfig>,
+    pub vm_playground_config: Option<ExperimentalVmPlaygroundConfig>,
     pub commitment_generator: Option<CommitmentGeneratorConfig>,
     pub snapshot_recovery: Option<SnapshotRecoveryConfig>,
     pub pruning: Option<PruningConfig>,

core/lib/config/src/configs/mod.rs

@@ -9,7 +9,7 @@ pub use self::{
     database::{DBConfig, PostgresConfig},
     eth_sender::{EthConfig, GasAdjusterConfig},
     eth_watch::EthWatchConfig,
-    experimental::ExperimentalDBConfig,
+    experimental::{ExperimentalDBConfig, ExperimentalVmPlaygroundConfig},
     external_price_api_client::ExternalPriceApiClientConfig,
     fri_proof_compressor::FriProofCompressorConfig,
     fri_prover::FriProverConfig,

core/lib/config/src/configs/vm_runner.rs

@@ -1,7 +1,7 @@
 use serde::Deserialize;
 use zksync_basic_types::L1BatchNumber;
 
-#[derive(Debug, Deserialize, Clone, PartialEq, Default)]
+#[derive(Debug, Deserialize, Clone, PartialEq)]
 pub struct ProtectiveReadsWriterConfig {
     /// Path to the RocksDB data directory that serves state cache.
     #[serde(default = "ProtectiveReadsWriterConfig::default_db_path")]
@@ -18,7 +18,7 @@ impl ProtectiveReadsWriterConfig {
     }
 }
 
-#[derive(Debug, Deserialize, Clone, PartialEq, Default)]
+#[derive(Debug, Deserialize, Clone, PartialEq)]
 pub struct BasicWitnessInputProducerConfig {
     /// Path to the RocksDB data directory that serves state cache.
     #[serde(default = "BasicWitnessInputProducerConfig::default_db_path")]

core/lib/config/src/testonly.rs

@@ -6,6 +6,7 @@ use zksync_basic_types::{
     commitment::L1BatchCommitmentMode,
     network::Network,
     protocol_version::{ProtocolSemanticVersion, ProtocolVersionId, VersionPatch},
+    vm::FastVmMode,
     L1BatchNumber, L1ChainId, L2ChainId,
 };
 use zksync_consensus_utils::EncodeDist;
@@ -293,6 +294,21 @@ impl Distribution<configs::ExperimentalDBConfig> for EncodeDist {
     }
 }
 
+impl Distribution<configs::ExperimentalVmPlaygroundConfig> for EncodeDist {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> configs::ExperimentalVmPlaygroundConfig {
+        configs::ExperimentalVmPlaygroundConfig {
+            fast_vm_mode: match rng.gen_range(0..3) {
+                0 => FastVmMode::Old,
+                1 => FastVmMode::New,
+                _ => FastVmMode::Shadow,
+            },
+            db_path: self.sample(rng),
+            first_processed_batch: L1BatchNumber(rng.gen()),
+            reset: self.sample(rng),
+        }
+    }
+}
+
 impl Distribution<configs::database::DBConfig> for EncodeDist {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> configs::database::DBConfig {
         configs::database::DBConfig {
@@ -1038,6 +1054,7 @@ impl Distribution<configs::GeneralConfig> for EncodeDist {
             da_dispatcher_config: self.sample(rng),
             protective_reads_writer_config: self.sample(rng),
             basic_witness_input_producer_config: self.sample(rng),
+            vm_playground_config: self.sample(rng),
             commitment_generator: self.sample(rng),
             snapshot_recovery: self.sample(rng),
             pruning: self.sample(rng),

core/lib/env_config/src/vm_runner.rs

@@ -1,4 +1,6 @@
-use zksync_config::configs::{BasicWitnessInputProducerConfig, ProtectiveReadsWriterConfig};
+use zksync_config::configs::{
+    BasicWitnessInputProducerConfig, ExperimentalVmPlaygroundConfig, ProtectiveReadsWriterConfig,
+};
 
 use crate::{envy_load, FromEnv};
 
@@ -13,3 +15,69 @@ impl FromEnv for BasicWitnessInputProducerConfig {
         envy_load("vm_runner.bwip", "VM_RUNNER_BWIP_")
     }
 }
+
+impl FromEnv for ExperimentalVmPlaygroundConfig {
+    fn from_env() -> anyhow::Result<Self> {
+        envy_load("vm_runner.playground", "VM_RUNNER_PLAYGROUND_")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use zksync_basic_types::{vm::FastVmMode, L1BatchNumber};
+
+    use super::*;
+    use crate::test_utils::EnvMutex;
+
+    static MUTEX: EnvMutex = EnvMutex::new();
+
+    #[test]
+    fn bwip_config_from_env() {
+        let mut lock = MUTEX.lock();
+        let config = r#"
+            VM_RUNNER_BWIP_DB_PATH=/db/bwip
+            VM_RUNNER_BWIP_WINDOW_SIZE=50
+            VM_RUNNER_BWIP_FIRST_PROCESSED_BATCH=123
+        "#;
+        lock.set_env(config);
+
+        let config = BasicWitnessInputProducerConfig::from_env().unwrap();
+        assert_eq!(config.db_path, "/db/bwip");
+        assert_eq!(config.window_size, 50);
+        assert_eq!(config.first_processed_batch, L1BatchNumber(123));
+    }
+
+    #[test]
+    fn playground_config_from_env() {
+        let mut lock = MUTEX.lock();
+        let config = r#"
+            VM_RUNNER_PLAYGROUND_FAST_VM_MODE=shadow
+            VM_RUNNER_PLAYGROUND_DB_PATH=/db/vm_playground
+            VM_RUNNER_PLAYGROUND_FIRST_PROCESSED_BATCH=123
+            VM_RUNNER_PLAYGROUND_RESET=true
+        "#;
+        lock.set_env(config);
+
+        let config = ExperimentalVmPlaygroundConfig::from_env().unwrap();
+        assert_eq!(config.fast_vm_mode, FastVmMode::Shadow);
+        assert_eq!(config.db_path, "/db/vm_playground");
+        assert_eq!(config.first_processed_batch, L1BatchNumber(123));
+        assert!(config.reset);
+
+        lock.remove_env(&["VM_RUNNER_PLAYGROUND_RESET"]);
+        let config = ExperimentalVmPlaygroundConfig::from_env().unwrap();
+        assert!(!config.reset);
+
+        lock.remove_env(&["VM_RUNNER_PLAYGROUND_FIRST_PROCESSED_BATCH"]);
+        let config = ExperimentalVmPlaygroundConfig::from_env().unwrap();
+        assert_eq!(config.first_processed_batch, L1BatchNumber(0));
+
+        lock.remove_env(&["VM_RUNNER_PLAYGROUND_FAST_VM_MODE"]);
+        let config = ExperimentalVmPlaygroundConfig::from_env().unwrap();
+        assert_eq!(config.fast_vm_mode, FastVmMode::Old);
+
+        lock.remove_env(&["VM_RUNNER_PLAYGROUND_DB_PATH"]);
+        let config = ExperimentalVmPlaygroundConfig::from_env().unwrap();
+        assert!(!config.db_path.is_empty());
+    }
+}