[validator] Changing module cache for VM validator (#15714)

Changes module caching logic for VM validator (mempool validation running prologues) to make it
work correctly with module publishing. When new modules are published and committed, the module
cache will detect the change and reload the new module.

aptos-move/aptos-vm/src/aptos_vm.rs

@@ -119,7 +119,7 @@ use move_vm_metrics::{Timer, VM_TIMER};
 use move_vm_runtime::{
     logging::expect_no_verification_errors,
     module_traversal::{TraversalContext, TraversalStorage},
-    RuntimeEnvironment, WithRuntimeEnvironment,
+    ModuleStorage, RuntimeEnvironment, WithRuntimeEnvironment,
 };
 use move_vm_types::gas::{GasMeter, UnmeteredGasMeter};
 use num_cpus;
@@ -1824,7 +1824,7 @@ impl AptosVM {
         &self,
         session: &mut SessionExt,
         resolver: &impl AptosMoveResolver,
-        module_storage: &impl AptosModuleStorage,
+        module_storage: &impl ModuleStorage,
         transaction: &SignedTransaction,
         transaction_data: &TransactionMetadata,
         log_context: &AdapterLogSchema,
@@ -2553,7 +2553,7 @@ impl AptosVM {
         &self,
         session: &mut SessionExt,
         resolver: &impl AptosMoveResolver,
-        module_storage: &impl AptosModuleStorage,
+        module_storage: &impl ModuleStorage,
         payload: &TransactionPayload,
         txn_data: &TransactionMetadata,
         log_context: &AdapterLogSchema,
@@ -2899,7 +2899,7 @@ impl VMValidator for AptosVM {
         &self,
         transaction: SignedTransaction,
         state_view: &impl StateView,
-        module_storage: &impl AptosCodeStorage,
+        module_storage: &impl ModuleStorage,
     ) -> VMValidatorResult {
         let _timer = TXN_VALIDATION_SECONDS.start_timer();
         let log_context = AdapterLogSchema::new(state_view.id(), 0);
@@ -3057,7 +3057,7 @@ pub(crate) fn is_account_init_for_sponsored_transaction(
     txn_data: &TransactionMetadata,
     features: &Features,
     resolver: &impl AptosMoveResolver,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
 ) -> VMResult<bool> {
     if features.is_enabled(FeatureFlag::SPONSORED_AUTOMATIC_ACCOUNT_V1_CREATION)
         && txn_data.fee_payer.is_some()
@@ -3084,7 +3084,7 @@ pub(crate) fn is_account_init_for_sponsored_transaction(
 pub(crate) fn fetch_module_metadata_for_struct_tag(
     struct_tag: &StructTag,
     resolver: &impl AptosMoveResolver,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
 ) -> VMResult<Vec<Metadata>> {
     if module_storage.is_enabled() {
         module_storage.fetch_existing_module_metadata(&struct_tag.address, &struct_tag.module)

aptos-move/aptos-vm/src/gas.rs

@@ -12,11 +12,9 @@ use aptos_logger::{enabled, Level};
 use aptos_memory_usage_tracker::MemoryTrackedGasMeter;
 use aptos_types::on_chain_config::Features;
 use aptos_vm_logging::{log_schema::AdapterLogSchema, speculative_log, speculative_warn};
-use aptos_vm_types::{
-    module_and_script_storage::module_storage::AptosModuleStorage,
-    storage::{space_pricing::DiskSpacePricing, StorageGasParameters},
-};
+use aptos_vm_types::storage::{space_pricing::DiskSpacePricing, StorageGasParameters};
 use move_core_types::vm_status::{StatusCode, VMStatus};
+use move_vm_runtime::ModuleStorage;
 
 /// This is used until gas version 18, which introduces a configurable entry for this.
 const MAXIMUM_APPROVED_TRANSACTION_SIZE_LEGACY: u64 = 1024 * 1024;
@@ -47,7 +45,7 @@ pub(crate) fn check_gas(
     gas_params: &AptosGasParameters,
     gas_feature_version: u64,
     resolver: &impl AptosMoveResolver,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
     txn_metadata: &TransactionMetadata,
     features: &Features,
     is_approved_gov_script: bool,

aptos-move/aptos-vm/src/keyless_validation.rs

@@ -15,14 +15,14 @@ use aptos_types::{
     transaction::authenticator::{EphemeralPublicKey, EphemeralSignature},
     vm_status::{StatusCode, VMStatus},
 };
-use aptos_vm_types::module_and_script_storage::module_storage::AptosModuleStorage;
 use ark_bn254::Bn254;
 use ark_groth16::PreparedVerifyingKey;
 use move_binary_format::errors::Location;
 use move_core_types::{
     account_address::AccountAddress, language_storage::CORE_CODE_ADDRESS,
     move_resource::MoveStructType,
 };
+use move_vm_runtime::ModuleStorage;
 use serde::Deserialize;
 
 macro_rules! value_deserialization_error {
@@ -36,15 +36,15 @@ macro_rules! value_deserialization_error {
 
 fn get_resource_on_chain<T: MoveStructType + for<'a> Deserialize<'a>>(
     resolver: &impl AptosMoveResolver,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
 ) -> anyhow::Result<T, VMStatus> {
     get_resource_on_chain_at_addr(&CORE_CODE_ADDRESS, resolver, module_storage)
 }
 
 fn get_resource_on_chain_at_addr<T: MoveStructType + for<'a> Deserialize<'a>>(
     addr: &AccountAddress,
     resolver: &impl AptosMoveResolver,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
 ) -> anyhow::Result<T, VMStatus> {
     let metadata = fetch_module_metadata_for_struct_tag(&T::struct_tag(), resolver, module_storage)
         .map_err(|e| e.into_vm_status())?;
@@ -87,21 +87,21 @@ fn get_jwks_onchain(resolver: &impl AptosMoveResolver) -> anyhow::Result<Patched
 fn get_federated_jwks_onchain(
     resolver: &impl AptosMoveResolver,
     jwk_addr: &AccountAddress,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
 ) -> anyhow::Result<FederatedJWKs, VMStatus> {
     get_resource_on_chain_at_addr::<FederatedJWKs>(jwk_addr, resolver, module_storage)
 }
 
 pub(crate) fn get_groth16_vk_onchain(
     resolver: &impl AptosMoveResolver,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
 ) -> anyhow::Result<Groth16VerificationKey, VMStatus> {
     get_resource_on_chain::<Groth16VerificationKey>(resolver, module_storage)
 }
 
 fn get_configs_onchain(
     resolver: &impl AptosMoveResolver,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
 ) -> anyhow::Result<Configuration, VMStatus> {
     get_resource_on_chain::<Configuration>(resolver, module_storage)
 }
@@ -160,7 +160,7 @@ pub(crate) fn validate_authenticators(
     authenticators: &Vec<(AnyKeylessPublicKey, KeylessSignature)>,
     features: &Features,
     resolver: &impl AptosMoveResolver,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
 ) -> Result<(), VMStatus> {
     let mut with_zk = false;
     for (pk, sig) in authenticators {

aptos-move/aptos-vm/src/lib.rs

@@ -139,7 +139,7 @@ use aptos_types::{
     },
     vm_status::VMStatus,
 };
-use aptos_vm_types::module_and_script_storage::code_storage::AptosCodeStorage;
+use move_vm_runtime::ModuleStorage;
 use std::{marker::Sync, sync::Arc};
 pub use verifier::view_function::determine_is_view;
 
@@ -150,7 +150,7 @@ pub trait VMValidator {
         &self,
         transaction: SignedTransaction,
         state_view: &impl StateView,
-        module_storage: &impl AptosCodeStorage,
+        module_storage: &impl ModuleStorage,
     ) -> VMValidatorResult;
 }
 

aptos-move/aptos-vm/src/transaction_validation.rs

@@ -17,7 +17,6 @@ use aptos_types::{
     on_chain_config::Features, transaction::Multisig,
 };
 use aptos_vm_logging::log_schema::AdapterLogSchema;
-use aptos_vm_types::module_and_script_storage::module_storage::AptosModuleStorage;
 use fail::fail_point;
 use move_binary_format::errors::VMResult;
 use move_core_types::{
@@ -28,7 +27,9 @@ use move_core_types::{
     value::{serialize_values, MoveValue},
     vm_status::{AbortLocation, StatusCode, VMStatus},
 };
-use move_vm_runtime::{logging::expect_no_verification_errors, module_traversal::TraversalContext};
+use move_vm_runtime::{
+    logging::expect_no_verification_errors, module_traversal::TraversalContext, ModuleStorage,
+};
 use move_vm_types::gas::UnmeteredGasMeter;
 use once_cell::sync::Lazy;
 
@@ -85,7 +86,7 @@ impl TransactionValidation {
 
 pub(crate) fn run_script_prologue(
     session: &mut SessionExt,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
     txn_data: &TransactionMetadata,
     features: &Features,
     log_context: &AdapterLogSchema,
@@ -232,7 +233,7 @@ pub(crate) fn run_script_prologue(
 /// match that hash.
 pub(crate) fn run_multisig_prologue(
     session: &mut SessionExt,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
     txn_data: &TransactionMetadata,
     payload: &Multisig,
     features: &Features,
@@ -272,7 +273,7 @@ pub(crate) fn run_multisig_prologue(
 
 fn run_epilogue(
     session: &mut SessionExt,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
     gas_remaining: Gas,
     fee_statement: FeeStatement,
     txn_data: &TransactionMetadata,
@@ -377,7 +378,7 @@ fn run_epilogue(
 
 fn emit_fee_statement(
     session: &mut SessionExt,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
     fee_statement: FeeStatement,
     traversal_context: &mut TraversalContext,
 ) -> VMResult<()> {
@@ -398,7 +399,7 @@ fn emit_fee_statement(
 /// in the `ACCOUNT_MODULE` on chain.
 pub(crate) fn run_success_epilogue(
     session: &mut SessionExt,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
     gas_remaining: Gas,
     fee_statement: FeeStatement,
     features: &Features,
@@ -431,7 +432,7 @@ pub(crate) fn run_success_epilogue(
 /// stored in the `ACCOUNT_MODULE` on chain.
 pub(crate) fn run_failure_epilogue(
     session: &mut SessionExt,
-    module_storage: &impl AptosModuleStorage,
+    module_storage: &impl ModuleStorage,
     gas_remaining: Gas,
     fee_statement: FeeStatement,
     features: &Features,

aptos-move/block-executor/src/code_cache.rs

@@ -75,7 +75,7 @@ impl<'a, T: Transaction, S: TStateView<Key = T::Key>> ModuleCache for LatestView
         deserialized_code: Self::Deserialized,
         extension: Arc<Self::Extension>,
         version: Self::Version,
-    ) -> VMResult<()> {
+    ) -> VMResult<Arc<ModuleCode<Self::Deserialized, Self::Verified, Self::Extension>>> {
         self.as_module_cache().insert_deserialized_module(
             key,
             deserialized_code,

third_party/move/move-vm/types/src/code/cache/module_cache.rs

@@ -92,13 +92,14 @@ pub trait ModuleCache {
     ///   1. returns an error if the version of existing entry is higher,
     ///   2. does not perform the insertion if the version is the same,
     ///   3. inserts the new code if the new version is higher.
+    /// Returns the newly inserted (or existing) module at the specified key.
     fn insert_deserialized_module(
         &self,
         key: Self::Key,
         deserialized_code: Self::Deserialized,
         extension: Arc<Self::Extension>,
         version: Self::Version,
-    ) -> VMResult<()>;
+    ) -> VMResult<Arc<ModuleCode<Self::Deserialized, Self::Verified, Self::Extension>>>;
 
     /// Stores verified code at specified version to the module cache if there was no entry
     /// associated with this key before. If module cache already contains an entry, then:
@@ -231,6 +232,12 @@ where
             .into_iter()
             .map(|(k, m)| (k, m.into_module_code()))
     }
+
+    /// Returns the version of the module stored in cache. Used for tests only.
+    #[cfg(any(test, feature = "testing"))]
+    pub fn get_module_version(&self, key: &K) -> Option<V> {
+        self.module_cache.borrow().get(key).map(|m| m.version())
+    }
 }
 
 impl<K, DC, VC, E, V> ModuleCache for UnsyncModuleCache<K, DC, VC, E, V>
@@ -251,23 +258,29 @@ where
         deserialized_code: Self::Deserialized,
         extension: Arc<Self::Extension>,
         version: Self::Version,
-    ) -> VMResult<()> {
+    ) -> VMResult<Arc<ModuleCode<Self::Deserialized, Self::Verified, Self::Extension>>> {
         use hashbrown::hash_map::Entry::*;
 
         match self.module_cache.borrow_mut().entry(key) {
             Occupied(mut entry) => match version.cmp(&entry.get().version()) {
                 Ordering::Less => Err(version_too_small_error!()),
-                Ordering::Equal => Ok(()),
+                Ordering::Equal => Ok(entry.get().module_code().clone()),
                 Ordering::Greater => {
-                    let module = ModuleCode::from_deserialized(deserialized_code, extension);
-                    entry.insert(VersionedModuleCode::new(module, version));
-                    Ok(())
+                    let versioned_module = VersionedModuleCode::new(
+                        ModuleCode::from_deserialized(deserialized_code, extension),
+                        version,
+                    );
+                    let module = versioned_module.module_code().clone();
+                    entry.insert(versioned_module);
+                    Ok(module)
                 },
             },
             Vacant(entry) => {
                 let module = ModuleCode::from_deserialized(deserialized_code, extension);
-                entry.insert(VersionedModuleCode::new(module, version));
-                Ok(())
+                Ok(entry
+                    .insert(VersionedModuleCode::new(module, version))
+                    .module_code()
+                    .clone())
             },
         }
     }
@@ -399,23 +412,29 @@ where
         deserialized_code: Self::Deserialized,
         extension: Arc<Self::Extension>,
         version: Self::Version,
-    ) -> VMResult<()> {
+    ) -> VMResult<Arc<ModuleCode<Self::Deserialized, Self::Verified, Self::Extension>>> {
         use dashmap::mapref::entry::Entry::*;
 
         match self.module_cache.entry(key) {
             Occupied(mut entry) => match version.cmp(&entry.get().version()) {
                 Ordering::Less => Err(version_too_small_error!()),
-                Ordering::Equal => Ok(()),
+                Ordering::Equal => Ok(entry.get().module_code().clone()),
                 Ordering::Greater => {
-                    let module = ModuleCode::from_deserialized(deserialized_code, extension);
-                    entry.insert(CachePadded::new(VersionedModuleCode::new(module, version)));
-                    Ok(())
+                    let versioned_module = VersionedModuleCode::new(
+                        ModuleCode::from_deserialized(deserialized_code, extension),
+                        version,
+                    );
+                    let module = versioned_module.module_code().clone();
+                    entry.insert(CachePadded::new(versioned_module));
+                    Ok(module)
                 },
             },
             Vacant(entry) => {
                 let module = ModuleCode::from_deserialized(deserialized_code, extension);
-                entry.insert(CachePadded::new(VersionedModuleCode::new(module, version)));
-                Ok(())
+                Ok(entry
+                    .insert(CachePadded::new(VersionedModuleCode::new(module, version)))
+                    .module_code()
+                    .clone())
             },
         }
     }

third_party/move/move-vm/types/src/code/errors.rs

@@ -14,7 +14,7 @@ macro_rules! panic_error {
 
 #[macro_export]
 macro_rules! module_storage_error {
-    ($addr:ident, $name:ident, $err:ident) => {
+    ($addr:expr, $name:expr, $err:ident) => {
         move_binary_format::errors::PartialVMError::new(
             move_core_types::vm_status::StatusCode::STORAGE_ERROR,
         )

vm-validator/Cargo.toml

@@ -20,8 +20,11 @@ aptos-types = { workspace = true }
 aptos-vm = { workspace = true }
 aptos-vm-environment = { workspace = true }
 aptos-vm-logging = { workspace = true }
-aptos-vm-types = { workspace = true }
 fail = { workspace = true }
+move-binary-format = { workspace = true }
+move-core-types = { workspace = true }
+move-vm-runtime = { workspace = true }
+move-vm-types = { workspace = true }
 rand = { workspace = true }
 
 [dev-dependencies]
@@ -31,9 +34,9 @@ aptos-db = { workspace = true }
 aptos-executor-test-helpers = { workspace = true }
 aptos-gas-schedule = { workspace = true, features = ["testing"] } 
 aptos-temppath = { workspace = true }
-aptos-types = { workspace = true }
+aptos-types = { workspace = true, features = ["testing"] }
 aptos-vm-genesis = { workspace = true }
-move-core-types = { workspace = true }
+move-vm-types = { workspace = true, features = ["testing"] }
 rand = { workspace = true }
 
 [features]