Merge pull request #405 from jsturtevant/precompile-on-oci-loading

Store Compiled Modules when using OCI Wasm Images

CONTRIBUTING.md

@@ -82,6 +82,12 @@ test instance::wasitest::test_delete_after_create ... ok
 test instance::wasitest::test_wasi ... ok
 ```
 
+Run individual test via cargo adding `RUST_LOG=trace` (adjust the level of logging as needed) to see shim output. Also adjust the test name as needed.
+
+```
+RUST_LOG=DEBUG cargo test --package containerd-shim-wasmtime --lib -- wasmtime_tests::test_hello_world --exact --nocapture
+```
+
 ### End to End tests
 
 The e2e test run on [k3s](https://k3s.io/) and [kind](https://kind.sigs.k8s.io/).  A test image is built using [oci-tar-builder](./crates/oci-tar-builder/) and is loaded onto the clusters.  This test image is not pushed to an external registry so be sure to use the Makefile targets to build the image and load it on the cluster.

Cargo.toml

@@ -25,6 +25,7 @@ chrono = { version = "0.4", default-features = false, features = ["clock"] }
 containerd-shim = "0.6.0"
 containerd-shim-wasm = { path = "crates/containerd-shim-wasm", version = "0.4.0" }
 containerd-shim-wasm-test-modules = { path = "crates/containerd-shim-wasm-test-modules", version = "0.3.1"}
+oci-tar-builder = { path = "crates/oci-tar-builder", version = "0.3.1" }
 crossbeam = { version = "0.8.4", default-features = false }
 env_logger = "0.10"
 libc = "0.2.153"

Makefile

@@ -238,6 +238,13 @@ test/k8s/deploy-workload-oci-%: test/k8s/clean test/k8s/cluster-% dist/img-oci.t
 	# verify that we are still running after some time
 	sleep 5s
 	kubectl --context=kind-$(KIND_CLUSTER_NAME) wait deployment wasi-demo --for condition=Available=True --timeout=5s
+	@if [ "$*" = "wasmtime" ]; then \
+		set -e; \
+		echo "checking for pre-compiled label and ensuring can scale"; \
+		docker exec $(KIND_CLUSTER_NAME)-control-plane ctr -n k8s.io i ls | grep "runwasi.io/precompiled"; \
+		kubectl --context=kind-$(KIND_CLUSTER_NAME) scale deployment wasi-demo --replicas=4; \
+		kubectl --context=kind-$(KIND_CLUSTER_NAME) wait deployment wasi-demo --for condition=Available=True --timeout=5s; \
+	fi
 
 .PHONY: test/k8s-%
 test/k8s-%: test/k8s/deploy-workload-%
@@ -289,6 +296,13 @@ test/k3s-oci-%: dist/img-oci.tar bin/k3s dist-%
 	sleep 5s
 	sudo bin/k3s kubectl wait deployment wasi-demo --for condition=Available=True --timeout=5s
 	sudo bin/k3s kubectl get pods -o wide
+	@if [ "$*" = "wasmtime" ]; then \
+		set -e; \
+		echo "checking for pre-compiled label and ensuring can scale"; \
+		sudo bin/k3s ctr -n k8s.io i ls | grep "runwasi.io/precompiled"; \
+		sudo bin/k3s kubectl scale deployment wasi-demo --replicas=4; \
+		sudo bin/k3s kubectl wait deployment wasi-demo --for condition=Available=True --timeout=5s; \
+	fi
 	sudo bin/k3s kubectl delete -f test/k8s/deploy.oci.yaml
 	sudo bin/k3s kubectl wait deployment wasi-demo --for delete --timeout=60s
 
@@ -299,6 +313,7 @@ test/k3s/clean: bin/k3s/clean;
 clean:
 	-rm -rf dist
 	-rm -rf bin
+	-rm -rf test/k8s/_out
 	-$(MAKE) test-image/clean
 	-$(MAKE) test/k8s/clean
 	-$(MAKE) test/k3s/clean

crates/containerd-shim-wasm/Cargo.toml

@@ -16,6 +16,7 @@ anyhow = { workspace = true }
 chrono = { workspace = true }
 containerd-shim = { workspace = true }
 containerd-shim-wasm-test-modules = { workspace = true, optional = true }
+oci-tar-builder = { workspace = true, optional = true }
 crossbeam = { workspace = true }
 env_logger = { workspace = true, optional = true }
 git-version = "0.3.9"
@@ -32,6 +33,9 @@ wat = { workspace = true }
 tokio = { version = "1.36.0", features = [ "full" ] }
 futures = { version = "0.3.30" }
 wasmparser = "0.121.0"
+tokio-stream = { version = "0.1" }
+prost-types = "0.11" # should match version in containerd-shim
+sha256 = "1.4.0"
 
 [target.'cfg(unix)'.dependencies]
 caps = "0.5"
@@ -51,8 +55,9 @@ ttrpc-codegen = { version = "0.4.2", optional = true }
 containerd-shim-wasm-test-modules = { workspace = true }
 env_logger = { workspace = true }
 tempfile = { workspace = true }
+oci-tar-builder = { workspace = true}
 
 [features]
-testing = ["dep:containerd-shim-wasm-test-modules", "dep:env_logger", "dep:tempfile"]
+testing = ["dep:containerd-shim-wasm-test-modules", "dep:env_logger", "dep:tempfile", "dep:oci-tar-builder"]
 generate_bindings = ["ttrpc-codegen"]
 generate_doc = []

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

@@ -34,6 +34,7 @@ pub trait RuntimeContext {
 }
 
 /// The source for a WASI module / components.
+#[derive(Debug)]
 pub enum Source<'a> {
     // The WASI module is a file in the file system.
     File(PathBuf),

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

@@ -1,7 +1,7 @@
 use std::fs::File;
 use std::io::Read;
 
-use anyhow::{Context, Result};
+use anyhow::{bail, Context, Result};
 
 use super::Source;
 use crate::container::{PathResolve, RuntimeContext};
@@ -52,4 +52,27 @@ pub trait Engine: Clone + Send + Sync + 'static {
     fn supported_layers_types() -> &'static [&'static str] {
         &["application/vnd.bytecodealliance.wasm.component.layer.v0+wasm"]
     }
+
+    /// Precompiles a module that is in the WASM OCI layer format
+    /// This is used to precompile a module before it is run and will be called if can_precompile returns true.
+    /// It is called only the first time a module is run and the resulting bytes will be cached in the containerd content store.  
+    /// The cached, precompiled module will be reloaded on subsequent runs.
+    fn precompile(&self, _layers: &[Vec<u8>]) -> Result<Vec<u8>> {
+        bail!("precompilation not supported for this runtime")
+    }
+
+    /// 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
+    }
 }