feat(wasix): Improve module loading performance and API correctness

examples/wasi.rs

@@ -43,7 +43,9 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     println!("Compiling module...");
     // Let's compile the Wasm module.
-    let module = runtime.load_module_sync(&wasm_bytes[..])?;
+    let data =
+        wasmer_wasix::runtime::module_cache::HashedModuleData::new_sha256(wasm_bytes.clone());
+    let module = runtime.load_hashed_module_sync(data)?;
 
     // Create a pipe for the module's stdout.
     let (stdout_tx, mut stdout_rx) = Pipe::channel();

lib/package/src/utils.rs

@@ -62,6 +62,11 @@ pub fn from_disk(path: impl AsRef<Path>) -> Result<Container, WasmerPackageError
     }
 }
 
+/// Check if the data looks like a webc.
+pub fn is_container(bytes: &[u8]) -> bool {
+    is_tarball(std::io::Cursor::new(bytes)) || webc::detect(bytes).is_ok()
+}
+
 pub fn from_bytes(bytes: impl Into<Bytes>) -> Result<Container, WasmerPackageError> {
     let bytes: Bytes = bytes.into();
 

lib/virtual-fs/src/lib.rs

@@ -412,6 +412,16 @@ pub trait VirtualFile:
         })
     }
 
+    /// Get the full contents of this file as an [`OwnedBuffer`].
+    ///
+    /// **NOTE**: Only implement this if the file is already available in-memory
+    /// and can be cloned cheaply!
+    ///
+    /// Allows consumers to do zero-copy cloning of the underlying data.
+    fn as_owned_buffer(&self) -> Option<OwnedBuffer> {
+        None
+    }
+
     /// Polls the file for when there is data to be read
     fn poll_read_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>>;
 

lib/virtual-fs/src/mem_fs/file.rs

@@ -507,6 +507,17 @@ impl VirtualFile for FileHandle {
         self.cursor = cursor;
         Ok(())
     }
+
+    fn as_owned_buffer(&self) -> Option<OwnedBuffer> {
+        let fs = self.filesystem.inner.read().ok()?;
+
+        let inode = fs.storage.get(self.inode)?;
+        match inode {
+            Node::ReadOnlyFile(f) => Some(f.file.buffer.clone()),
+            Node::CustomFile(f) => f.file.lock().ok()?.as_owned_buffer(),
+            _ => None,
+        }
+    }
 }
 
 #[cfg(test)]

lib/virtual-fs/src/webc_volume_fs.rs

@@ -251,6 +251,10 @@ impl VirtualFile for File {
     ) -> Poll<std::io::Result<usize>> {
         Poll::Ready(Err(std::io::ErrorKind::PermissionDenied.into()))
     }
+
+    fn as_owned_buffer(&self) -> Option<SharedBytes> {
+        Some(self.content.get_ref().clone())
+    }
 }
 
 impl AsyncRead for File {

lib/wasix/src/bin_factory/exec.rs

@@ -6,6 +6,7 @@ use crate::{
         TaskJoinHandle,
     },
     runtime::{
+        module_cache::HashedModuleData,
         task_manager::{
             TaskWasm, TaskWasmRecycle, TaskWasmRecycleProperties, TaskWasmRunProperties,
         },
@@ -43,7 +44,7 @@ pub async fn spawn_exec(
 
 #[tracing::instrument(level = "trace", skip_all, fields(%name))]
 pub async fn spawn_exec_wasm(
-    wasm: &[u8],
+    wasm: HashedModuleData,
     name: &str,
     env: WasiEnv,
     runtime: &Arc<dyn Runtime + Send + Sync + 'static>,
@@ -85,10 +86,10 @@ pub fn package_command_by_name<'a>(
 
 pub async fn spawn_load_module(
     name: &str,
-    wasm: &[u8],
+    wasm: HashedModuleData,
     runtime: &Arc<dyn Runtime + Send + Sync + 'static>,
 ) -> Result<Module, SpawnError> {
-    match runtime.load_module(wasm).await {
+    match runtime.load_hashed_module(wasm).await {
         Ok(module) => Ok(module),
         Err(err) => {
             tracing::error!(

lib/wasix/src/bin_factory/mod.rs

@@ -8,6 +8,7 @@ use std::{
 };
 
 use anyhow::Context;
+use shared_buffer::OwnedBuffer;
 use virtual_fs::{AsyncReadExt, FileSystem};
 use wasmer::FunctionEnvMut;
 use wasmer_package::utils::from_bytes;
@@ -24,6 +25,7 @@ pub use self::{
 };
 use crate::{
     os::{command::Commands, task::TaskJoinHandle},
+    runtime::module_cache::HashedModuleData,
     Runtime, SpawnError, WasiEnv,
 };
 
@@ -84,7 +86,8 @@ impl BinFactory {
             // Execute
             match executable {
                 Executable::Wasm(bytes) => {
-                    spawn_exec_wasm(&bytes, name.as_str(), env, &self.runtime).await
+                    let data = HashedModuleData::new_xxhash(bytes.clone());
+                    spawn_exec_wasm(data, name.as_str(), env, &self.runtime).await
                 }
                 Executable::BinaryPackage(pkg) => {
                     // Get the command that is going to be executed
@@ -170,7 +173,7 @@ impl BinFactory {
 }
 
 pub enum Executable {
-    Wasm(bytes::Bytes),
+    Wasm(OwnedBuffer),
     BinaryPackage(BinaryPackage),
 }
 
@@ -185,18 +188,35 @@ async fn load_executable_from_filesystem(
         .open(path)
         .context("Unable to open the file")?;
 
-    let mut data = Vec::with_capacity(f.size() as usize);
-    f.read_to_end(&mut data).await.context("Read failed")?;
+    // Fast path if the file is fully available in memory.
+    // Prevents redundant copying of the file data.
+    if let Some(buf) = f.as_owned_buffer() {
+        if wasmer_package::utils::is_container(buf.as_slice()) {
+            let bytes = buf.clone().into_bytes();
+            if let Ok(container) = from_bytes(bytes.clone()) {
+                let pkg = BinaryPackage::from_webc(&container, rt)
+                    .await
+                    .context("Unable to load the package")?;
+
+                return Ok(Executable::BinaryPackage(pkg));
+            }
+        }
 
-    let bytes: bytes::Bytes = data.into();
+        Ok(Executable::Wasm(buf))
+    } else {
+        let mut data = Vec::with_capacity(f.size() as usize);
+        f.read_to_end(&mut data).await.context("Read failed")?;
 
-    if let Ok(container) = from_bytes(bytes.clone()) {
-        let pkg = BinaryPackage::from_webc(&container, rt)
-            .await
-            .context("Unable to load the package")?;
+        let bytes: bytes::Bytes = data.into();
 
-        Ok(Executable::BinaryPackage(pkg))
-    } else {
-        Ok(Executable::Wasm(bytes))
+        if let Ok(container) = from_bytes(bytes.clone()) {
+            let pkg = BinaryPackage::from_webc(&container, rt)
+                .await
+                .context("Unable to load the package")?;
+
+            Ok(Executable::BinaryPackage(pkg))
+        } else {
+            Ok(Executable::Wasm(OwnedBuffer::from_bytes(bytes)))
+        }
     }
 }

lib/wasix/src/os/command/builtins/cmd_wasmer.rs

@@ -3,9 +3,10 @@ use std::{any::Any, path::PathBuf, sync::Arc};
 use crate::{
     bin_factory::spawn_exec_wasm,
     os::task::{OwnedTaskStatus, TaskJoinHandle},
-    runtime::task_manager::InlineWaker,
+    runtime::{module_cache::HashedModuleData, task_manager::InlineWaker},
     SpawnError,
 };
+use shared_buffer::OwnedBuffer;
 use virtual_fs::{AsyncReadExt, FileSystem};
 use wasmer::FunctionEnvMut;
 use wasmer_package::utils::from_bytes;
@@ -50,6 +51,12 @@ impl CmdWasmer {
     }
 }
 
+#[derive(Debug, Clone)]
+enum Executable {
+    Wasm(OwnedBuffer),
+    BinaryPackage(BinaryPackage),
+}
+
 impl CmdWasmer {
     async fn run(
         &self,
@@ -59,11 +66,6 @@ impl CmdWasmer {
         what: Option<String>,
         mut args: Vec<String>,
     ) -> Result<TaskJoinHandle, SpawnError> {
-        pub enum Executable {
-            Wasm(bytes::Bytes),
-            BinaryPackage(BinaryPackage),
-        }
-
         // If the first argument is a '--' then skip it
         if args.first().map(|a| a.as_str()) == Some("--") {
             args = args.into_iter().skip(1).collect();
@@ -102,7 +104,7 @@ impl CmdWasmer {
 
                     Executable::BinaryPackage(pkg)
                 } else {
-                    Executable::Wasm(bytes)
+                    Executable::Wasm(OwnedBuffer::from_bytes(bytes))
                 }
             } else if let Ok(pkg) = self.get_package(&what).await {
                 Executable::BinaryPackage(pkg)
@@ -136,7 +138,10 @@ impl CmdWasmer {
                     // Now run the module
                     spawn_exec(binary, name, env, &self.runtime).await
                 }
-                Executable::Wasm(bytes) => spawn_exec_wasm(&bytes, name, env, &self.runtime).await,
+                Executable::Wasm(bytes) => {
+                    let data = HashedModuleData::new_xxhash(bytes);
+                    spawn_exec_wasm(data, name, env, &self.runtime).await
+                }
             }
         } else {
             let _ = unsafe { stderr_write(parent_ctx, HELP_RUN.as_bytes()) }.await;