[aptos-vm] Handle runtime environment nicely

aptos-move/aptos-vm-types/Cargo.toml

@@ -18,7 +18,6 @@ anyhow = { workspace = true }
 aptos-aggregator = { workspace = true }
 aptos-gas-algebra = { workspace = true }
 aptos-gas-schedule = { workspace = true }
-aptos-native-interface = { workspace = true }
 aptos-types = { workspace = true }
 bcs = { workspace = true }
 bytes = { workspace = true }

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

@@ -112,7 +112,7 @@ use move_vm_runtime::{
     logging::expect_no_verification_errors,
     module_linker_error,
     module_traversal::{TraversalContext, TraversalStorage},
-    TemporaryModuleStorage, UnreachableCodeStorage,
+    RuntimeEnvironment, TemporaryModuleStorage, UnreachableCodeStorage,
 };
 use move_vm_types::gas::{GasMeter, UnmeteredGasMeter};
 use num_cpus;
@@ -312,6 +312,11 @@ impl AptosVM {
         self.move_vm.env.chain_id()
     }
 
+    #[inline(always)]
+    pub fn runtime_environment(&self) -> &RuntimeEnvironment {
+        self.move_vm.runtime_environment()
+    }
+
     /// 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());

aptos-move/aptos-vm/src/block_executor/vm_wrapper.rs

@@ -20,13 +20,20 @@ use aptos_vm_types::{
 };
 use fail::fail_point;
 use move_core_types::vm_status::{StatusCode, VMStatus};
+use move_vm_runtime::{RuntimeEnvironment, WithRuntimeEnvironment};
 use std::sync::Arc;
 
 pub(crate) struct AptosExecutorTask {
     vm: AptosVM,
     id: StateViewId,
 }
 
+impl WithRuntimeEnvironment for AptosExecutorTask {
+    fn runtime_environment(&self) -> &RuntimeEnvironment {
+        self.vm.runtime_environment()
+    }
+}
+
 impl ExecutorTask for AptosExecutorTask {
     type Environment = Arc<Environment>;
     type Error = VMStatus;

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

@@ -154,14 +154,14 @@ impl MoveVmExt {
             .features()
             .is_enabled(FeatureFlag::DISALLOW_USER_NATIVES);
 
-        // TODO(loader_v2): Instead, propagate environment down to Block-STM.
+        // TODO(loader_v2): Refresh environment when configs change.
         set_runtime_environment(
             vm_config.clone(),
             aptos_natives_with_builder(&mut builder, inject_create_signer_for_gov_sim),
         );
         let vm = if env.features().is_loader_v2_enabled() {
             // TODO(loader_v2): For now re-create the VM every time. Later we can have a
-            //                  single VM created once.
+            //                  single VM created once, which also holds the environment.
             MoveVM::new_with_runtime_environment(fetch_runtime_environment())
         } else {
             WarmVmCache::get_warm_vm(

aptos-move/block-executor/Cargo.toml

@@ -20,7 +20,6 @@ aptos-infallible = { workspace = true }
 aptos-logger = { workspace = true }
 aptos-metrics-core = { workspace = true }
 aptos-mvhashmap = { workspace = true }
-aptos-framework = { workspace = true }
 aptos-types = { workspace = true }
 aptos-vm-logging = { workspace = true }
 aptos-vm-types = { workspace = true }

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

@@ -45,9 +45,7 @@ use aptos_types::{
     write_set::{TransactionWrite, WriteOp},
 };
 use aptos_vm_logging::{alert, clear_speculative_txn_logs, init_speculative_logs, prelude::*};
-use aptos_vm_types::{
-    change_set::randomly_check_layout_matches, environment::fetch_runtime_environment,
-};
+use aptos_vm_types::change_set::randomly_check_layout_matches;
 use bytes::Bytes;
 use claims::assert_none;
 use core::panic;
@@ -826,7 +824,6 @@ where
         scheduler: &Scheduler,
         // TODO: should not need to pass base view.
         base_view: &S,
-        runtime_environment: &RuntimeEnvironment,
         start_shared_counter: u32,
         shared_counter: &AtomicU32,
         shared_commit_state: &ExplicitSyncWrapper<BlockGasLimitProcessor<T>>,
@@ -838,6 +835,9 @@ where
         let executor = E::init(env.clone(), base_view);
         drop(init_timer);
 
+        // Shared environment used by each executor.
+        let runtime_environment = executor.runtime_environment();
+
         let _timer = WORK_WITH_TASK_SECONDS.start_timer();
         let mut scheduler_task = SchedulerTask::Retry;
 
@@ -987,7 +987,6 @@ where
         let scheduler = Scheduler::new(num_txns);
 
         let timer = RAYON_EXECUTION_SECONDS.start_timer();
-        let runtime_environment = fetch_runtime_environment();
         self.executor_thread_pool.scope(|s| {
             for _ in 0..num_workers {
                 s.spawn(|_| {
@@ -998,7 +997,6 @@ where
                         &versioned_cache,
                         &scheduler,
                         base_view,
-                        &runtime_environment,
                         start_shared_counter,
                         &shared_counter,
                         &shared_commit_state,
@@ -1135,6 +1133,8 @@ where
         let executor = E::init(env, base_view);
         drop(init_timer);
 
+        let runtime_environment = executor.runtime_environment();
+
         let start_counter = gen_id_start_value(true);
         let counter = RefCell::new(start_counter);
         let unsync_map = UnsyncMap::new();
@@ -1147,11 +1147,10 @@ where
         let last_input_output: TxnLastInputOutput<T, E::Output, E::Error> =
             TxnLastInputOutput::new(num_txns as TxnIndex);
 
-        let runtime_environment = fetch_runtime_environment();
         for (idx, txn) in signature_verified_block.iter().enumerate() {
             let latest_view = LatestView::<T, S, X>::new(
                 base_view,
-                runtime_environment.as_ref(),
+                runtime_environment,
                 ViewState::Unsync(SequentialState::new(&unsync_map, start_counter, &counter)),
                 idx as TxnIndex,
             );
@@ -1321,7 +1320,7 @@ where
                     // Apply the writes.
                     let resource_write_set = output.resource_write_set();
                     Self::apply_output_sequential(
-                        runtime_environment.as_ref(),
+                        runtime_environment,
                         &unsync_map,
                         &output,
                         resource_write_set.clone(),

aptos-move/block-executor/src/proptest_types/types.rs

@@ -33,6 +33,7 @@ use aptos_vm_types::{
 use bytes::Bytes;
 use claims::{assert_ge, assert_le, assert_ok};
 use move_core_types::{identifier::IdentStr, value::MoveTypeLayout};
+use move_vm_runtime::{RuntimeEnvironment, WithRuntimeEnvironment};
 use move_vm_types::delayed_values::delayed_field_id::DelayedFieldID;
 use once_cell::sync::OnceCell;
 use proptest::{arbitrary::Arbitrary, collection::vec, prelude::*, proptest, sample::Index};
@@ -843,12 +844,23 @@ impl<V: Into<Vec<u8>> + Arbitrary + Clone + Debug + Eq + Sync + Send> Transactio
 // Mock transaction executor implementation.
 ///////////////////////////////////////////////////////////////////////////
 
-#[derive(Default)]
-pub(crate) struct MockTask<K, E>(PhantomData<(K, E)>);
+pub(crate) struct MockTask<K, E> {
+    runtime_environment: RuntimeEnvironment,
+    phantom_data: PhantomData<(K, E)>,
+}
 
 impl<K, E> MockTask<K, E> {
     pub fn new() -> Self {
-        Self(PhantomData)
+        Self {
+            runtime_environment: RuntimeEnvironment::test(),
+            phantom_data: PhantomData,
+        }
+    }
+}
+
+impl<K, E> WithRuntimeEnvironment for MockTask<K, E> {
+    fn runtime_environment(&self) -> &RuntimeEnvironment {
+        &self.runtime_environment
     }
 }
 

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

@@ -19,6 +19,7 @@ use aptos_vm_types::{
     resolver::{TExecutorView, TResourceGroupView},
 };
 use move_core_types::{value::MoveTypeLayout, vm_status::StatusCode};
+use move_vm_runtime::WithRuntimeEnvironment;
 use std::{
     collections::{BTreeMap, HashSet},
     fmt::Debug,
@@ -52,7 +53,7 @@ pub struct Accesses<K> {
 
 /// Trait for single threaded transaction executor.
 // TODO: Sync should not be required. Sync is only introduced because this trait occurs as a phantom type of executor struct.
-pub trait ExecutorTask: Sync {
+pub trait ExecutorTask: Sync + WithRuntimeEnvironment {
     /// Type of transaction and its associated key and value.
     type Txn: Transaction;
 

aptos-move/e2e-move-tests/src/tests/code_publishing.rs

@@ -11,7 +11,7 @@ use aptos_types::{
     account_address::{create_resource_address, AccountAddress},
     move_utils::MemberId,
     on_chain_config::{FeatureFlag, Features, OnChainConfig},
-    transaction::{AbortInfo, ExecutionStatus, TransactionPayload},
+    transaction::{ExecutionStatus, TransactionPayload},
 };
 use claims::assert_ok;
 use move_core_types::{

aptos-move/mvhashmap/Cargo.toml

@@ -26,8 +26,6 @@ derivative = { workspace = true }
 move-binary-format = { workspace = true }
 move-core-types = { workspace = true }
 move-vm-types = { workspace = true }
-move-vm-runtime = { workspace = true }
-parking_lot = { workspace = true }
 serde = { workspace = true }
 
 [dev-dependencies]

aptos-move/vm-genesis/Cargo.toml

@@ -22,7 +22,6 @@ aptos-vm = { workspace = true }
 aptos-vm-types = { workspace = true }
 bcs = { workspace = true }
 bytes = { workspace = true }
-claims = { workspace = true }
 move-core-types = { workspace = true }
 move-vm-runtime = { workspace = true }
 move-vm-types = { workspace = true }

third_party/move/move-vm/runtime/src/lib.rs

@@ -36,7 +36,7 @@ pub use loader::{LoadedFunction, Module, Script};
 pub use storage::{
     code_storage::{ambassador_impl_CodeStorage, script_hash, CodeStorage},
     dummy::use_loader_v2_based_on_env,
-    environment::RuntimeEnvironment,
+    environment::{RuntimeEnvironment, WithRuntimeEnvironment},
     implementations::{
         unreachable_code_storage::UnreachableCodeStorage,
         unsync_code_storage::{IntoUnsyncCodeStorage, UnsyncCodeStorage},

third_party/move/move-vm/runtime/src/storage/environment.rs

@@ -180,6 +180,6 @@ impl RuntimeEnvironment {
 }
 
 /// Represents any type that contains a [RuntimeEnvironment].
-pub trait WithEnvironment {
+pub trait WithRuntimeEnvironment {
     fn runtime_environment(&self) -> &RuntimeEnvironment;
 }

third_party/move/move-vm/runtime/src/storage/implementations/unsync_code_storage.rs

@@ -4,7 +4,7 @@
 use crate::{
     script_hash,
     storage::{
-        environment::WithEnvironment,
+        environment::WithRuntimeEnvironment,
         implementations::unsync_module_storage::IntoUnsyncModuleStorage,
         module_storage::{ambassador_impl_ModuleStorage, ModuleBytesStorage},
     },
@@ -61,7 +61,7 @@ impl<'e, B: ModuleBytesStorage> IntoUnsyncCodeStorage<'e, B> for B {
     }
 }
 
-impl<M: ModuleStorage + WithEnvironment> UnsyncCodeStorage<M> {
+impl<M: ModuleStorage + WithRuntimeEnvironment> UnsyncCodeStorage<M> {
     /// Create a new storage with no scripts. There are no constrains on which
     /// modules exist in module storage.
     fn new(module_storage: M) -> Self {
@@ -123,7 +123,7 @@ impl<M: ModuleStorage + WithEnvironment> UnsyncCodeStorage<M> {
     }
 }
 
-impl<M: ModuleStorage + WithEnvironment> CodeStorage for UnsyncCodeStorage<M> {
+impl<M: ModuleStorage + WithRuntimeEnvironment> CodeStorage for UnsyncCodeStorage<M> {
     fn deserialize_and_cache_script(
         &self,
         serialized_script: &[u8],
@@ -174,7 +174,7 @@ impl<M: ModuleStorage + WithEnvironment> CodeStorage for UnsyncCodeStorage<M> {
 }
 
 #[cfg(test)]
-impl<M: ModuleStorage + WithEnvironment> UnsyncCodeStorage<M> {
+impl<M: ModuleStorage + WithRuntimeEnvironment> UnsyncCodeStorage<M> {
     fn matches<P: Fn(&ScriptStorageEntry) -> bool>(
         &self,
         script_hashes: impl IntoIterator<Item = [u8; 32]>,

third_party/move/move-vm/runtime/src/storage/implementations/unsync_module_storage.rs

@@ -4,7 +4,7 @@
 use crate::{
     module_cyclic_dependency_error, module_linker_error,
     storage::{
-        environment::{RuntimeEnvironment, WithEnvironment},
+        environment::{RuntimeEnvironment, WithRuntimeEnvironment},
         module_storage::{ModuleBytesStorage, ModuleStorage},
     },
     Module,
@@ -281,7 +281,7 @@ impl<'e, B: ModuleBytesStorage> UnsyncModuleStorage<'e, B> {
     }
 }
 
-impl<'e, B: ModuleBytesStorage> WithEnvironment for UnsyncModuleStorage<'e, B> {
+impl<'e, B: ModuleBytesStorage> WithRuntimeEnvironment for UnsyncModuleStorage<'e, B> {
     fn runtime_environment(&self) -> &RuntimeEnvironment {
         self.runtime_environment
     }