wasmtime: add caching of executables and artifacts (#11532)

This largely mirrors the code in near_vm_runner module. I heard some
people pondering what it would be like to use a higher quality backend.
Outside LLVM, Cranelift is by far the next in line in produced code
quality. Since we already have wasmtime in place, might as well wire it
up completely for a full experience.

Based on top of #11529
Part of #11319

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

@@ -110,6 +110,12 @@ pub enum CompilationError {
     WasmerCompileError {
         msg: String,
     },
+    /// This is for defense in depth.
+    /// We expect our runtime-independent preparation code to fully catch all invalid wasms,
+    /// but, if it ever misses something we’ll emit this error
+    WasmtimeCompileError {
+        msg: String,
+    },
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, BorshDeserialize, BorshSerialize)]
@@ -342,6 +348,9 @@ impl fmt::Display for CompilationError {
             CompilationError::WasmerCompileError { msg } => {
                 write!(f, "Wasmer compilation error: {}", msg)
             }
+            CompilationError::WasmtimeCompileError { msg } => {
+                write!(f, "Wasmtime compilation error: {}", msg)
+            }
         }
     }
 }

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

@@ -173,7 +173,7 @@ impl NearVM {
             .engine
             .compile_universal(&prepared_code, &self)
             .map_err(|err| {
-                tracing::error!(?err, "near_vm failed to compile the prepared code (this is defense-in-depth, the error was recovered from but should be reported to pagoda)");
+                tracing::error!(?err, "near_vm failed to compile the prepared code (this is defense-in-depth, the error was recovered from but should be reported to the developers)");
                 CompilationError::WasmerCompileError { msg: err.to_string() }
             })?;
         crate::metrics::compilation_duration(VMKind::NearVm, start.elapsed());

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

@@ -1,12 +1,16 @@
 use crate::errors::ContractPrecompilatonResult;
 use crate::logic::errors::{
-    CompilationError, FunctionCallError, MethodResolveError, PrepareError, VMLogicError,
-    VMRunnerError, WasmTrap,
+    CacheError, CompilationError, FunctionCallError, MethodResolveError, PrepareError,
+    VMLogicError, VMRunnerError, WasmTrap,
 };
 use crate::logic::types::PromiseResult;
 use crate::logic::Config;
 use crate::logic::{External, MemSlice, MemoryLike, VMContext, VMLogic, VMOutcome};
-use crate::{imports, prepare, ContractCode, ContractRuntimeCache};
+use crate::runner::VMResult;
+use crate::{
+    get_contract_cache_key, imports, prepare, CompiledContract, CompiledContractInfo, ContractCode,
+    ContractRuntimeCache, NoContractRuntimeCache,
+};
 use near_parameters::vm::VMKind;
 use near_parameters::RuntimeFeesConfig;
 use std::borrow::Cow;
@@ -121,49 +125,131 @@ impl IntoVMError for anyhow::Error {
 }
 
 #[allow(clippy::needless_pass_by_ref_mut)]
-pub fn get_engine(config: &mut wasmtime::Config) -> Engine {
+pub fn get_engine(config: &wasmtime::Config) -> Engine {
     Engine::new(config).unwrap()
 }
 
+pub(crate) fn default_wasmtime_config(config: &Config) -> wasmtime::Config {
+    let features =
+        crate::features::WasmFeatures::from(config.limit_config.contract_prepare_version);
+    let mut config = wasmtime::Config::from(features);
+    config.max_wasm_stack(1024 * 1024 * 1024); // wasm stack metering is implemented by instrumentation, we don't want wasmtime to trap before that
+    config
+}
+
 pub(crate) fn wasmtime_vm_hash() -> u64 {
     // TODO: take into account compiler and engine used to compile the contract.
     64
 }
 
 pub(crate) struct WasmtimeVM {
     config: Config,
+    engine: wasmtime::Engine,
 }
 
 impl WasmtimeVM {
     pub(crate) fn new(config: Config) -> Self {
-        Self { config }
+        Self { engine: get_engine(&default_wasmtime_config(&config)), config }
     }
 
-    pub(crate) fn default_wasmtime_config(&self) -> wasmtime::Config {
-        let features =
-            crate::features::WasmFeatures::from(self.config.limit_config.contract_prepare_version);
-        let mut config = wasmtime::Config::from(features);
-        config.max_wasm_stack(1024 * 1024 * 1024); // wasm stack metering is implemented by instrumentation, we don't want wasmtime to trap before that
-        config
+    #[tracing::instrument(target = "vm", level = "debug", "WasmtimeVM::compile_uncached", skip_all)]
+    fn compile_uncached(&self, code: &ContractCode) -> Result<Vec<u8>, CompilationError> {
+        let start = std::time::Instant::now();
+        let prepared_code = prepare::prepare_contract(code.code(), &self.config, VMKind::Wasmtime)
+            .map_err(CompilationError::PrepareError)?;
+        let serialized = self.engine.precompile_module(&prepared_code).map_err(|err| {
+            tracing::error!(?err, "wasmtime failed to compile the prepared code (this is defense-in-depth, the error was recovered from but should be reported to the developers)");
+            CompilationError::WasmtimeCompileError { msg: err.to_string() }
+        });
+        crate::metrics::compilation_duration(VMKind::Wasmtime, start.elapsed());
+        serialized
     }
-}
 
-impl crate::runner::VM for WasmtimeVM {
-    fn run(
+    fn compile_and_cache(
         &self,
-        method_name: &str,
+        code: &ContractCode,
+        cache: &dyn ContractRuntimeCache,
+    ) -> Result<Result<Vec<u8>, CompilationError>, CacheError> {
+        let serialized_or_error = self.compile_uncached(code);
+        let key = get_contract_cache_key(*code.hash(), &self.config);
+        let record = CompiledContractInfo {
+            wasm_bytes: code.code().len() as u64,
+            compiled: match &serialized_or_error {
+                Ok(serialized) => CompiledContract::Code(serialized.clone()),
+                Err(err) => CompiledContract::CompileModuleError(err.clone()),
+            },
+        };
+        cache.put(&key, record).map_err(CacheError::WriteError)?;
+        Ok(serialized_or_error)
+    }
+
+    fn with_compiled_and_loaded(
+        &self,
+        cache: &dyn ContractRuntimeCache,
         ext: &mut dyn External,
         context: &VMContext,
         fees_config: &RuntimeFeesConfig,
         promise_results: &[PromiseResult],
-        _cache: Option<&dyn ContractRuntimeCache>,
-    ) -> Result<VMOutcome, VMRunnerError> {
-        let Some(code) = ext.get_contract() else {
-            return Err(VMRunnerError::ContractCodeNotPresent);
-        };
-        let mut config = self.default_wasmtime_config();
-        let engine = get_engine(&mut config);
-        let mut store = Store::new(&engine, ());
+        method_name: &str,
+        closure: impl FnOnce(VMLogic, Memory, Store<()>, Module) -> Result<VMOutcome, VMRunnerError>,
+    ) -> VMResult<VMOutcome> {
+        let code_hash = ext.code_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(
+            code_hash,
+            || {
+                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 {
+                        return Err(VMRunnerError::ContractCodeNotPresent);
+                    };
+                    return Ok(to_any((
+                        code.code().len() as u64,
+                        match self.compile_and_cache(&code, cache)? {
+                            Ok(serialized_module) => Ok(unsafe {
+                                Module::deserialize(&self.engine, serialized_module)
+                                    .map_err(|err| VMRunnerError::LoadingError(err.to_string()))?
+                            }),
+                            Err(err) => Err(err),
+                        },
+                    )));
+                };
+                match &compiled_contract_info.compiled {
+                    CompiledContract::CompileModuleError(err) => Ok::<_, VMRunnerError>(to_any((
+                        compiled_contract_info.wasm_bytes,
+                        Err(err.clone()),
+                    ))),
+                    CompiledContract::Code(serialized_module) => {
+                        unsafe {
+                            // (UN-)SAFETY: the `serialized_module` must have been produced by
+                            // a prior call to `serialize`.
+                            //
+                            // In practice this is not necessarily true. One could have
+                            // forgotten to change the cache key when upgrading the version of
+                            // the near_vm library or the database could have had its data
+                            // corrupted while at rest.
+                            //
+                            // There should definitely be some validation in near_vm to ensure
+                            // we load what we think we load.
+                            let module = Module::deserialize(&self.engine, &serialized_module)
+                                .map_err(|err| VMRunnerError::LoadingError(err.to_string()))?;
+                            Ok(to_any((compiled_contract_info.wasm_bytes, Ok(module))))
+                        }
+                    }
+                }
+            },
+            move |value| {
+                let &(wasm_bytes, ref downcast) = value
+                    .downcast_ref::<MemoryCacheType>()
+                    .expect("downcast should always succeed");
+
+                (wasm_bytes, downcast.clone())
+            },
+        )?;
+
+        let mut store = Store::new(&self.engine, ());
         let mut memory = WasmtimeMemory::new(
             &mut store,
             self.config.limit_config.initial_memory_pages,
@@ -173,83 +259,109 @@ impl crate::runner::VM for WasmtimeVM {
         let memory_copy = memory.0;
         let mut logic =
             VMLogic::new(ext, context, &self.config, fees_config, promise_results, &mut memory);
-
-        let result = logic.before_loading_executable(method_name, code.code().len() as u64);
+        let result = logic.before_loading_executable(method_name, wasm_bytes);
         if let Err(e) = result {
             return Ok(VMOutcome::abort(logic, e));
         }
-
-        let prepared_code =
-            match prepare::prepare_contract(code.code(), &self.config, VMKind::Wasmtime) {
-                Ok(code) => code,
-                Err(err) => return Ok(VMOutcome::abort(logic, FunctionCallError::from(err))),
-            };
-        let start = std::time::Instant::now();
-        let module = match Module::new(&engine, prepared_code) {
-            Ok(module) => module,
-            Err(err) => return Ok(VMOutcome::abort(logic, err.into_vm_error()?)),
-        };
-        crate::metrics::compilation_duration(VMKind::Wasmtime, start.elapsed());
-        let mut linker = Linker::new(&engine);
-
-        let result = logic.after_loading_executable(code.code().len() as u64);
-        if let Err(e) = result {
-            return Ok(VMOutcome::abort(logic, e));
-        }
-        link(&mut linker, memory_copy, &store, &mut logic);
-        match module.get_export(method_name) {
-            Some(export) => match export {
-                Func(func_type) => {
-                    if func_type.params().len() != 0 || func_type.results().len() != 0 {
-                        let err = FunctionCallError::MethodResolveError(
-                            MethodResolveError::MethodInvalidSignature,
-                        );
-                        return Ok(VMOutcome::abort_but_nop_outcome_in_old_protocol(logic, err));
-                    }
+        match module_result {
+            Ok(module) => {
+                let result = logic.after_loading_executable(wasm_bytes);
+                if let Err(e) = result {
+                    return Ok(VMOutcome::abort(logic, e));
                 }
-                _ => {
-                    return Ok(VMOutcome::abort_but_nop_outcome_in_old_protocol(
-                        logic,
-                        FunctionCallError::MethodResolveError(MethodResolveError::MethodNotFound),
-                    ));
-                }
-            },
-            None => {
-                return Ok(VMOutcome::abort_but_nop_outcome_in_old_protocol(
-                    logic,
-                    FunctionCallError::MethodResolveError(MethodResolveError::MethodNotFound),
-                ));
+                closure(logic, memory_copy, store, module)
             }
+            Err(e) => Ok(VMOutcome::abort(logic, FunctionCallError::CompilationError(e))),
         }
-        match linker.instantiate(&mut store, &module) {
-            Ok(instance) => match instance.get_func(&mut store, method_name) {
-                Some(func) => match func.typed::<(), ()>(&mut store) {
-                    Ok(run) => match run.call(&mut store, ()) {
-                        Ok(_) => Ok(VMOutcome::ok(logic)),
-                        Err(err) => Ok(VMOutcome::abort(logic, err.into_vm_error()?)),
+    }
+}
+
+impl crate::runner::VM for WasmtimeVM {
+    fn run(
+        &self,
+        method_name: &str,
+        ext: &mut dyn External,
+        context: &VMContext,
+        fees_config: &RuntimeFeesConfig,
+        promise_results: &[PromiseResult],
+        cache: Option<&dyn ContractRuntimeCache>,
+    ) -> Result<VMOutcome, VMRunnerError> {
+        let cache = cache.unwrap_or(&NoContractRuntimeCache);
+        self.with_compiled_and_loaded(
+            cache,
+            ext,
+            context,
+            fees_config,
+            promise_results,
+            method_name,
+            |mut logic, memory, mut store, module| {
+                let mut linker = Linker::new(&(&self.engine));
+                link(&mut linker, memory, &store, &mut logic);
+                match module.get_export(method_name) {
+                    Some(export) => match export {
+                        Func(func_type) => {
+                            if func_type.params().len() != 0 || func_type.results().len() != 0 {
+                                let err = FunctionCallError::MethodResolveError(
+                                    MethodResolveError::MethodInvalidSignature,
+                                );
+                                return Ok(VMOutcome::abort_but_nop_outcome_in_old_protocol(
+                                    logic, err,
+                                ));
+                            }
+                        }
+                        _ => {
+                            return Ok(VMOutcome::abort_but_nop_outcome_in_old_protocol(
+                                logic,
+                                FunctionCallError::MethodResolveError(
+                                    MethodResolveError::MethodNotFound,
+                                ),
+                            ));
+                        }
+                    },
+                    None => {
+                        return Ok(VMOutcome::abort_but_nop_outcome_in_old_protocol(
+                            logic,
+                            FunctionCallError::MethodResolveError(
+                                MethodResolveError::MethodNotFound,
+                            ),
+                        ));
+                    }
+                }
+                match linker.instantiate(&mut store, &module) {
+                    Ok(instance) => match instance.get_func(&mut store, method_name) {
+                        Some(func) => match func.typed::<(), ()>(&mut store) {
+                            Ok(run) => match run.call(&mut store, ()) {
+                                Ok(_) => Ok(VMOutcome::ok(logic)),
+                                Err(err) => Ok(VMOutcome::abort(logic, err.into_vm_error()?)),
+                            },
+                            Err(err) => Ok(VMOutcome::abort(logic, err.into_vm_error()?)),
+                        },
+                        None => {
+                            return Ok(VMOutcome::abort_but_nop_outcome_in_old_protocol(
+                                logic,
+                                FunctionCallError::MethodResolveError(
+                                    MethodResolveError::MethodNotFound,
+                                ),
+                            ));
+                        }
                     },
                     Err(err) => Ok(VMOutcome::abort(logic, err.into_vm_error()?)),
-                },
-                None => {
-                    return Ok(VMOutcome::abort_but_nop_outcome_in_old_protocol(
-                        logic,
-                        FunctionCallError::MethodResolveError(MethodResolveError::MethodNotFound),
-                    ));
                 }
             },
-            Err(err) => Ok(VMOutcome::abort(logic, err.into_vm_error()?)),
-        }
+        )
     }
 
     fn precompile(
         &self,
-        _code: &ContractCode,
-        _cache: &dyn ContractRuntimeCache,
+        code: &ContractCode,
+        cache: &dyn ContractRuntimeCache,
     ) -> Result<
         Result<ContractPrecompilatonResult, CompilationError>,
         crate::logic::errors::CacheError,
     > {
-        Ok(Ok(ContractPrecompilatonResult::CacheNotAvailable))
+        Ok(self
+            .compile_and_cache(code, cache)?
+            .map(|_| ContractPrecompilatonResult::ContractCompiled))
     }
 }
 

runtime/runtime/src/conversions.rs

@@ -49,6 +49,9 @@ mod compilation_error {
                 },
                 From::PrepareError(pe) => Self::PrepareError(super::Convert::convert(pe)),
                 From::WasmerCompileError { msg } => Self::WasmerCompileError { msg },
+                // Intentionally converting into "Wasmer" error here in order to avoid
+                // this particular detail being visible to the protocol unnecessarily.
+                From::WasmtimeCompileError { msg } => Self::WasmerCompileError { msg },
             }
         }
     }