[refactoring][fix] Add module cache for txn validation

  - Refactored module/code storages to own `E: WithRuntimeEnvironment`
  - Added a cache for mempool validation

api/test-context/src/test_context.rs

@@ -48,7 +48,7 @@ use aptos_types::{
     },
 };
 use aptos_vm::AptosVM;
-use aptos_vm_validator::vm_validator::VMValidator;
+use aptos_vm_validator::vm_validator::PooledVMValidator;
 use bytes::Bytes;
 use hyper::{HeaderMap, Response};
 use rand::SeedableRng;
@@ -172,7 +172,7 @@ pub fn new_test_context_inner(
         db_bootstrapper::maybe_bootstrap::<AptosVM>(&db_rw, &genesis, genesis_waypoint).unwrap();
     assert!(ret.is_some());
 
-    let mempool = MockSharedMempool::new_in_runtime(&db_rw, VMValidator::new(db.clone()));
+    let mempool = MockSharedMempool::new_in_runtime(&db_rw, PooledVMValidator::new(db.clone(), 1));
 
     node_config
         .storage

aptos-move/aptos-debugger/Cargo.toml

@@ -29,7 +29,6 @@ aptos-vm-types = { workspace = true }
 bcs = { workspace = true }
 clap = { workspace = true }
 itertools = { workspace = true }
-move-vm-runtime = { workspace = true }
 regex = { workspace = true }
 reqwest = { workspace = true }
 tokio = { workspace = true }

aptos-move/aptos-debugger/src/aptos_debugger.rs

@@ -30,7 +30,6 @@ use aptos_vm_environment::environment::AptosEnvironment;
 use aptos_vm_logging::log_schema::AdapterLogSchema;
 use aptos_vm_types::{module_and_script_storage::AsAptosCodeStorage, output::VMOutput};
 use itertools::Itertools;
-use move_vm_runtime::WithRuntimeEnvironment;
 use std::{path::Path, sync::Arc, time::Instant};
 
 pub struct AptosDebugger {
@@ -123,7 +122,7 @@ impl AptosDebugger {
         let env = AptosEnvironment::new(&state_view);
         let vm = AptosVM::new(env.clone(), &state_view);
         let resolver = state_view.as_move_resolver();
-        let code_storage = state_view.as_aptos_code_storage(env.runtime_environment());
+        let code_storage = state_view.as_aptos_code_storage(env);
 
         let (status, output, gas_profiler) = vm.execute_user_transaction_with_modified_gas_meter(
             &resolver,

aptos-move/aptos-release-builder/src/simulate.rs

@@ -70,10 +70,7 @@ use move_core_types::{
     language_storage::{ModuleId, StructTag},
     move_resource::MoveResource,
 };
-use move_vm_runtime::{
-    module_traversal::{TraversalContext, TraversalStorage},
-    WithRuntimeEnvironment,
-};
+use move_vm_runtime::module_traversal::{TraversalContext, TraversalStorage};
 use move_vm_types::{gas::UnmeteredGasMeter, resolver::ModuleResolver};
 use once_cell::sync::Lazy;
 use parking_lot::Mutex;
@@ -471,7 +468,7 @@ fn force_end_epoch(state_view: &SimulationStateView<impl StateView>) -> Result<(
     let env = AptosEnvironment::new_with_injected_create_signer_for_gov_sim(&state_view);
     let vm = AptosVM::new(env.clone(), &state_view);
     let resolver = state_view.as_move_resolver();
-    let module_storage = state_view.as_aptos_code_storage(env.runtime_environment());
+    let module_storage = state_view.as_aptos_code_storage(env);
 
     let gas_schedule =
         GasScheduleV2::fetch_config(&state_view).context("failed to fetch gas schedule v2")?;
@@ -627,7 +624,7 @@ pub async fn simulate_multistep_proposal(
         let log_context = AdapterLogSchema::new(state_view.id(), 0);
 
         let resolver = state_view.as_move_resolver();
-        let code_storage = state_view.as_aptos_code_storage(env.runtime_environment());
+        let code_storage = state_view.as_aptos_code_storage(env);
 
         let (_vm_status, vm_output) = vm.execute_user_transaction(
             &resolver,

aptos-move/aptos-vm-profiling/src/bins/run_move.rs

@@ -188,7 +188,7 @@ fn main() -> Result<()> {
     let mut sess = vm.new_session_with_extensions(&storage, extensions);
 
     let traversal_storage = TraversalStorage::new();
-    let code_storage = storage.as_unsync_code_storage(&runtime_environment);
+    let code_storage = storage.as_unsync_code_storage(runtime_environment);
 
     let args: Vec<Vec<u8>> = vec![];
     match entrypoint {

aptos-move/aptos-vm-types/src/module_and_script_storage/state_view_adapter.rs

@@ -13,16 +13,16 @@ use move_binary_format::{
 use move_core_types::{account_address::AccountAddress, identifier::IdentStr, metadata::Metadata};
 use move_vm_runtime::{
     ambassador_impl_CodeStorage, ambassador_impl_ModuleStorage,
-    ambassador_impl_WithRuntimeEnvironment, AsUnsyncCodeStorage, CodeStorage, Module,
-    ModuleStorage, RuntimeEnvironment, Script, UnsyncCodeStorage, UnsyncModuleStorage,
+    ambassador_impl_WithRuntimeEnvironment, AsUnsyncCodeStorage, BorrowedOrOwned, CodeStorage,
+    Module, ModuleStorage, RuntimeEnvironment, Script, UnsyncCodeStorage, UnsyncModuleStorage,
     WithRuntimeEnvironment,
 };
 use move_vm_types::{code::ModuleBytesStorage, module_storage_error};
 use std::sync::Arc;
 
 /// Avoids orphan rule to implement [ModuleBytesStorage] for [StateView].
 struct StateViewAdapter<'s, S> {
-    state_view: &'s S,
+    state_view: BorrowedOrOwned<'s, S>,
 }
 
 impl<'s, S: StateView> ModuleBytesStorage for StateViewAdapter<'s, S> {
@@ -42,37 +42,52 @@ impl<'s, S: StateView> ModuleBytesStorage for StateViewAdapter<'s, S> {
 /// directly by clients - only via [AsAptosCodeStorage] trait. Can be used to resolve both modules
 /// and cached scripts.
 #[derive(Delegate)]
-#[delegate(WithRuntimeEnvironment, where = "S: StateView")]
-#[delegate(ModuleStorage, where = "S: StateView")]
-#[delegate(CodeStorage, where = "S: StateView")]
-pub struct AptosCodeStorageAdapter<'s, S> {
-    storage: UnsyncCodeStorage<UnsyncModuleStorage<'s, StateViewAdapter<'s, S>>>,
+#[delegate(
+    WithRuntimeEnvironment,
+    where = "S: StateView, E: WithRuntimeEnvironment"
+)]
+#[delegate(ModuleStorage, where = "S: StateView, E: WithRuntimeEnvironment")]
+#[delegate(CodeStorage, where = "S: StateView, E: WithRuntimeEnvironment")]
+pub struct AptosCodeStorageAdapter<'s, S, E> {
+    storage: UnsyncCodeStorage<UnsyncModuleStorage<'s, StateViewAdapter<'s, S>, E>>,
 }
 
-impl<'s, S: StateView> AptosCodeStorageAdapter<'s, S> {
+impl<'s, S: StateView, E: WithRuntimeEnvironment> AptosCodeStorageAdapter<'s, S, E> {
     /// Creates new instance of [AptosCodeStorageAdapter] built on top of the passed state view and
     /// the provided runtime environment.
-    fn new(state_view: &'s S, runtime_environment: &'s RuntimeEnvironment) -> Self {
-        let adapter = StateViewAdapter { state_view };
+    fn from_borrowed(state_view: &'s S, runtime_environment: E) -> Self {
+        let adapter = StateViewAdapter {
+            state_view: BorrowedOrOwned::Borrowed(state_view),
+        };
         let storage = adapter.into_unsync_code_storage(runtime_environment);
         Self { storage }
     }
 
-    /// Returns the state view used by [UnsyncCodeStorage] as the raw byte storage.
-    fn state_view(&self) -> &'s S {
-        self.storage.module_storage().byte_storage().state_view
+    /// Creates new instance of [AptosCodeStorageAdapter] capturing the passed state view and the
+    /// provided environment.
+    fn from_owned(state_view: S, runtime_environment: E) -> Self {
+        let adapter = StateViewAdapter {
+            state_view: BorrowedOrOwned::Owned(state_view),
+        };
+        let storage = adapter.into_unsync_code_storage(runtime_environment);
+        Self { storage }
     }
 }
 
-impl<'s, S: StateView> AptosModuleStorage for AptosCodeStorageAdapter<'s, S> {
+impl<'s, S: StateView, E: WithRuntimeEnvironment> AptosModuleStorage
+    for AptosCodeStorageAdapter<'s, S, E>
+{
     fn fetch_state_value_metadata(
         &self,
         address: &AccountAddress,
         module_name: &IdentStr,
     ) -> PartialVMResult<Option<StateValueMetadata>> {
         let state_key = StateKey::module(address, module_name);
         Ok(self
-            .state_view()
+            .storage
+            .module_storage()
+            .byte_storage()
+            .state_view
             .get_state_value(&state_key)
             .map_err(|err| module_storage_error!(address, module_name, err).to_partial())?
             .map(|state_value| state_value.into_metadata()))
@@ -82,18 +97,19 @@ impl<'s, S: StateView> AptosModuleStorage for AptosCodeStorageAdapter<'s, S> {
 /// Allows to treat the state view as a code storage with scripts and modules. The main use case is
 /// when a transaction or a Move function has to be executed outside the long-living environment or
 /// block executor, e.g., for single transaction simulation, in Aptos debugger, etc.
-pub trait AsAptosCodeStorage<'s, S> {
-    fn as_aptos_code_storage(
-        &'s self,
-        runtime_environment: &'s RuntimeEnvironment,
-    ) -> AptosCodeStorageAdapter<'s, S>;
+pub trait AsAptosCodeStorage<'s, S, E> {
+    fn as_aptos_code_storage(&'s self, runtime_environment: E)
+        -> AptosCodeStorageAdapter<'s, S, E>;
+
+    fn into_aptos_code_storage(self, runtime_environment: E) -> AptosCodeStorageAdapter<'s, S, E>;
 }
 
-impl<'s, S: StateView> AsAptosCodeStorage<'s, S> for S {
-    fn as_aptos_code_storage(
-        &'s self,
-        runtime_environment: &'s RuntimeEnvironment,
-    ) -> AptosCodeStorageAdapter<S> {
-        AptosCodeStorageAdapter::new(self, runtime_environment)
+impl<'s, S: StateView, E: WithRuntimeEnvironment> AsAptosCodeStorage<'s, S, E> for S {
+    fn as_aptos_code_storage(&'s self, runtime_environment: E) -> AptosCodeStorageAdapter<S, E> {
+        AptosCodeStorageAdapter::from_borrowed(self, runtime_environment)
+    }
+
+    fn into_aptos_code_storage(self, runtime_environment: E) -> AptosCodeStorageAdapter<'s, S, E> {
+        AptosCodeStorageAdapter::from_owned(self, runtime_environment)
     }
 }

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

@@ -229,11 +229,11 @@ impl AptosVM {
         let _timer = TIMER.timer_with(&["AptosVM::new"]);
 
         let resolver = state_view.as_move_resolver();
-        let move_vm = MoveVmExt::new(env, &resolver);
+        let move_vm = MoveVmExt::new(env.clone(), &resolver);
 
         // We use an `Option` to handle the VK not being set on-chain, or an incorrect VK being set
         // via governance (although, currently, we do check for that in `keyless_account.move`).
-        let module_storage = state_view.as_aptos_code_storage(move_vm.env.runtime_environment());
+        let module_storage = state_view.as_aptos_code_storage(env);
         let pvk = keyless_validation::get_groth16_vk_onchain(&resolver, &module_storage)
             .ok()
             .and_then(|vk| vk.try_into().ok());
@@ -301,6 +301,11 @@ impl AptosVM {
         self.move_vm.env.runtime_environment()
     }
 
+    #[inline(always)]
+    pub fn environment(&self) -> AptosEnvironment {
+        self.move_vm.env.clone()
+    }
+
     /// Sets execution concurrency level when invoked the first time.
     pub fn set_concurrency_level_once(mut concurrency_level: usize) {
         concurrency_level = min(concurrency_level, num_cpus::get());
@@ -2436,7 +2441,7 @@ impl AptosVM {
         max_gas_amount: u64,
     ) -> ViewFunctionOutput {
         let env = AptosEnvironment::new(state_view);
-        let vm = AptosVM::new(env, state_view);
+        let vm = AptosVM::new(env.clone(), state_view);
 
         let log_context = AdapterLogSchema::new(state_view.id(), 0);
 
@@ -2462,7 +2467,7 @@ impl AptosVM {
         );
 
         let resolver = state_view.as_move_resolver();
-        let module_storage = state_view.as_aptos_code_storage(vm.runtime_environment());
+        let module_storage = state_view.as_aptos_code_storage(env);
 
         let mut session = vm.new_session(&resolver, SessionId::Void, None);
         let execution_result = Self::execute_view_function_in_vm(
@@ -2849,6 +2854,7 @@ impl VMValidator for AptosVM {
         &self,
         transaction: SignedTransaction,
         state_view: &impl StateView,
+        module_storage: &impl AptosCodeStorage,
     ) -> VMValidatorResult {
         let _timer = TXN_VALIDATION_SECONDS.start_timer();
         let log_context = AdapterLogSchema::new(state_view.id(), 0);
@@ -2899,8 +2905,6 @@ impl VMValidator for AptosVM {
         let txn_data = TransactionMetadata::new(&txn);
 
         let resolver = self.as_move_resolver(&state_view);
-        let module_storage = state_view.as_aptos_code_storage(self.runtime_environment());
-
         let is_approved_gov_script = is_approved_gov_script(&resolver, &txn, &txn_data);
 
         let mut session = self.new_session(
@@ -2915,7 +2919,7 @@ impl VMValidator for AptosVM {
         let (counter_label, result) = match self.validate_signed_transaction(
             &mut session,
             &resolver,
-            &module_storage,
+            module_storage,
             &txn,
             &txn_data,
             &log_context,
@@ -2963,11 +2967,11 @@ impl AptosSimulationVM {
         );
 
         let env = AptosEnvironment::new(state_view);
-        let vm = Self::new(env, state_view);
+        let vm = Self::new(env.clone(), state_view);
         let log_context = AdapterLogSchema::new(state_view.id(), 0);
 
         let resolver = state_view.as_move_resolver();
-        let code_storage = state_view.as_aptos_code_storage(vm.0.runtime_environment());
+        let code_storage = state_view.as_aptos_code_storage(env);
 
         let (vm_status, vm_output) =
             vm.0.execute_user_transaction(&resolver, &code_storage, transaction, &log_context);

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

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

aptos-move/aptos-vm/src/move_vm_ext/write_op_converter.rs

@@ -311,7 +311,6 @@ mod tests {
         identifier::Identifier,
         language_storage::{StructTag, TypeTag},
     };
-    use move_vm_runtime::WithRuntimeEnvironment;
 
     fn raw_metadata(v: u64) -> StateValueMetadata {
         StateValueMetadata::legacy(v, &CurrentTimeMicroseconds { microseconds: v })
@@ -419,7 +418,7 @@ mod tests {
         ]));
         let resolver = state_view.as_move_resolver();
         let env = AptosEnvironment::new(&state_view);
-        let code_storage = state_view.as_aptos_code_storage(env.runtime_environment());
+        let code_storage = state_view.as_aptos_code_storage(env);
         // Storage slot metadata is enabled on the mainnet.
         let woc = WriteOpConverter::new(&resolver, true);
 

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

@@ -102,7 +102,7 @@ impl AptosVM {
             .change_set_configs;
 
         let resolver = state_view.as_move_resolver();
-        let module_storage = state_view.as_aptos_code_storage(self.runtime_environment());
+        let module_storage = state_view.as_aptos_code_storage(self.environment());
 
         let traversal_storage = TraversalStorage::new();
         self.failed_transaction_cleanup(

aptos-move/block-executor/Cargo.toml

@@ -13,8 +13,8 @@ repository = { workspace = true }
 rust-version = { workspace = true }
 
 [dependencies]
-anyhow = { workspace = true }
 ambassador = { workspace = true }
+anyhow = { workspace = true }
 aptos-aggregator = { workspace = true }
 aptos-drop-helper = { workspace = true }
 aptos-infallible = { workspace = true }