Split the spawn functions so they can be invoked separately in a more performant way

Author: john-sharratt

lib/wasix/src/bin_factory/exec.rs

@@ -30,6 +30,24 @@ pub async fn spawn_exec(
     env: WasiEnv,
     runtime: &Arc<dyn Runtime + Send + Sync + 'static>,
 ) -> Result<TaskJoinHandle, SpawnError> {
+    // Load the WASM
+    let wasm = spawn_load_wasm(&env, &binary, name).await?;
+
+    // Load the module
+    let module = spawn_load_module(&env, name, wasm, runtime).await?;
+
+    // Spawn union the file system
+    spawn_union_fs(&env, &binary).await?;
+
+    // Now run the module
+    spawn_exec_module(module, env, runtime)
+}
+
+pub async fn spawn_load_wasm<'a>(
+    env: &WasiEnv,
+    binary: &'a BinaryPackage,
+    name: &str,
+) -> Result<&'a [u8], SpawnError> {
     let wasm = if let Some(cmd) = binary.get_command(name) {
         cmd.atom.as_ref()
     } else if let Some(wasm) = binary.entrypoint_bytes() {
@@ -44,7 +62,15 @@ pub async fn spawn_exec(
         env.on_exit(Some(Errno::Noexec.into())).await;
         return Err(SpawnError::CompileError);
     };
+    Ok(wasm)
+}
 
+pub async fn spawn_load_module(
+    env: &WasiEnv,
+    name: &str,
+    wasm: &[u8],
+    runtime: &Arc<dyn Runtime + Send + Sync + 'static>,
+) -> Result<Module, SpawnError> {
     let module = match runtime.load_module(wasm).await {
         Ok(module) => module,
         Err(err) => {
@@ -57,7 +83,10 @@ pub async fn spawn_exec(
             return Err(SpawnError::CompileError);
         }
     };
+    Ok(module)
+}
 
+pub async fn spawn_union_fs(env: &WasiEnv, binary: &BinaryPackage) -> Result<(), SpawnError> {
     // If the file system has not already been union'ed then do so
     env.state
         .fs
@@ -68,9 +97,7 @@ pub async fn spawn_exec(
             SpawnError::FileSystemError
         })?;
     tracing::debug!("{:?}", env.state.fs);
-
-    // Now run the module
-    spawn_exec_module(module, env, runtime)
+    Ok(())
 }
 
 pub fn spawn_exec_module(

lib/wasix/src/bin_factory/mod.rs

@@ -14,7 +14,9 @@ mod exec;
 
 pub use self::{
     binary_package::*,
-    exec::{run_exec, spawn_exec, spawn_exec_module},
+    exec::{
+        run_exec, spawn_exec, spawn_exec_module, spawn_load_module, spawn_load_wasm, spawn_union_fs,
+    },
 };
 use crate::{os::command::Commands, Runtime};
 

lib/wasix/src/runtime/mod.rs

@@ -104,8 +104,9 @@ where
     fn load_module<'a>(&'a self, wasm: &'a [u8]) -> BoxFuture<'a, anyhow::Result<Module>> {
         let engine = self.engine();
         let module_cache = self.module_cache();
+        let hash = ModuleHash::hash(wasm);
 
-        let task = async move { load_module(&engine, &module_cache, wasm).await };
+        let task = async move { load_module(&engine, &module_cache, wasm, hash).await };
 
         Box::pin(task)
     }
@@ -150,16 +151,16 @@ pub async fn load_module(
     engine: &wasmer::Engine,
     module_cache: &(dyn ModuleCache + Send + Sync),
     wasm: &[u8],
+    wasm_hash: ModuleHash,
 ) -> Result<Module, anyhow::Error> {
-    let hash = ModuleHash::hash(wasm);
-    let result = module_cache.load(hash, engine).await;
+    let result = module_cache.load(wasm_hash, engine).await;
 
     match result {
         Ok(module) => return Ok(module),
         Err(CacheError::NotFound) => {}
         Err(other) => {
             tracing::warn!(
-                %hash,
+                %wasm_hash,
                 error=&other as &dyn std::error::Error,
                 "Unable to load the cached module",
             );
@@ -168,9 +169,9 @@ pub async fn load_module(
 
     let module = Module::new(&engine, wasm)?;
 
-    if let Err(e) = module_cache.save(hash, engine, &module).await {
+    if let Err(e) = module_cache.save(wasm_hash, engine, &module).await {
         tracing::warn!(
-            %hash,
+            %wasm_hash,
             error=&e as &dyn std::error::Error,
             "Unable to cache the compiled module",
         );
@@ -549,7 +550,9 @@ impl Runtime for OverriddenRuntime {
         if self.engine.is_some() || self.module_cache.is_some() {
             let engine = self.engine();
             let module_cache = self.module_cache();
-            let task = async move { load_module(&engine, &module_cache, wasm).await };
+            let hash = ModuleHash::hash(wasm);
+
+            let task = async move { load_module(&engine, &module_cache, wasm, hash).await };
             Box::pin(task)
         } else {
             self.inner.load_module(wasm)