Merge pull request #3841 from wasmerio/module-cache

Introduce a module cache abstraction

lib/api/src/externals/function.rs

@@ -60,7 +60,7 @@ mod private {
     //! Sealing the HostFunctionKind because it shouldn't be implemented
     //! by any type outside.
     //! See:
-    //! https://rust-lang.github.io/api-guidelines/future-proofing.html#c-sealed
+    //! <https://rust-lang.github.io/api-guidelines/future-proofing.html#c-sealed>
     pub trait HostFunctionKindSealed {}
     impl HostFunctionKindSealed for super::WithEnv {}
     impl HostFunctionKindSealed for super::WithoutEnv {}

lib/cli/src/commands/run/wasi.rs

@@ -14,6 +14,7 @@ use wasmer_wasix::{
     os::{tty_sys::SysTty, TtyBridge},
     runners::MappedDirectory,
     runtime::{
+        module_cache::{FileSystemCache, ModuleCache},
         resolver::{PackageResolver, RegistryResolver},
         task_manager::tokio::TokioTaskManager,
     },
@@ -220,12 +221,17 @@ impl Wasi {
             rt.set_tty(tty);
         }
 
-        rt.set_engine(Some(engine));
+        let wasmer_home = WasmerConfig::get_wasmer_dir().map_err(anyhow::Error::msg)?;
 
         let resolver = self
-            .prepare_resolver()
+            .prepare_resolver(&wasmer_home)
             .context("Unable to prepare the package resolver")?;
-        rt.set_resolver(resolver);
+        let module_cache = wasmer_wasix::runtime::module_cache::in_memory()
+            .and_then(FileSystemCache::new(wasmer_home.join("compiled")));
+
+        rt.set_resolver(resolver)
+            .set_module_cache(module_cache)
+            .set_engine(Some(engine));
 
         Ok(rt)
     }
@@ -273,8 +279,8 @@ impl Wasi {
         })
     }
 
-    fn prepare_resolver(&self) -> Result<impl PackageResolver> {
-        let mut resolver = wapm_resolver()?;
+    fn prepare_resolver(&self, wasmer_home: &Path) -> Result<impl PackageResolver> {
+        let mut resolver = wapm_resolver(wasmer_home)?;
 
         for path in &self.include_webcs {
             let pkg = preload_webc(path)
@@ -286,14 +292,13 @@ impl Wasi {
     }
 }
 
-fn wapm_resolver() -> Result<RegistryResolver, anyhow::Error> {
+fn wapm_resolver(wasmer_home: &Path) -> Result<RegistryResolver, anyhow::Error> {
     // FIXME(Michael-F-Bryan): Ideally, all of this would in the
     // RegistryResolver::from_env() constructor, but we don't want to add
     // wasmer-registry as a dependency of wasmer-wasix just yet.
-    let wasmer_home = WasmerConfig::get_wasmer_dir().map_err(anyhow::Error::msg)?;
-    let cache_dir = wasmer_registry::get_webc_dir(&wasmer_home);
+    let cache_dir = wasmer_registry::get_webc_dir(wasmer_home);
     let config =
-        wasmer_registry::WasmerConfig::from_file(&wasmer_home).map_err(anyhow::Error::msg)?;
+        wasmer_registry::WasmerConfig::from_file(wasmer_home).map_err(anyhow::Error::msg)?;
 
     let registry = config.registry.get_graphql_url();
     let registry = registry

lib/cli/src/commands/run_unstable.rs

@@ -454,17 +454,24 @@ impl TargetOnDisk {
         let leading_bytes = &buffer[..bytes_read];
 
         if wasmer::is_wasm(leading_bytes) {
-            Ok(TargetOnDisk::WebAssemblyBinary(path))
-        } else if webc::detect(leading_bytes).is_ok() {
-            Ok(TargetOnDisk::Webc(path))
-        } else if path.extension() == Some("wat".as_ref()) {
-            Ok(TargetOnDisk::Wat(path))
-        } else {
-            #[cfg(feature = "compiler")]
-            if ArtifactBuild::is_deserializable(leading_bytes) {
-                return Ok(TargetOnDisk::Artifact(path));
-            }
-            anyhow::bail!("Unable to determine how to execute \"{}\"", path.display());
+            return Ok(TargetOnDisk::WebAssemblyBinary(path));
+        }
+
+        if webc::detect(leading_bytes).is_ok() {
+            return Ok(TargetOnDisk::Webc(path));
+        }
+
+        #[cfg(feature = "compiler")]
+        if ArtifactBuild::is_deserializable(leading_bytes) {
+            return Ok(TargetOnDisk::Artifact(path));
+        }
+
+        // If we can't figure out the file type based on its content, fall back
+        // to checking the extension.
+
+        match path.extension().and_then(|s| s.to_str()) {
+            Some("wat") => Ok(TargetOnDisk::Wat(path)),
+            _ => anyhow::bail!("Unable to determine how to execute \"{}\"", path.display()),
         }
     }
 

lib/compiler/src/engine/inner.rs

@@ -220,7 +220,7 @@ impl Engine {
     ///
     /// # Safety
     ///
-    /// See [`crate::Module::deserialize_from_file`].
+    /// See [`Artifact::deserialize`].
     #[cfg(not(target_arch = "wasm32"))]
     pub unsafe fn deserialize_from_file(
         &self,

lib/types/src/serialize.rs

@@ -207,7 +207,7 @@ pub struct MetadataHeader {
 impl MetadataHeader {
     /// Current ABI version. Increment this any time breaking changes are made
     /// to the format of the serialized data.
-    const CURRENT_VERSION: u32 = 4;
+    pub const CURRENT_VERSION: u32 = 4;
 
     /// Magic number to identify wasmer metadata.
     const MAGIC: [u8; 8] = *b"WASMER\0\0";

lib/types/src/stack/frame.rs

@@ -1,13 +1,9 @@
 use crate::SourceLoc;
 
-/// Description of a frame in a backtrace for a [`RuntimeError`](crate::RuntimeError).
+/// Description of a frame in a backtrace.
 ///
-/// Whenever a WebAssembly trap occurs an instance of [`RuntimeError`]
-/// is created. Each [`RuntimeError`] has a backtrace of the
-/// WebAssembly frames that led to the trap, and each frame is
-/// described by this structure.
-///
-/// [`RuntimeError`]: crate::RuntimeError
+/// Each runtime error includes a backtrace of the WebAssembly frames that led
+/// to the trap, and each frame is described by this structure.
 #[derive(Debug, Clone)]
 pub struct FrameInfo {
     /// The name of the module

lib/vm/src/instance/mod.rs

@@ -989,7 +989,7 @@ impl Drop for VMInstance {
 }
 
 impl VMInstance {
-    /// Create a new `VMInstance` pointing at a new [`InstanceRef`].
+    /// Create a new `VMInstance` pointing at a new [`Instance`].
     ///
     /// # Safety
     ///

lib/vm/src/vmcontext.rs

@@ -272,7 +272,7 @@ pub struct VMGlobalImport {
 /// # Safety
 /// This data is safe to share between threads because it's plain data that
 /// is the user's responsibility to synchronize. Additionally, all operations
-/// on `from` are thread-safe through the use of a mutex in [`Global`].
+/// on `from` are thread-safe through the use of a mutex in [`VMGlobal`].
 unsafe impl Send for VMGlobalImport {}
 /// # Safety
 /// This data is safe to share between threads because it's plain data that

lib/wasi/Cargo.toml

@@ -53,14 +53,16 @@ wai-bindgen-wasmer = { path = "../wai-bindgen-wasmer", version = "0.4.0", featur
 heapless = "0.7.16"
 once_cell = "1.17.0"
 pin-project = "1.0.12"
+semver = "1.0.17"
+dashmap = "5.4.0"
+tempfile = "3.4.0"
 # Used by the WCGI runner
 hyper = { version = "0.14", features = ["server", "stream"], optional = true }
 wcgi = { version = "0.1.2", optional = true }
 wcgi-host = { version = "0.1.2", optional = true }
 tower-http = { version = "0.4.0", features = ["trace", "util", "catch-panic", "cors"], optional = true }
 tower = { version = "0.4.13", features = ["make", "util"], optional = true }
 url = "2.3.1"
-semver = "1.0.17"
 
 [target.'cfg(not(target_arch = "riscv64"))'.dependencies.reqwest]
 version = "0.11"
@@ -89,7 +91,6 @@ wasm-bindgen = "0.2.74"
 [dev-dependencies]
 wasmer = { path = "../api", version = "=3.3.0", default-features = false, features = ["wat", "js-serializable-module"] }
 tokio = { version = "1", features = [ "sync", "macros", "rt" ], default_features = false }
-tempfile = "3.4.0"
 
 [target.'cfg(target_arch = "wasm32")'.dev-dependencies]
 wasm-bindgen-test = "0.3.0"

lib/wasi/src/bin_factory/binary_package.rs

@@ -1,20 +1,20 @@
-use std::sync::{Arc, Mutex, RwLock};
+use std::sync::{Arc, RwLock};
 
 use derivative::*;
 use once_cell::sync::OnceCell;
 use semver::Version;
 use virtual_fs::FileSystem;
 use webc::compat::SharedBytes;
 
-use super::hash_of_binary;
+use crate::runtime::module_cache::ModuleHash;
 
 #[derive(Derivative, Clone)]
 #[derivative(Debug)]
 pub struct BinaryPackageCommand {
     name: String,
     #[derivative(Debug = "ignore")]
     pub(crate) atom: SharedBytes,
-    hash: OnceCell<String>,
+    hash: OnceCell<ModuleHash>,
 }
 
 impl BinaryPackageCommand {
@@ -38,23 +38,24 @@ impl BinaryPackageCommand {
         &self.atom
     }
 
-    pub fn hash(&self) -> &str {
-        self.hash.get_or_init(|| hash_of_binary(self.atom()))
+    pub fn hash(&self) -> &ModuleHash {
+        self.hash.get_or_init(|| ModuleHash::sha256(self.atom()))
     }
 }
 
 /// A WebAssembly package that has been loaded into memory.
 ///
-/// You can crate a [`BinaryPackage`] using [`crate::bin_factory::ModuleCache`]
-/// or [`crate::wapm::parse_static_webc()`].
+/// You can crate a [`BinaryPackage`] using a
+/// [`crate::runtime::resolver::PackageResolver`] or
+/// [`crate::wapm::parse_static_webc()`].
 #[derive(Derivative, Clone)]
 #[derivative(Debug)]
 pub struct BinaryPackage {
     pub package_name: String,
     pub when_cached: Option<u128>,
     #[derivative(Debug = "ignore")]
     pub entry: Option<SharedBytes>,
-    pub hash: Arc<Mutex<Option<String>>>,
+    pub hash: OnceCell<ModuleHash>,
     pub webc_fs: Option<Arc<dyn FileSystem + Send + Sync + 'static>>,
     pub commands: Arc<RwLock<Vec<BinaryPackageCommand>>>,
     pub uses: Vec<String>,
@@ -64,15 +65,13 @@ pub struct BinaryPackage {
 }
 
 impl BinaryPackage {
-    pub fn hash(&self) -> String {
-        let mut hash = self.hash.lock().unwrap();
-        if hash.is_none() {
+    pub fn hash(&self) -> ModuleHash {
+        *self.hash.get_or_init(|| {
             if let Some(entry) = self.entry.as_ref() {
-                hash.replace(hash_of_binary(entry.as_ref()));
+                ModuleHash::sha256(entry)
             } else {
-                hash.replace(hash_of_binary(&self.package_name));
+                ModuleHash::sha256(self.package_name.as_bytes())
             }
-        }
-        hash.as_ref().unwrap().clone()
+        })
     }
 }

lib/wasi/src/bin_factory/exec.rs

@@ -2,31 +2,33 @@ use std::{pin::Pin, sync::Arc};
 
 use crate::{
     os::task::{thread::WasiThreadRunGuard, TaskJoinHandle},
+    runtime::module_cache::ModuleCache,
     VirtualBusError, WasiRuntimeError,
 };
 use futures::Future;
 use tracing::*;
 use wasmer::{FunctionEnvMut, Instance, Memory, Module, Store};
 use wasmer_wasix_types::wasi::Errno;
 
-use super::{BinFactory, BinaryPackage, ModuleCache};
+use super::{BinFactory, BinaryPackage};
 use crate::{
     import_object_for_all_wasi_versions, runtime::SpawnType, SpawnedMemory, WasiEnv,
     WasiFunctionEnv, WasiRuntime,
 };
 
-pub fn spawn_exec(
+#[tracing::instrument(level = "trace", skip_all, fields(%name, %binary.package_name))]
+pub async fn spawn_exec(
     binary: BinaryPackage,
     name: &str,
     store: Store,
     env: WasiEnv,
     runtime: &Arc<dyn WasiRuntime + Send + Sync + 'static>,
-    compiled_modules: &ModuleCache,
 ) -> Result<TaskJoinHandle, VirtualBusError> {
-    // The deterministic id for this engine
-    let compiler = store.engine().deterministic_id();
+    let key = binary.hash();
+
+    let compiled_modules = runtime.module_cache();
+    let module = compiled_modules.load(key, store.engine()).await.ok();
 
-    let module = compiled_modules.get_compiled_module(&store, binary.hash().as_str(), compiler);
     let module = match (module, binary.entry.as_ref()) {
         (Some(a), _) => a,
         (None, Some(entry)) => {
@@ -43,7 +45,15 @@ pub fn spawn_exec(
                 env.blocking_cleanup(Some(Errno::Noexec.into()));
             }
             let module = module?;
-            compiled_modules.set_compiled_module(binary.hash().as_str(), compiler, &module);
+
+            if let Err(e) = compiled_modules.save(key, store.engine(), &module).await {
+                tracing::debug!(
+                    %key,
+                    package_name=%binary.package_name,
+                    error=&e as &dyn std::error::Error,
+                    "Unable to save the compiled module",
+                );
+            }
             module
         }
         (None, None) => {
@@ -205,14 +215,7 @@ impl BinFactory {
             let binary = binary?;
 
             // Execute
-            spawn_exec(
-                binary,
-                name.as_str(),
-                store,
-                env,
-                &self.runtime,
-                &self.cache,
-            )
+            spawn_exec(binary, name.as_str(), store, env, &self.runtime).await
         })
     }
 

lib/wasi/src/bin_factory/mod.rs

@@ -10,34 +10,25 @@ use virtual_fs::{AsyncReadExt, FileSystem};
 
 mod binary_package;
 mod exec;
-mod module_cache;
-
-use sha2::*;
 
 pub use self::{
     binary_package::*,
     exec::{spawn_exec, spawn_exec_module},
-    module_cache::ModuleCache,
 };
 use crate::{os::command::Commands, WasiRuntime};
 
 #[derive(Debug, Clone)]
 pub struct BinFactory {
     pub(crate) commands: Commands,
     runtime: Arc<dyn WasiRuntime + Send + Sync + 'static>,
-    pub(crate) cache: Arc<ModuleCache>,
     pub(crate) local: Arc<RwLock<HashMap<String, Option<BinaryPackage>>>>,
 }
 
 impl BinFactory {
-    pub fn new(
-        compiled_modules: Arc<ModuleCache>,
-        runtime: Arc<dyn WasiRuntime + Send + Sync + 'static>,
-    ) -> BinFactory {
+    pub fn new(runtime: Arc<dyn WasiRuntime + Send + Sync + 'static>) -> BinFactory {
         BinFactory {
-            commands: Commands::new_with_builtins(runtime.clone(), compiled_modules.clone()),
+            commands: Commands::new_with_builtins(runtime.clone()),
             runtime,
-            cache: compiled_modules,
             local: Arc::new(RwLock::new(HashMap::new())),
         }
     }
@@ -118,10 +109,3 @@ async fn load_package_from_filesystem(
 
     Ok(pkg)
 }
-
-pub fn hash_of_binary(data: impl AsRef<[u8]>) -> String {
-    let mut hasher = Sha256::default();
-    hasher.update(data.as_ref());
-    let hash = hasher.finalize();
-    hex::encode(&hash[..])
-}