Introduce a module cache abstraction

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/wasi/Cargo.toml

@@ -61,6 +61,7 @@ tower-http = { version = "0.4.0", features = ["trace", "util", "catch-panic", "c
 tower = { version = "0.4.13", features = ["make", "util"], optional = true }
 url = "2.3.1"
 semver = "1.0.17"
+dashmap = "5.4.0"
 
 [target.'cfg(not(target_arch = "riscv64"))'.dependencies.reqwest]
 version = "0.11"

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::Key;
 
 #[derive(Derivative, Clone)]
 #[derivative(Debug)]
 pub struct BinaryPackageCommand {
     name: String,
     #[derivative(Debug = "ignore")]
     pub(crate) atom: SharedBytes,
-    hash: OnceCell<String>,
+    hash: OnceCell<Key>,
 }
 
 impl BinaryPackageCommand {
@@ -38,8 +38,8 @@ impl BinaryPackageCommand {
         &self.atom
     }
 
-    pub fn hash(&self) -> &str {
-        self.hash.get_or_init(|| hash_of_binary(self.atom()))
+    pub fn hash(&self) -> &Key {
+        self.hash.get_or_init(|| Key::sha256(self.atom()))
     }
 }
 
@@ -54,7 +54,7 @@ pub struct BinaryPackage {
     pub when_cached: Option<u128>,
     #[derivative(Debug = "ignore")]
     pub entry: Option<SharedBytes>,
-    pub hash: Arc<Mutex<Option<String>>>,
+    pub hash: OnceCell<Key>,
     pub webc_fs: Option<Arc<dyn FileSystem + Send + Sync + 'static>>,
     pub commands: Arc<RwLock<Vec<BinaryPackageCommand>>>,
     pub uses: Vec<String>,
@@ -64,15 +64,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) -> Key {
+        *self.hash.get_or_init(|| {
             if let Some(entry) = self.entry.as_ref() {
-                hash.replace(hash_of_binary(entry.as_ref()));
+                Key::sha256(entry)
             } else {
-                hash.replace(hash_of_binary(&self.package_name));
+                Key::sha256(self.package_name.as_bytes())
             }
-        }
-        hash.as_ref().unwrap().clone()
+        })
     }
 }

lib/wasi/src/bin_factory/exec.rs

@@ -2,31 +2,36 @@ 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 module = compiled_modules.get_compiled_module(&store, binary.hash().as_str(), compiler);
+    let key = binary.hash().combined_with(compiler);
+
+    let compiled_modules = runtime.module_cache();
+    let module = compiled_modules.load(key, store.engine()).await.ok();
+
     let module = match (module, binary.entry.as_ref()) {
         (Some(a), _) => a,
         (None, Some(entry)) => {
@@ -43,7 +48,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, &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 +218,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[..])
-}