Store Compiled Modules when using OCI Wasm Images

crates/containerd-shim-wasm/src/container/engine.rs

@@ -7,16 +7,9 @@ use super::Source;
 use crate::container::{PathResolve, RuntimeContext};
 use crate::sandbox::Stdio;
 
-/// RuntimeInfo contains the name and version of the runtime that is running
-#[derive(Clone, Debug, Default)]
-pub struct RuntimeInfo {
-    pub name: &'static str,
-    pub version: &'static str,
-}
-
 pub trait Engine: Clone + Send + Sync + 'static {
     /// The name to use for this engine
-    fn info() -> &'static RuntimeInfo;
+    fn name() -> &'static str;
 
     /// Run a WebAssembly container
     fn run_wasi(&self, ctx: &impl RuntimeContext, stdio: Stdio) -> Result<i32>;
@@ -68,8 +61,18 @@ pub trait Engine: Clone + Send + Sync + 'static {
         bail!("precompilation not supported for this runtime")
     }
 
-    /// Precompile a module
-    fn can_precompile() -> bool {
-        false
+    /// Can_precompile lets the shim know if the runtime supports precompilation.
+    /// When it returns Some(unique_string) the `unique_string` will be used as a cache key for the precompiled module.
+    ///
+    /// `unique_string` should at least include the version of the shim running but could include other information such as a hash
+    /// of the version and cpu type and other important information in the validation of being able to use precompiled module.  
+    /// If the string doesn't match then the module will be recompiled and cached with the new `unique_string`.
+    ///
+    /// This string will be used in the following way:
+    /// "runwasi.io/precompiled/<Engine.name()>/<unique_string>"
+    ///
+    /// When it returns None the runtime will not be asked to precompile the module.  This is the default value.
+    fn can_precompile(&self) -> Option<String> {
+        None
     }
 }

crates/containerd-shim-wasm/src/container/mod.rs

@@ -17,7 +17,7 @@ mod wasm;
 
 pub(crate) use context::WasiContext;
 pub use context::{Entrypoint, RuntimeContext, Source};
-pub use engine::{Engine, RuntimeInfo};
+pub use engine::Engine;
 pub use instance::Instance;
 pub use path::PathResolve;
 pub use wasm::WasmBinaryType;

crates/containerd-shim-wasm/src/container/tests.rs

@@ -1,6 +1,5 @@
 use anyhow::bail;
 
-use super::engine::RuntimeInfo;
 use crate::container::{Engine, RuntimeContext, Stdio};
 use crate::sys::container::instance::Instance;
 use crate::testing::WasiTest;
@@ -9,11 +8,8 @@ use crate::testing::WasiTest;
 struct EngineFailingValidation;
 
 impl Engine for EngineFailingValidation {
-    fn info() -> &'static RuntimeInfo {
-        &RuntimeInfo {
-            name: "wasi_instance",
-            version: "0.0.0",
-        }
+    fn name() -> &'static str {
+        "wasi_instance"
     }
     fn can_handle(&self, _ctx: &impl RuntimeContext) -> anyhow::Result<()> {
         bail!("can't handle");

crates/containerd-shim-wasm/src/sandbox/containerd/client.rs

@@ -25,7 +25,7 @@ use tokio_stream::wrappers::ReceiverStream;
 use tonic::{Code, Request};
 
 use super::lease::LeaseGuard;
-use crate::container::{Engine, RuntimeInfo};
+use crate::container::Engine;
 use crate::sandbox::error::{Error as ShimError, Result};
 use crate::sandbox::oci::{self, WasmLayer};
 use crate::with_lease;
@@ -139,10 +139,10 @@ impl Client {
         &self,
         data: Vec<u8>,
         original_digest: String,
-        info: &RuntimeInfo,
+        label: &str,
     ) -> Result<WriteContent> {
         let expected = format!("sha256:{}", digest(data.clone()));
-        let reference = format!("precompile-{}-{}-{}", info.name, info.version, expected);
+        let reference = format!("precompile-{}", label);
         let lease = self.lease(reference.clone())?;
 
         let digest = self.rt.block_on(async {
@@ -195,10 +195,9 @@ impl Client {
             // by returning the offset of the content that was found.
             let data_to_write = data[response.offset as usize..].to_vec();
 
-            let label = precompile_label(info);
             // Write and commit at same time
             let mut labels = HashMap::new();
-            labels.insert(label, original_digest.clone());
+            labels.insert(label.to_string(), original_digest.clone());
             let commit_request = WriteContentRequest {
                 action: WriteAction::Commit.into(),
                 total: len,
@@ -405,10 +404,14 @@ impl Client {
         log::info!("found manifest with WASM OCI image format.");
         // This label is unique across runtimes and version of the shim running
         // a precompiled component/module will not work across different runtimes or versions
-        let label = precompile_label(T::info());
-        match image.labels.get(&label) {
-            Some(precompile_digest) if T::can_precompile() => {
-                log::info!("found precompiled label: {} ", &label);
+        let (can_precompile, precompile_id) = match engine.can_precompile() {
+            Some(precompile_id) => (true, precompile_label(T::name(), &precompile_id)),
+            None => (false, "".to_string()),
+        };
+
+        match image.labels.get(&precompile_id) {
+            Some(precompile_digest) if can_precompile => {
+                log::info!("found precompiled label: {} ", &precompile_id);
                 match self.read_content(precompile_digest) {
                     Ok(precompiled) => {
                         log::info!("found precompiled module in cache: {} ", &precompile_digest);
@@ -441,17 +444,17 @@ impl Client {
             return Ok((vec![], platform));
         }
 
-        if T::can_precompile() {
+        if can_precompile {
             log::info!("precompiling module");
             let precompiled = engine.precompile(layers.as_slice())?;
             log::info!("precompiling module: {}", image_digest.clone());
             let precompiled_content =
-                self.save_content(precompiled.clone(), image_digest.clone(), T::info())?;
+                self.save_content(precompiled.clone(), image_digest.clone(), &precompile_id)?;
 
             log::debug!("updating image with compiled content digest");
             image
                 .labels
-                .insert(label, precompiled_content.digest.clone());
+                .insert(precompile_id, precompiled_content.digest.clone());
             self.update_image(image)?;
 
             // The original image is considered a root object, by adding a ref to the new compiled content
@@ -486,8 +489,8 @@ impl Client {
     }
 }
 
-fn precompile_label(info: &RuntimeInfo) -> String {
-    format!("{}/{}/{}", PRECOMPILE_PREFIX, info.name, info.version)
+fn precompile_label(name: &str, version: &str) -> String {
+    format!("{}/{}/{}", PRECOMPILE_PREFIX, name, version)
 }
 
 fn is_wasm_layer(media_type: &MediaType, supported_layer_types: &[&str]) -> bool {
@@ -510,30 +513,17 @@ mod tests {
         let expected = digest(data.clone());
         let expected = format!("sha256:{}", expected);
 
+        let label = precompile_label("test", "hasdfh");
         let returned = client
-            .save_content(
-                data,
-                "original".to_string(),
-                &RuntimeInfo {
-                    name: "test",
-                    version: "0.0.0",
-                },
-            )
+            .save_content(data, "original".to_string(), &label)
             .unwrap();
         assert_eq!(expected, returned.digest.clone());
 
         let data = client.read_content(returned.digest.clone()).unwrap();
         assert_eq!(data, b"hello world");
 
         client
-            .save_content(
-                data.clone(),
-                "original".to_string(),
-                &RuntimeInfo {
-                    name: "test",
-                    version: "0.0.0",
-                },
-            )
+            .save_content(data.clone(), "original".to_string(), &label)
             .expect_err("Should not be able to save when lease is open");
 
         // need to drop the lease to be able to create a second one
@@ -542,14 +532,7 @@ mod tests {
 
         // a second call should be successful since it already exists
         let returned = client
-            .save_content(
-                data,
-                "original".to_string(),
-                &RuntimeInfo {
-                    name: "test",
-                    version: "0.0.0",
-                },
-            )
+            .save_content(data, "original".to_string(), &label)
             .unwrap();
         assert_eq!(expected, returned.digest);
 

crates/containerd-shim-wasm/src/sys/unix/container/executor.rs

@@ -36,7 +36,7 @@ impl<E: Engine> LibcontainerExecutor for Executor<E> {
     fn validate(&self, spec: &Spec) -> Result<(), ExecutorValidationError> {
         // We can handle linux container. We delegate wasm container to the engine.
         match self.inner(spec) {
-            InnerExecutor::CantHandle => Err(ExecutorValidationError::CantHandle(E::info().name)),
+            InnerExecutor::CantHandle => Err(ExecutorValidationError::CantHandle(E::name())),
             _ => Ok(()),
         }
     }
@@ -45,7 +45,7 @@ impl<E: Engine> LibcontainerExecutor for Executor<E> {
         // If it looks like a linux container, run it as a linux container.
         // Otherwise, run it as a wasm container
         match self.inner(spec) {
-            InnerExecutor::CantHandle => Err(LibcontainerExecutorError::CantHandle(E::info().name)),
+            InnerExecutor::CantHandle => Err(LibcontainerExecutorError::CantHandle(E::name())),
             InnerExecutor::Linux => {
                 log::info!("executing linux container");
                 self.stdio.take().redirect().unwrap();

crates/containerd-shim-wasm/src/sys/unix/container/instance.rs

@@ -39,7 +39,7 @@ impl<E: Engine> SandboxInstance for Instance<E> {
         let engine = cfg.get_engine();
         let bundle = cfg.get_bundle().to_path_buf();
         let namespace = cfg.get_namespace();
-        let rootdir = Path::new(DEFAULT_CONTAINER_ROOT_DIR).join(E::info().name);
+        let rootdir = Path::new(DEFAULT_CONTAINER_ROOT_DIR).join(E::name());
         let rootdir = determine_rootdir(&bundle, &namespace, rootdir)?;
         let stdio = Stdio::init_from_cfg(cfg)?;
 

crates/containerd-shim-wasm/src/testing.rs

@@ -354,7 +354,7 @@ pub mod oci_helpers {
     pub fn get_image_label() -> Result<(String, String)> {
         let mut grep = Command::new("grep")
             .arg("-ohE")
-            .arg("runwasi.io/precompiled/wasmtime/0.3.1=.*")
+            .arg("runwasi.io/precompiled/.*")
             .stdout(Stdio::piped())
             .stdin(Stdio::piped())
             .spawn()?;

crates/containerd-shim-wasmedge/src/instance.rs

@@ -1,7 +1,5 @@
 use anyhow::{Context, Result};
-use containerd_shim_wasm::container::{
-    Engine, Entrypoint, Instance, RuntimeContext, RuntimeInfo, Stdio,
-};
+use containerd_shim_wasm::container::{Engine, Entrypoint, Instance, RuntimeContext, Stdio};
 use wasmedge_sdk::config::{ConfigBuilder, HostRegistrationConfigOptions};
 use wasmedge_sdk::plugin::PluginManager;
 use wasmedge_sdk::VmBuilder;
@@ -27,11 +25,8 @@ impl Default for WasmEdgeEngine {
 }
 
 impl Engine for WasmEdgeEngine {
-    fn info() -> &'static RuntimeInfo {
-        &RuntimeInfo {
-            name: "wasmedge",
-            version: env!("CARGO_PKG_VERSION"),
-        }
+    fn name() -> &'static str {
+        "wasmedge"
     }
 
     fn run_wasi(&self, ctx: &impl RuntimeContext, stdio: Stdio) -> Result<i32> {

crates/containerd-shim-wasmer/src/instance.rs

@@ -1,7 +1,5 @@
 use anyhow::Result;
-use containerd_shim_wasm::container::{
-    Engine, Entrypoint, Instance, RuntimeContext, RuntimeInfo, Stdio,
-};
+use containerd_shim_wasm::container::{Engine, Entrypoint, Instance, RuntimeContext, Stdio};
 use wasmer::{Module, Store};
 use wasmer_wasix::virtual_fs::host_fs::FileSystem;
 use wasmer_wasix::{WasiEnv, WasiError};
@@ -14,11 +12,8 @@ pub struct WasmerEngine {
 }
 
 impl Engine for WasmerEngine {
-    fn info() -> &'static RuntimeInfo {
-        &RuntimeInfo {
-            name: "wasmer",
-            version: env!("CARGO_PKG_VERSION"),
-        }
+    fn name() -> &'static str {
+        "wasmer"
     }
 
     fn run_wasi(&self, ctx: &impl RuntimeContext, stdio: Stdio) -> Result<i32> {

crates/containerd-shim-wasmtime/src/instance.rs

@@ -1,32 +1,49 @@
+use std::collections::hash_map::DefaultHasher;
 use std::fs::File;
-use std::sync::OnceLock;
+use std::hash::{Hash, Hasher};
+use std::marker::PhantomData;
 
 use anyhow::{bail, Context, Result};
 use containerd_shim_wasm::container::{
-    Engine, Entrypoint, Instance, RuntimeContext, RuntimeInfo, Stdio, WasmBinaryType,
+    Engine, Entrypoint, Instance, RuntimeContext, Stdio, WasmBinaryType,
 };
 use wasi_common::I32Exit;
 use wasmtime::component::{self as wasmtime_component, Component, ResourceTable};
-use wasmtime::{Module, Precompiled, Store};
+use wasmtime::{Config, Module, Precompiled, Store};
 use wasmtime_wasi::preview2::{self as wasi_preview2};
 use wasmtime_wasi::{self as wasi_preview1, Dir};
 
-pub type WasmtimeInstance = Instance<WasmtimeEngine>;
+pub type WasmtimeInstance = Instance<WasmtimeEngine<DefaultConfig>>;
 
 #[derive(Clone)]
-pub struct WasmtimeEngine {
+pub struct WasmtimeEngine<T: WasiConfig + Clone + Send + Sync> {
     engine: wasmtime::Engine,
+    config_type: PhantomData<T>,
 }
 
-impl Default for WasmtimeEngine {
-    fn default() -> Self {
+#[derive(Clone)]
+pub struct DefaultConfig {}
+
+impl WasiConfig for DefaultConfig {
+    fn new_config() -> Config {
         let mut config = wasmtime::Config::new();
-        config.parallel_compilation(false);
         config.wasm_component_model(true); // enable component linking
+        config
+    }
+}
+
+pub trait WasiConfig {
+    fn new_config() -> Config;
+}
+
+impl<T: WasiConfig + Clone + Sync + Send> Default for WasmtimeEngine<T> {
+    fn default() -> Self {
+        let config = T::new_config();
         Self {
             engine: wasmtime::Engine::new(&config)
                 .context("failed to create wasmtime engine")
                 .unwrap(),
+            config_type: PhantomData,
         }
     }
 }
@@ -57,17 +74,9 @@ impl wasmtime_wasi::preview2::WasiView for WasiCtx {
     }
 }
 
-fn version_info() -> &'static RuntimeInfo {
-    static INFO: OnceLock<RuntimeInfo> = OnceLock::new();
-    INFO.get_or_init(|| RuntimeInfo {
-        name: "wasmtime",
-        version: option_env!("WASMTIME_VERSION_INFO").unwrap_or(env!("CARGO_PKG_VERSION")),
-    })
-}
-
-impl Engine for WasmtimeEngine {
-    fn info() -> &'static RuntimeInfo {
-        version_info()
+impl<T: WasiConfig + Clone + Sync + Send + 'static> Engine for WasmtimeEngine<T> {
+    fn name() -> &'static str {
+        "wasmtime"
     }
 
     fn run_wasi(&self, ctx: &impl RuntimeContext, stdio: Stdio) -> Result<i32> {
@@ -110,12 +119,16 @@ impl Engine for WasmtimeEngine {
         }
     }
 
-    fn can_precompile() -> bool {
-        true
+    fn can_precompile(&self) -> Option<String> {
+        let mut hasher = DefaultHasher::new();
+        self.engine
+            .precompile_compatibility_hash()
+            .hash(&mut hasher);
+        Some(hasher.finish().to_string())
     }
 }
 
-impl WasmtimeEngine {
+impl<T: std::clone::Clone + Sync + WasiConfig + Send + 'static> WasmtimeEngine<T> {
     /// Execute a wasm module.
     ///
     /// This function adds wasi_preview1 to the linker and can be utilized