runtime: start lifting contract preparation up through tx runtime layers

runtime/near-vm-runner/src/lib.rs

@@ -36,7 +36,7 @@ pub use code::ContractCode;
 #[cfg(feature = "metrics")]
 pub use metrics::{report_metrics, reset_metrics};
 pub use profile::ProfileDataV3;
-pub use runner::{run, prepare, PreparedContract, VM};
+pub use runner::{run, prepare, PreparedContract, Contract, VM};
 
 /// This is public for internal experimentation use only, and should otherwise be considered an
 /// implementation detail of `near-vm-runner`.

runtime/near-vm-runner/src/logic/dependencies.rs

@@ -488,10 +488,4 @@ pub trait External {
     ///
     /// Panics if `ReceiptIndex` is invalid.
     fn get_receipt_receiver(&self, receipt_index: ReceiptIndex) -> &AccountId;
-
-    /// Hash of the contract for the current account.
-    fn code_hash(&self) -> CryptoHash;
-
-    /// Get the contract code
-    fn get_contract(&self) -> Option<std::sync::Arc<crate::ContractCode>>;
 }

runtime/near-vm-runner/src/logic/mocks/mock_external.rs

@@ -320,12 +320,14 @@ impl External for MockedExternal {
             _ => panic!("not a valid receipt index!"),
         }
     }
+}
 
-    fn code_hash(&self) -> CryptoHash {
+impl crate::Contract for MockedExternal {
+    fn hash(&self) -> CryptoHash {
         self.code_hash
     }
 
-    fn get_contract(&self) -> Option<Arc<ContractCode>> {
+    fn get_code(&self) -> Option<Arc<ContractCode>> {
         self.code.clone()
     }
 }

runtime/near-vm-runner/src/near_vm_runner/runner.rs

@@ -9,7 +9,7 @@ use crate::logic::{Config, ExecutionResultState, External, VMContext, VMLogic, V
 use crate::near_vm_runner::{NearVmCompiler, NearVmEngine};
 use crate::runner::VMResult;
 use crate::{
-    get_contract_cache_key, imports, CompiledContract, ContractCode, ContractRuntimeCache,
+    get_contract_cache_key, imports, CompiledContract, Contract, ContractCode, ContractRuntimeCache,
 };
 use crate::{prepare, NoContractRuntimeCache};
 use memoffset::offset_of;
@@ -211,7 +211,7 @@ impl NearVM {
     fn with_compiled_and_loaded(
         self: Box<Self>,
         cache: &dyn ContractRuntimeCache,
-        ext: &dyn External,
+        contract: &dyn Contract,
         context: &VMContext,
         closure: impl FnOnce(ExecutionResultState, &VMArtifact, Box<Self>) -> VMResult<PreparedContract>,
     ) -> VMResult<PreparedContract> {
@@ -221,7 +221,7 @@ impl NearVM {
         // To identify a cache hit from either in-memory and on-disk cache correctly, we first assume that we have a cache hit here,
         // and then we set it to false when we fail to find any entry and decide to compile (by calling compile_and_cache below).
         let mut is_cache_hit = true;
-        let code_hash = ext.code_hash();
+        let code_hash = contract.hash();
         let (wasm_bytes, artifact_result) = cache.memory_cache().try_lookup(
             code_hash,
             || {
@@ -235,7 +235,7 @@ impl NearVM {
                 let key = get_contract_cache_key(code_hash, &self.config);
                 let cache_record = cache.get(&key).map_err(CacheError::ReadError)?;
                 let Some(compiled_contract_info) = cache_record else {
-                    let Some(code) = ext.get_contract() else {
+                    let Some(code) = contract.get_code() else {
                         return Err(VMRunnerError::ContractCodeNotPresent);
                     };
                     let _span =
@@ -574,13 +574,16 @@ impl<'a> finite_wasm::wasmparser::VisitOperator<'a> for GasCostCfg {
 impl crate::runner::VM for NearVM {
     fn prepare(
         self: Box<Self>,
-        ext: &dyn External,
+        contract: &dyn Contract,
         context: &VMContext,
         cache: Option<&dyn ContractRuntimeCache>,
     ) -> Box<dyn crate::PreparedContract> {
         let cache = cache.unwrap_or(&NoContractRuntimeCache);
-        let prepd =
-            self.with_compiled_and_loaded(cache, ext, context, |result_state, artifact, vm| {
+        let prepd = self.with_compiled_and_loaded(
+            cache,
+            contract,
+            context,
+            |result_state, artifact, vm| {
                 let memory = NearVmMemory::new(
                     vm.config.limit_config.initial_memory_pages,
                     vm.config.limit_config.max_memory_pages,
@@ -601,7 +604,8 @@ impl crate::runner::VM for NearVM {
                     artifact: Arc::clone(artifact),
                     vm,
                 }))
-            });
+            },
+        );
         Box::new(prepd)
     }
 

runtime/near-vm-runner/src/runner.rs

@@ -23,7 +23,6 @@ use std::sync::Arc;
 /// validators, even when a guest error occurs, or else their state will diverge.
 pub(crate) type VMResult<T = VMOutcome> = Result<T, VMRunnerError>;
 
-
 /// Prepare the contract for execution.
 ///
 /// The returned value does some work in preparation to execute the contract, without executing any
@@ -33,14 +32,14 @@ pub(crate) type VMResult<T = VMOutcome> = Result<T, VMRunnerError>;
 ///
 /// Contract preparation and execution need not to be executed on the same thread.
 #[tracing::instrument(target = "vm", level = "debug", "prepare", skip_all, fields(
-    code.hash = %ext.code_hash(),
+    code.hash = %contract.hash(),
     method_name,
     vm_kind = ?wasm_config.vm_kind,
     burnt_gas = tracing::field::Empty,
     compute_usage = tracing::field::Empty,
 ))]
 pub fn prepare(
-    ext: &(dyn External + Send),
+    contract: &dyn Contract,
     context: &VMContext,
     wasm_config: Arc<Config>,
     cache: Option<&dyn ContractRuntimeCache>,
@@ -49,7 +48,7 @@ pub fn prepare(
     let runtime = vm_kind
         .runtime(wasm_config)
         .unwrap_or_else(|| panic!("the {vm_kind:?} runtime has not been enabled at compile time"));
-    runtime.prepare(ext, context, cache)
+    runtime.prepare(contract, context, cache)
 }
 
 /// Validate and run the specified contract.
@@ -69,7 +68,6 @@ pub fn prepare(
 /// The gas cost for contract preparation will be subtracted by the VM
 /// implementation.
 #[tracing::instrument(target = "vm", level = "debug", "run", skip_all, fields(
-    code.hash = %ext.code_hash(),
     method_name,
     burnt_gas = tracing::field::Empty,
     compute_usage = tracing::field::Empty,
@@ -107,6 +105,18 @@ pub trait PreparedContract: Send {
     ) -> VMResult;
 }
 
+/// Trait encapsulating access to the contract's WASM source code.
+pub trait Contract {
+    /// Hash of the contract for the current account.
+    fn hash(&self) -> near_primitives_core::hash::CryptoHash;
+
+    /// Get the contract code.
+    ///
+    /// The runtime might not call this if it finds e.g. a compiled contract inside the supplied
+    /// cache.
+    fn get_code(&self) -> Option<std::sync::Arc<crate::ContractCode>>;
+}
+
 pub trait VM {
     /// Prepare a contract for execution.
     ///
@@ -120,7 +130,7 @@ pub trait VM {
     /// will be reported when the returned value is `run`.
     fn prepare(
         self: Box<Self>,
-        ext: &dyn External,
+        ext: &dyn Contract,
         context: &VMContext,
         cache: Option<&dyn ContractRuntimeCache>,
     ) -> Box<dyn PreparedContract>;

runtime/near-vm-runner/src/wasmer2_runner.rs

@@ -7,9 +7,9 @@ use crate::logic::gas_counter::FastGasCounter;
 use crate::logic::{
     Config, ExecutionResultState, External, MemSlice, MemoryLike, VMContext, VMLogic, VMOutcome,
 };
-use crate::prepare;
 use crate::runner::VMResult;
 use crate::{get_contract_cache_key, imports, ContractCode};
+use crate::{prepare, Contract};
 use memoffset::offset_of;
 use near_parameters::vm::VMKind;
 use near_parameters::RuntimeFeesConfig;
@@ -578,12 +578,12 @@ impl crate::runner::VM for Wasmer2VM {
 
     fn prepare(
         self: Box<Self>,
-        ext: &dyn External,
+        contract: &dyn Contract,
         context: &VMContext,
         cache: Option<&dyn ContractRuntimeCache>,
     ) -> Box<dyn crate::PreparedContract> {
         type Result = VMResult<PreparedContract>;
-        let Some(code) = ext.get_contract() else {
+        let Some(code) = contract.get_code() else {
             return Box::new(Result::Err(VMRunnerError::ContractCodeNotPresent));
         };
         let mut result_state = ExecutionResultState::new(&context, Arc::clone(&self.config));

runtime/near-vm-runner/src/wasmer_runner.rs

@@ -5,7 +5,7 @@ use crate::logic::errors::{
 };
 use crate::logic::{ExecutionResultState, External, VMContext, VMLogic, VMLogicError, VMOutcome};
 use crate::logic::{MemSlice, MemoryLike};
-use crate::prepare;
+use crate::{prepare, Contract};
 use crate::runner::VMResult;
 use crate::{get_contract_cache_key, imports, ContractCode};
 use near_parameters::vm::{Config, VMKind};
@@ -429,12 +429,12 @@ impl crate::runner::VM for Wasmer0VM {
 
     fn prepare(
         self: Box<Self>,
-        ext: &dyn External,
+        contract: &dyn Contract,
         context: &VMContext,
         cache: Option<&dyn ContractRuntimeCache>,
     ) -> Box<dyn crate::PreparedContract> {
         type Result = VMResult<PreparedContract>;
-        let Some(code) = ext.get_contract() else {
+        let Some(code) = contract.get_code() else {
             return Box::new(Result::Err(VMRunnerError::ContractCodeNotPresent));
         };
         if !cfg!(target_arch = "x86") && !cfg!(target_arch = "x86_64") {

runtime/near-vm-runner/src/wasmtime_runner.rs

@@ -7,8 +7,8 @@ use crate::logic::{Config, ExecutionResultState};
 use crate::logic::{External, MemSlice, MemoryLike, VMContext, VMLogic, VMOutcome};
 use crate::runner::VMResult;
 use crate::{
-    get_contract_cache_key, imports, prepare, CompiledContract, CompiledContractInfo, ContractCode,
-    ContractRuntimeCache, NoContractRuntimeCache,
+    get_contract_cache_key, imports, prepare, CompiledContract, CompiledContractInfo, Contract,
+    ContractCode, ContractRuntimeCache, NoContractRuntimeCache,
 };
 use near_parameters::vm::VMKind;
 use near_parameters::RuntimeFeesConfig;
@@ -186,11 +186,11 @@ impl WasmtimeVM {
     fn with_compiled_and_loaded(
         &self,
         cache: &dyn ContractRuntimeCache,
-        ext: &dyn External,
+        contract: &dyn Contract,
         context: &VMContext,
         closure: impl FnOnce(ExecutionResultState, Module) -> VMResult<PreparedContract>,
     ) -> VMResult<PreparedContract> {
-        let code_hash = ext.code_hash();
+        let code_hash = contract.hash();
         type MemoryCacheType = (u64, Result<Module, CompilationError>);
         let to_any = |v: MemoryCacheType| -> Box<dyn std::any::Any + Send> { Box::new(v) };
         let (wasm_bytes, module_result) = cache.memory_cache().try_lookup(
@@ -199,7 +199,7 @@ impl WasmtimeVM {
                 let key = get_contract_cache_key(code_hash, &self.config);
                 let cache_record = cache.get(&key).map_err(CacheError::ReadError)?;
                 let Some(compiled_contract_info) = cache_record else {
-                    let Some(code) = ext.get_contract() else {
+                    let Some(code) = contract.get_code() else {
                         return Err(VMRunnerError::ContractCodeNotPresent);
                     };
                     return Ok(to_any((
@@ -283,12 +283,12 @@ impl crate::runner::VM for WasmtimeVM {
 
     fn prepare(
         self: Box<Self>,
-        ext: &dyn External,
+        code: &dyn Contract,
         context: &VMContext,
         cache: Option<&dyn ContractRuntimeCache>,
     ) -> Box<dyn crate::PreparedContract> {
         let cache = cache.unwrap_or(&NoContractRuntimeCache);
-        let prepd = self.with_compiled_and_loaded(cache, ext, context, |result_state, module| {
+        let prepd = self.with_compiled_and_loaded(cache, code, context, |result_state, module| {
             match module.get_export(&context.method) {
                 Some(export) => match export {
                     Func(func_type) => {