vm: split preparation and running of a contract (#11667)

This is a very exciting step forward! Finally we got up to the point
where we can do some work in preparing the contract to run separately
from actual running of the contract. And all of this is encapsulated in
a very neat API that gives out `Send + 'static` types for users to pass
around between threads or whatever so that they can pipeline these
processes.

It will remain to see whether the requirement to have `&External` and
`&VMContext` in both calls is a problem, and how much of a problem it
is, but that might be very well solvable with some scoped threads or
smart use of channels, or even just `Arc<Mutex>`, especially since both
of these structures generally tend to be unique to a contract
execution...

Part of #11319

runtime/near-vm-runner/fuzz/fuzz_targets/diffrunner.rs

@@ -20,7 +20,8 @@ libfuzzer_sys::fuzz_target!(|module: ArbitraryModule| {
 
 fn run_fuzz(code: &ContractCode, vm_kind: VMKind) -> VMOutcome {
     let mut fake_external = MockedExternal::with_code(code.clone_for_tests());
-    let mut context = create_context(vec![]);
+    let method_name = find_entry_point(code).unwrap_or_else(|| "main".to_string());
+    let mut context = create_context(&method_name, vec![]);
     context.prepaid_gas = 10u64.pow(14);
     let config_store = RuntimeConfigStore::new(None);
     let config = config_store.get_config(PROTOCOL_VERSION);
@@ -29,15 +30,11 @@ fn run_fuzz(code: &ContractCode, vm_kind: VMKind) -> VMOutcome {
     wasm_config.limit_config.contract_prepare_version =
         near_vm_runner::logic::ContractPrepareVersion::V2;
 
-    let method_name = find_entry_point(code).unwrap_or_else(|| "main".to_string());
-    let res = vm_kind.runtime(wasm_config.into()).unwrap().run(
-        &method_name,
-        &mut fake_external,
-        &context,
-        fees,
-        [].into(),
-        None,
-    );
+    let res = vm_kind
+        .runtime(wasm_config.into())
+        .unwrap()
+        .prepare(&fake_external, &context, None)
+        .run(&mut fake_external, &context, fees);
 
     // Remove the VMError message details as they can differ between runtimes
     // TODO: maybe there's actually things we could check for equality here too?

runtime/near-vm-runner/fuzz/fuzz_targets/runner.rs

@@ -18,16 +18,17 @@ libfuzzer_sys::fuzz_target!(|module: ArbitraryModule| {
 
 fn run_fuzz(code: &ContractCode, config: Arc<RuntimeConfig>) -> VMOutcome {
     let mut fake_external = MockedExternal::with_code(code.clone_for_tests());
-    let mut context = create_context(vec![]);
+    let method_name = find_entry_point(code).unwrap_or_else(|| "main".to_string());
+    let mut context = create_context(&method_name, vec![]);
     context.prepaid_gas = 10u64.pow(14);
     let mut wasm_config = near_parameters::vm::Config::clone(&config.wasm_config);
     wasm_config.limit_config.wasmer2_stack_limit = i32::MAX; // If we can crash wasmer2 even without the secondary stack limit it's still good to know
     let vm_kind = config.wasm_config.vm_kind;
     let fees = Arc::clone(&config.fees);
-    let method_name = find_entry_point(code).unwrap_or_else(|| "main".to_string());
     vm_kind
         .runtime(wasm_config.into())
         .unwrap()
-        .run(&method_name, &mut fake_external, &context, fees, [].into(), None)
+        .prepare(&fake_external, &context, None)
+        .run(&mut fake_external, &context, fees)
         .unwrap_or_else(|err| panic!("fatal error: {err:?}"))
 }

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

@@ -30,13 +30,15 @@ pub fn find_entry_point(contract: &ContractCode) -> Option<String> {
     None
 }
 
-pub fn create_context(input: Vec<u8>) -> VMContext {
+pub fn create_context(method: &str, input: Vec<u8>) -> VMContext {
     VMContext {
         current_account_id: "alice".parse().unwrap(),
         signer_account_id: "bob".parse().unwrap(),
         signer_account_pk: vec![0, 1, 2, 3, 4],
         predecessor_account_id: "carol".parse().unwrap(),
+        method: method.into(),
         input,
+        promise_results: Vec::new().into(),
         block_height: 10,
         block_timestamp: 42,
         epoch_height: 1,

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, VM};
+pub use runner::{run, PreparedContract, 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/context.rs

@@ -1,4 +1,4 @@
-use super::types::PublicKey;
+use super::types::{PromiseResult, PublicKey};
 use near_primitives_core::config::ViewConfig;
 use near_primitives_core::types::{
     AccountId, Balance, BlockHeight, EpochHeight, Gas, StorageUsage,
@@ -20,9 +20,14 @@ pub struct VMContext {
     /// If this execution is the result of direct execution of transaction then it
     /// is equal to `signer_account_id`.
     pub predecessor_account_id: AccountId,
+    /// The name of the method to invoke.
+    pub method: String,
     /// The input to the contract call.
     /// Encoded as base64 string to be able to pass input in borsh binary format.
     pub input: Vec<u8>,
+    /// If this method execution is invoked directly as a callback by one or more contract calls
+    /// the results of the methods that made the callback are stored in this collection.
+    pub promise_results: std::sync::Arc<[PromiseResult]>,
     /// The current block height.
     pub block_height: BlockHeight,
     /// The current block timestamp (number of non-leap-nanoseconds since January 1, 1970 0:00:00 UTC).

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

@@ -35,7 +35,7 @@ fn base64(s: &[u8]) -> String {
 /// This is a subset of [`VMLogic`] that's strictly necessary to produce `VMOutcome`s.
 pub struct ExecutionResultState {
     /// All gas and economic parameters required during contract execution.
-    config: Arc<Config>,
+    pub(crate) config: Arc<Config>,
     /// Gas tracking for the current contract execution.
     gas_counter: GasCounter,
     /// Logs written by the runtime.
@@ -219,9 +219,6 @@ pub struct VMLogic<'a> {
     config: Arc<Config>,
     /// Fees charged for various operations that contract may execute.
     fees_config: Arc<RuntimeFeesConfig>,
-    /// If this method execution is invoked directly as a callback by one or more contract calls the
-    /// results of the methods that made the callback are stored in this collection.
-    promise_results: Arc<[PromiseResult]>,
     /// Pointer to the guest memory.
     memory: super::vmstate::Memory,
 
@@ -303,7 +300,6 @@ impl<'a> VMLogic<'a> {
         ext: &'a mut dyn External,
         context: &'a VMContext,
         fees_config: Arc<RuntimeFeesConfig>,
-        promise_results: Arc<[PromiseResult]>,
         result_state: ExecutionResultState,
         memory: impl MemoryLike + 'static,
     ) -> Self {
@@ -319,7 +315,6 @@ impl<'a> VMLogic<'a> {
             context,
             config,
             fees_config,
-            promise_results,
             memory: super::vmstate::Memory::new(memory),
             current_account_locked_balance,
             recorded_storage_counter,
@@ -2381,7 +2376,7 @@ impl<'a> VMLogic<'a> {
             }
             .into());
         }
-        Ok(self.promise_results.len() as _)
+        Ok(self.context.promise_results.len() as _)
     }
 
     /// If the current function is invoked by a callback we can access the execution results of the
@@ -2414,6 +2409,7 @@ impl<'a> VMLogic<'a> {
             );
         }
         match self
+            .context
             .promise_results
             .get(result_idx as usize)
             .ok_or(HostError::InvalidPromiseResultIndex { result_idx })?

runtime/near-vm-runner/src/logic/tests/promises.rs

@@ -19,7 +19,7 @@ fn test_promise_results() {
     ];
 
     let mut logic_builder = VMLogicBuilder::default();
-    logic_builder.promise_results = promise_results.into();
+    logic_builder.context.promise_results = promise_results.into();
     let mut logic = logic_builder.build();
 
     assert_eq!(logic.promise_results_count(), Ok(3), "Total count of registers must be 3");

runtime/near-vm-runner/src/logic/tests/vm_logic_builder.rs

@@ -1,6 +1,5 @@
 use crate::logic::mocks::mock_external::MockedExternal;
 use crate::logic::mocks::mock_memory::MockedMemory;
-use crate::logic::types::PromiseResult;
 use crate::logic::{Config, ExecutionResultState, MemSlice, VMContext, VMLogic};
 use crate::tests::test_vm_config;
 use near_parameters::RuntimeFeesConfig;
@@ -10,7 +9,6 @@ pub(super) struct VMLogicBuilder {
     pub ext: MockedExternal,
     pub config: Config,
     pub fees_config: RuntimeFeesConfig,
-    pub promise_results: Arc<[PromiseResult]>,
     pub memory: MockedMemory,
     pub context: VMContext,
 }
@@ -22,7 +20,6 @@ impl Default for VMLogicBuilder {
             fees_config: RuntimeFeesConfig::test(),
             ext: MockedExternal::default(),
             memory: MockedMemory::default(),
-            promise_results: [].into(),
             context: get_context(),
         }
     }
@@ -43,7 +40,6 @@ impl VMLogicBuilder {
             &mut self.ext,
             &self.context,
             Arc::new(self.fees_config.clone()),
-            Arc::clone(&self.promise_results),
             result_state,
             self.memory.clone(),
         ))
@@ -59,7 +55,6 @@ impl VMLogicBuilder {
             fees_config: RuntimeFeesConfig::free(),
             ext: MockedExternal::default(),
             memory: MockedMemory::default(),
-            promise_results: [].into(),
             context: get_context(),
         }
     }
@@ -71,7 +66,9 @@ fn get_context() -> VMContext {
         signer_account_id: "bob.near".parse().unwrap(),
         signer_account_pk: vec![0, 1, 2, 3, 4],
         predecessor_account_id: "carol.near".parse().unwrap(),
+        method: "VMLogicBuilder::method_not_specified".into(),
         input: vec![0, 1, 2, 3, 4],
+        promise_results: vec![].into(),
         block_height: 10,
         block_timestamp: 42,
         epoch_height: 1,

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

@@ -5,7 +5,6 @@ use crate::logic::errors::{
     CacheError, CompilationError, FunctionCallError, MethodResolveError, VMRunnerError, WasmTrap,
 };
 use crate::logic::gas_counter::FastGasCounter;
-use crate::logic::types::PromiseResult;
 use crate::logic::{Config, ExecutionResultState, External, VMContext, VMLogic, VMOutcome};
 use crate::near_vm_runner::{NearVmCompiler, NearVmEngine};
 use crate::runner::VMResult;
@@ -210,17 +209,12 @@ impl NearVM {
         skip_all
     )]
     fn with_compiled_and_loaded(
-        &self,
+        self: Box<Self>,
         cache: &dyn ContractRuntimeCache,
-        ext: &mut dyn External,
+        ext: &dyn External,
         context: &VMContext,
-        method_name: &str,
-        closure: impl FnOnce(
-            ExecutionResultState,
-            &mut dyn External,
-            &VMArtifact,
-        ) -> Result<VMOutcome, VMRunnerError>,
-    ) -> VMResult<VMOutcome> {
+        closure: impl FnOnce(ExecutionResultState, &VMArtifact, Box<Self>) -> VMResult<PreparedContract>,
+    ) -> VMResult<PreparedContract> {
         // (wasm code size, compilation result)
         type MemoryCacheType = (u64, Result<VMArtifact, CompilationError>);
         let to_any = |v: MemoryCacheType| -> Box<dyn std::any::Any + Send> { Box::new(v) };
@@ -307,19 +301,22 @@ impl NearVM {
         crate::metrics::record_compiled_contract_cache_lookup(is_cache_hit);
 
         let mut result_state = ExecutionResultState::new(&context, Arc::clone(&self.config));
-        let result = result_state.before_loading_executable(method_name, wasm_bytes);
+        let result = result_state.before_loading_executable(&context.method, wasm_bytes);
         if let Err(e) = result {
-            return Ok(VMOutcome::abort(result_state, e));
+            return Ok(PreparedContract::Outcome(VMOutcome::abort(result_state, e)));
         }
         match artifact_result {
             Ok(artifact) => {
                 let result = result_state.after_loading_executable(wasm_bytes);
                 if let Err(e) = result {
-                    return Ok(VMOutcome::abort(result_state, e));
+                    return Ok(PreparedContract::Outcome(VMOutcome::abort(result_state, e)));
                 }
-                closure(result_state, ext, &artifact)
+                closure(result_state, &artifact, self)
             }
-            Err(e) => Ok(VMOutcome::abort(result_state, FunctionCallError::CompilationError(e))),
+            Err(e) => Ok(PreparedContract::Outcome(VMOutcome::abort(
+                result_state,
+                FunctionCallError::CompilationError(e),
+            ))),
         }
     }
 
@@ -575,53 +572,37 @@ impl<'a> finite_wasm::wasmparser::VisitOperator<'a> for GasCostCfg {
 }
 
 impl crate::runner::VM for NearVM {
-    fn run(
-        &self,
-        method_name: &str,
-        ext: &mut dyn External,
+    fn prepare(
+        self: Box<Self>,
+        ext: &dyn External,
         context: &VMContext,
-        fees_config: Arc<RuntimeFeesConfig>,
-        promise_results: Arc<[PromiseResult]>,
         cache: Option<&dyn ContractRuntimeCache>,
-    ) -> Result<VMOutcome, VMRunnerError> {
+    ) -> Box<dyn crate::PreparedContract> {
         let cache = cache.unwrap_or(&NoContractRuntimeCache);
-        self.with_compiled_and_loaded(
-            cache,
-            ext,
-            context,
-            method_name,
-            |result_state, ext, artifact| {
+        let prepd =
+            self.with_compiled_and_loaded(cache, ext, context, |result_state, artifact, vm| {
                 let memory = NearVmMemory::new(
-                    self.config.limit_config.initial_memory_pages,
-                    self.config.limit_config.max_memory_pages,
+                    vm.config.limit_config.initial_memory_pages,
+                    vm.config.limit_config.max_memory_pages,
                 )
                 .expect("Cannot create memory for a contract call");
-                // FIXME: this mostly duplicates the `run_module` method.
-                // Note that we don't clone the actual backing memory, just increase the RC.
-                let vmmemory = memory.vm();
-                let mut logic =
-                    VMLogic::new(ext, context, fees_config, promise_results, result_state, memory);
-                let import = build_imports(
-                    vmmemory,
-                    &mut logic,
-                    Arc::clone(&self.config),
-                    artifact.engine(),
-                );
-                let entrypoint = match get_entrypoint_index(&*artifact, method_name) {
+                let entrypoint = match get_entrypoint_index(&*artifact, &context.method) {
                     Ok(index) => index,
                     Err(e) => {
-                        return Ok(VMOutcome::abort_but_nop_outcome_in_old_protocol(
-                            logic.result_state,
-                            e,
+                        return Ok(PreparedContract::Outcome(
+                            VMOutcome::abort_but_nop_outcome_in_old_protocol(result_state, e),
                         ))
                     }
                 };
-                match self.run_method(&artifact, import, entrypoint)? {
-                    Ok(()) => Ok(VMOutcome::ok(logic.result_state)),
-                    Err(err) => Ok(VMOutcome::abort(logic.result_state, err)),
-                }
-            },
-        )
+                Ok(PreparedContract::Ready(ReadyContract {
+                    memory,
+                    result_state,
+                    entrypoint,
+                    artifact: Arc::clone(artifact),
+                    vm,
+                }))
+            });
+        Box::new(prepd)
     }
 
     fn precompile(
@@ -638,6 +619,42 @@ impl crate::runner::VM for NearVM {
     }
 }
 
+struct ReadyContract {
+    memory: NearVmMemory,
+    result_state: ExecutionResultState,
+    entrypoint: FunctionIndex,
+    artifact: VMArtifact,
+    vm: Box<NearVM>,
+}
+
+#[allow(clippy::large_enum_variant)]
+enum PreparedContract {
+    Outcome(VMOutcome),
+    Ready(ReadyContract),
+}
+
+impl crate::PreparedContract for VMResult<PreparedContract> {
+    fn run(
+        self: Box<Self>,
+        ext: &mut dyn External,
+        context: &VMContext,
+        fees_config: Arc<RuntimeFeesConfig>,
+    ) -> VMResult {
+        let ReadyContract { memory, result_state, entrypoint, artifact, vm } = match (*self)? {
+            PreparedContract::Outcome(outcome) => return Ok(outcome),
+            PreparedContract::Ready(r) => r,
+        };
+        let config = Arc::clone(&result_state.config);
+        let vmmemory = memory.vm();
+        let mut logic = VMLogic::new(ext, context, fees_config, result_state, memory);
+        let import = build_imports(vmmemory, &mut logic, config, artifact.engine());
+        match vm.run_method(&artifact, import, entrypoint)? {
+            Ok(()) => Ok(VMOutcome::ok(logic.result_state)),
+            Err(err) => Ok(VMOutcome::abort(logic.result_state, err)),
+        }
+    }
+}
+
 pub(crate) struct NearVmImports<'engine, 'vmlogic, 'vmlogic_refs> {
     pub(crate) memory: VMMemory,
     config: Arc<Config>,