Use cache in runners

lib/cli/src/commands/run_unstable.rs

@@ -20,7 +20,8 @@ use tempfile::NamedTempFile;
 use url::Url;
 use wapm_targz_to_pirita::FileMap;
 use wasmer::{
-    DeserializeError, Function, Imports, Instance, Module, Store, Type, TypedFunction, Value,
+    DeserializeError, Engine, Function, Imports, Instance, Module, Store, Type, TypedFunction,
+    Value,
 };
 use wasmer_cache::Cache;
 use wasmer_compiler::ArtifactBuild;
@@ -71,10 +72,12 @@ impl RunUnstable {
 
         let (mut store, _) = self.store.get_store()?;
 
-        let mut cache = self.wasmer_home.module_cache();
-        let result = match target.load(&mut cache, &store)? {
+        let cache = self.wasmer_home.module_cache();
+        let result = match target.load(cache, &store)? {
             ExecutableTarget::WebAssembly(wasm) => self.execute_wasm(&target, &wasm, &mut store),
-            ExecutableTarget::Webc(container) => self.execute_webc(&target, &container, &mut store),
+            ExecutableTarget::Webc(container, cache) => {
+                self.execute_webc(&target, container, cache, &mut store)
+            }
         };
 
         if let Err(e) = &result {
@@ -111,7 +114,8 @@ impl RunUnstable {
     fn execute_webc(
         &self,
         target: &TargetOnDisk,
-        container: &WapmContainer,
+        container: WapmContainer,
+        mut cache: ModuleCache,
         store: &mut Store,
     ) -> Result<(), Error> {
         let id = match self.entrypoint.as_deref() {
@@ -129,41 +133,54 @@ impl RunUnstable {
         // without needing to go through and instantiate each one.
 
         let (store, _compiler_type) = self.store.get_store()?;
-        let mut runner = wasmer_wasix::runners::wasi::WasiRunner::new(store)
-            .with_args(self.args.clone())
-            .with_envs(self.wasi.env_vars.clone())
-            .with_mapped_directories(self.wasi.mapped_dirs.clone());
-        if self.wasi.forward_host_env {
-            runner.set_forward_host_env();
-        }
-        if runner.can_run_command(id, command).unwrap_or(false) {
-            return runner.run_cmd(container, id).context("WASI runner failed");
-        }
-
-        let (store, _compiler_type) = self.store.get_store()?;
-        let mut runner = wasmer_wasix::runners::emscripten::EmscriptenRunner::new(store);
-        runner.set_args(self.args.clone());
-        if runner.can_run_command(id, command).unwrap_or(false) {
-            return runner
-                .run_cmd(container, id)
-                .context("Emscripten runner failed");
-        }
+        match command.runner.as_ref() {
+            webc::metadata::annotations::EMSCRIPTEN_RUNNER_URI => {
+                let mut runner = wasmer_wasix::runners::emscripten::EmscriptenRunner::new(store);
+                runner.set_args(self.args.clone());
+                if runner.can_run_command(id, command).unwrap_or(false) {
+                    return runner
+                        .run_cmd(&container, id)
+                        .context("Emscripten runner failed");
+                }
+            }
+            webc::metadata::annotations::WCGI_RUNNER_URI => {
+                let mut runner = wasmer_wasix::runners::wcgi::WcgiRunner::new(id).with_compile(
+                    move |engine, bytes| {
+                        compile_wasm_cached("".to_string(), bytes, &mut cache, engine)
+                    },
+                );
 
-        let mut runner = wasmer_wasix::runners::wcgi::WcgiRunner::new(id);
-        let (store, _compiler_type) = self.store.get_store()?;
-        runner
-            .config()
-            .args(self.args.clone())
-            .store(store)
-            .addr(self.wcgi.addr)
-            .envs(self.wasi.env_vars.clone())
-            .map_directories(self.wasi.mapped_dirs.clone())
-            .callbacks(Callbacks::new(self.wcgi.addr));
-        if self.wasi.forward_host_env {
-            runner.config().forward_host_env();
-        }
-        if runner.can_run_command(id, command).unwrap_or(false) {
-            return runner.run_cmd(container, id).context("WCGI runner failed");
+                runner
+                    .config()
+                    .args(self.args.clone())
+                    .store(store)
+                    .addr(self.wcgi.addr)
+                    .envs(self.wasi.env_vars.clone())
+                    .map_directories(self.wasi.mapped_dirs.clone())
+                    .callbacks(Callbacks::new(self.wcgi.addr));
+                if self.wasi.forward_host_env {
+                    runner.config().forward_host_env();
+                }
+                if runner.can_run_command(id, command).unwrap_or(false) {
+                    return runner.run_cmd(&container, id).context("WCGI runner failed");
+                }
+            }
+            webc::metadata::annotations::WASI_RUNNER_URI => {
+                let mut runner = wasmer_wasix::runners::wasi::WasiRunner::new(store)
+                    .with_compile(move |engine, bytes| {
+                        compile_wasm_cached("".to_string(), bytes, &mut cache, engine)
+                    })
+                    .with_args(self.args.clone())
+                    .with_envs(self.wasi.env_vars.clone())
+                    .with_mapped_directories(self.wasi.mapped_dirs.clone());
+                if self.wasi.forward_host_env {
+                    runner.set_forward_host_env();
+                }
+                if runner.can_run_command(id, command).unwrap_or(false) {
+                    return runner.run_cmd(&container, id).context("WASI runner failed");
+                }
+            }
+            _ => {}
         }
 
         anyhow::bail!(
@@ -436,12 +453,12 @@ impl TargetOnDisk {
         }
     }
 
-    fn load(&self, cache: &mut ModuleCache, store: &Store) -> Result<ExecutableTarget, Error> {
+    fn load(&self, mut cache: ModuleCache, store: &Store) -> Result<ExecutableTarget, Error> {
         match self {
             TargetOnDisk::Webc(webc) => {
                 // As an optimisation, try to use the mmapped version first.
                 if let Ok(container) = WapmContainer::from_path(webc.clone()) {
-                    return Ok(ExecutableTarget::Webc(container));
+                    return Ok(ExecutableTarget::Webc(container, cache));
                 }
 
                 // Otherwise, fall back to the version that reads everything
@@ -450,7 +467,7 @@ impl TargetOnDisk {
                     .with_context(|| format!("Unable to read \"{}\"", webc.display()))?;
                 let container = WapmContainer::from_bytes(bytes.into())?;
 
-                Ok(ExecutableTarget::Webc(container))
+                Ok(ExecutableTarget::Webc(container, cache))
             }
             TargetOnDisk::Directory(dir) => {
                 // FIXME: Runners should be able to load directories directly
@@ -461,12 +478,17 @@ impl TargetOnDisk {
                 let container = WapmContainer::from_bytes(webc.into())
                     .context("Unable to parse the generated WEBC file")?;
 
-                Ok(ExecutableTarget::Webc(container))
+                Ok(ExecutableTarget::Webc(container, cache))
             }
             TargetOnDisk::WebAssemblyBinary(path) => {
                 let wasm = std::fs::read(path)
                     .with_context(|| format!("Unable to read \"{}\"", path.display()))?;
-                let module = compile_wasm_cached(path, &wasm, cache, store)?;
+                let module = compile_wasm_cached(
+                    path.display().to_string(),
+                    &wasm,
+                    &mut cache,
+                    store.engine(),
+                )?;
                 Ok(ExecutableTarget::WebAssembly(module))
             }
             TargetOnDisk::Wat(path) => {
@@ -475,7 +497,12 @@ impl TargetOnDisk {
                 let wasm =
                     wasmer::wat2wasm(&wat).context("Unable to convert the WAT to WebAssembly")?;
 
-                let module = compile_wasm_cached(path, &wasm, cache, store)?;
+                let module = compile_wasm_cached(
+                    path.display().to_string(),
+                    &wasm,
+                    &mut cache,
+                    store.engine(),
+                )?;
                 Ok(ExecutableTarget::WebAssembly(module))
             }
             TargetOnDisk::Artifact(artifact) => {
@@ -490,38 +517,41 @@ impl TargetOnDisk {
 }
 
 fn compile_wasm_cached(
-    path: &Path,
+    name: String,
     wasm: &[u8],
     cache: &mut ModuleCache,
-    store: &Store,
+    engine: &Engine,
 ) -> Result<Module, Error> {
+    tracing::debug!("Trying to retrieve module from cache");
+
     let hash = wasmer_cache::Hash::generate(wasm);
+    tracing::debug!("Generated hash: {}", hash);
 
     unsafe {
-        match cache.load(store, hash) {
+        match cache.load(engine, hash) {
             Ok(m) => {
-                tracing::debug!(%hash, "Reusing a cached module");
+                tracing::debug!(%hash, "Module loaded from cache");
                 return Ok(m);
             }
             Err(DeserializeError::Io(e)) if e.kind() == ErrorKind::NotFound => {}
             Err(error) => {
                 tracing::warn!(
                     %hash,
                     error=&error as &dyn std::error::Error,
-                    wasm_path=%path.display(),
+                    name=%name,
                     "Unable to deserialize the cached module",
                 );
             }
         }
     }
 
-    let mut module = Module::new(store, wasm).context("Unable to load the module from a file")?;
-    module.set_name(path.display().to_string().as_str());
+    let mut module = Module::new(engine, wasm).context("Unable to load the module from a file")?;
+    module.set_name(&name);
 
     if let Err(e) = cache.store(hash, &module) {
         tracing::warn!(
             error=&e as &dyn std::error::Error,
-            wat=%path.display(),
+            wat=%name,
             key=%hash,
             "Unable to cache the compiled module",
         );
@@ -533,7 +563,7 @@ fn compile_wasm_cached(
 #[derive(Debug, Clone)]
 enum ExecutableTarget {
     WebAssembly(Module),
-    Webc(WapmContainer),
+    Webc(WapmContainer, ModuleCache),
 }
 
 fn generate_coredump(err: &Error, source: &Path, coredump_path: &Path) -> Result<(), Error> {

lib/wasi/Cargo.toml

@@ -58,7 +58,7 @@ pin-project = "1.0.12"
 # Used by the WCGI runner
 hyper = { version = "0.14", features = ["server", "stream"], optional = true }
 wcgi = { version = "0.1.1", optional = true }
-wcgi-host = { version = "0.1.0", optional = true }
+wcgi-host = { version = "0.1.1", 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 }
 

lib/wasi/src/runners/mod.rs

@@ -12,7 +12,7 @@ pub mod wcgi;
 
 pub use self::{
     container::{Bindings, WapmContainer, WebcParseError, WitBindings},
-    runner::Runner,
+    runner::{CompileModule, Runner},
 };
 
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]

lib/wasi/src/runners/runner.rs

@@ -1,7 +1,10 @@
+use crate::runners::WapmContainer;
 use anyhow::Error;
+use wasmer::{Engine, Module};
 use webc::metadata::Command;
 
-use crate::runners::WapmContainer;
+/// A compile module function
+pub type CompileModule = dyn FnMut(&Engine, &[u8]) -> Result<Module, Error>;
 
 /// Trait that all runners have to implement
 pub trait Runner {

lib/wasi/src/runners/wasi.rs

@@ -4,21 +4,23 @@ use std::sync::Arc;
 
 use anyhow::{Context, Error};
 use serde::{Deserialize, Serialize};
-use wasmer::{Module, Store};
+use wasmer::{Engine, Module, Store};
 use webc::metadata::{annotations::Wasi, Command};
 
 use crate::{
-    runners::{wasi_common::CommonWasiOptions, MappedDirectory, WapmContainer},
+    runners::{wasi_common::CommonWasiOptions, CompileModule, MappedDirectory, WapmContainer},
     PluggableRuntime, VirtualTaskManager, WasiEnvBuilder,
 };
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Serialize, Deserialize)]
 pub struct WasiRunner {
     wasi: CommonWasiOptions,
     #[serde(skip, default)]
     store: Store,
     #[serde(skip, default)]
     pub(crate) tasks: Option<Arc<dyn VirtualTaskManager>>,
+    #[serde(skip, default)]
+    compile: Option<Box<CompileModule>>,
 }
 
 impl WasiRunner {
@@ -28,9 +30,19 @@ impl WasiRunner {
             store,
             wasi: CommonWasiOptions::default(),
             tasks: None,
+            compile: None,
         }
     }
 
+    /// Sets the compile function
+    pub fn with_compile(
+        mut self,
+        compile: impl FnMut(&Engine, &[u8]) -> Result<Module, Error> + 'static,
+    ) -> Self {
+        self.compile = Some(Box::new(compile));
+        self
+    }
+
     /// Returns the current arguments for this `WasiRunner`
     pub fn get_args(&self) -> Vec<String> {
         self.wasi.args.clone()