WIP: Store compiled modules

Cargo.toml

@@ -18,6 +18,7 @@ oci-spec = { version = "0.5", features = ["runtime"] }
 thiserror = "1.0"
 serde_json = "1.0"
 nix = "0.25"
+sha256 = "1.1"
 
 [dev-dependencies]
 tempfile = "3.0"

crates/containerd-shim-wasm/Cargo.toml

@@ -25,6 +25,7 @@ clone3 = "0.2"
 libc = "0.2"
 caps = "0.1"
 proc-mounts = "0.3"
+sha256 = "1.1"
 
 [build-dependencies]
 ttrpc-codegen = { version = "0.3", optional = true }

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

@@ -0,0 +1,161 @@
+use anyhow::{Context, Error, Result};
+use serde::{Deserialize, Serialize};
+use sha256::try_digest;
+use std::fs::File;
+use std::io::prelude::*;
+use std::path::PathBuf;
+
+#[derive(Deserialize, Serialize, PartialEq, Clone)]
+pub struct Digest {
+    alg: String,
+    enc: String,
+}
+
+impl Digest {
+    fn new(algorithm: String, encoded: String) -> Self {
+        Self {
+            alg: algorithm,
+            enc: encoded,
+        }
+    }
+
+    fn algorithm(&self) -> &str {
+        &self.alg
+    }
+
+    fn encoded(&self) -> &str {
+        &self.enc
+    }
+}
+
+impl std::fmt::Display for Digest {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}:{}", self.algorithm(), self.encoded())
+    }
+}
+
+impl TryFrom<&str> for Digest {
+    type Error = Error;
+    fn try_from(s: &str) -> Result<Self> {
+        let mut parts = s.splitn(2, ':');
+        let algorithm = parts
+            .next()
+            .ok_or_else(|| Error::msg(format!("invalid digest format: {}", s)))?;
+        let hex = parts
+            .next()
+            .ok_or_else(|| Error::msg(format!("invalid digest format: {}", s)))?;
+        Ok(Self::new(algorithm.to_string(), hex.to_string()))
+    }
+}
+
+impl TryFrom<String> for Digest {
+    type Error = Error;
+    fn try_from(s: String) -> Result<Self> {
+        Self::try_from(s.as_str())
+    }
+}
+
+impl Into<String> for Digest {
+    fn into(self) -> String {
+        self.algorithm().to_owned() + ":" + self.encoded()
+    }
+}
+
+pub struct Store {
+    dir: PathBuf,
+}
+
+#[derive(Deserialize, Serialize)]
+pub struct Metadata {
+    pub digest: Digest,
+}
+
+impl Store {
+    pub fn new(dir: &str) -> Self {
+        let p = PathBuf::from(dir);
+        std::fs::create_dir_all(p.join("blobs/sha256")).unwrap();
+        std::fs::create_dir_all(p.join("ingests")).unwrap();
+        std::fs::create_dir_all(p.join("metadata/sha256")).unwrap();
+
+        Self {
+            dir: PathBuf::from(dir),
+        }
+    }
+
+    pub fn path(&self, dgst: &Digest) -> PathBuf {
+        self.dir
+            .join("blobs")
+            .join(dgst.algorithm())
+            .join(dgst.encoded())
+    }
+
+    pub fn metadata(&self, dgst: &Digest) -> Result<Metadata> {
+        let path = self
+            .dir
+            .join("metadata")
+            .join(dgst.algorithm())
+            .join(dgst.encoded());
+
+        let mut file = File::open(path).context("could not open metadata file")?;
+        let mut buf = String::new();
+        file.read_to_string(&mut buf)
+            .context("could not read metadata file")?;
+
+        Ok(serde_json::from_str(&buf).context("could not parse metadata file")?)
+    }
+
+    pub fn write_metadata(&self, dgst: &Digest, data: &Metadata) -> Result<()> {
+        let path = self
+            .dir
+            .join("metadata")
+            .join(dgst.algorithm())
+            .join(dgst.encoded());
+
+        let f = File::create(path)?;
+        serde_json::to_writer(f, data)?;
+        Ok(())
+    }
+
+    pub fn writer(&mut self, id: String) -> Result<ContentWriter> {
+        let ingest_path = self.dir.join("ingests").join(id);
+        let target = self.dir.join("blobs");
+        let f = File::create(&ingest_path)?;
+        let cw = ContentWriter {
+            f,
+            ingest_path,
+            target,
+        };
+        Ok(cw)
+    }
+}
+
+pub struct ContentWriter {
+    f: File,
+    ingest_path: PathBuf,
+    target: PathBuf,
+}
+
+impl ContentWriter {
+    pub fn write(&mut self, data: &[u8]) -> Result<usize> {
+        self.f.try_clone()?.write(data).context("write")
+    }
+
+    pub fn commit(&mut self, expected: Option<Digest>) -> Result<Digest> {
+        self.f.try_clone()?.flush().context("flush")?;
+
+        let dgst = try_digest(self.ingest_path.as_path()).context("digest")?;
+        let d: Digest = ("sha256:".to_owned() + &dgst).try_into()?;
+
+        if let Some(ex) = expected {
+            if d != ex {
+                return Err(Error::msg(format!("digest mismatch: {} != {}", d, ex)));
+            }
+        }
+
+        std::fs::rename(
+            &self.ingest_path,
+            &self.target.join(d.algorithm()).join(d.encoded()),
+        )?;
+        Ok(d)
+    }
+}

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

@@ -1,2 +1,3 @@
+pub mod content;
 pub mod sandbox;
 pub mod services;

src/instance.rs

@@ -1,23 +1,25 @@
 use std::fs::OpenOptions;
 use std::path::Path;
 use std::sync::mpsc::Sender;
-use std::sync::{Arc, Condvar, Mutex};
+use std::sync::{Arc, Condvar, Mutex, RwLock};
 use std::thread;
 
-use anyhow::Context;
+use anyhow::{Context, Error as AnyError};
 use chrono::{DateTime, Utc};
+use containerd_shim_wasm::content::{Metadata, Store as ContentStore};
 use containerd_shim_wasm::sandbox::error::Error;
 use containerd_shim_wasm::sandbox::exec;
 use containerd_shim_wasm::sandbox::oci;
 use containerd_shim_wasm::sandbox::{EngineGetter, Instance, InstanceConfig};
 use log::{debug, error};
+use sha256::try_digest;
 use wasmtime::{Engine, Linker, Module, Store};
 use wasmtime_wasi::{sync::file::File as WasiFile, WasiCtx, WasiCtxBuilder};
 
-use super::error::WasmtimeError;
 use super::oci_wasmtime;
 
 type ExitCode = (Mutex<Option<(u32, DateTime<Utc>)>>, Condvar);
+
 pub struct Wasi {
     exit_code: Arc<ExitCode>,
     engine: wasmtime::Engine,
@@ -28,6 +30,8 @@ pub struct Wasi {
     bundle: String,
 
     pidfd: Arc<Mutex<Option<exec::PidFD>>>,
+
+    store: Arc<RwLock<containerd_shim_wasm::content::Store>>,
 }
 
 #[cfg(test)]
@@ -82,11 +86,12 @@ fn load_spec(bundle: String) -> Result<oci::Spec, Error> {
 
 pub fn prepare_module(
     engine: wasmtime::Engine,
+    store: Arc<RwLock<containerd_shim_wasm::content::Store>>,
     spec: &oci::Spec,
     stdin_path: String,
     stdout_path: String,
     stderr_path: String,
-) -> Result<(WasiCtx, Module), WasmtimeError> {
+) -> Result<(WasiCtx, Module), AnyError> {
     debug!("opening rootfs");
     let rootfs = oci_wasmtime::get_rootfs(&spec)?;
     let args = oci::get_args(&spec);
@@ -129,9 +134,34 @@ pub fn prepare_module(
     let mod_path = oci::get_root(&spec).join(cmd);
 
     debug!("loading module from file");
-    let module = Module::from_file(&engine, mod_path)
-        .map_err(|err| Error::Others(format!("could not load module from file: {}", err)))?;
+    let mut s = store.write().unwrap();
+    let mod_dgst = try_digest(mod_path.as_path()).context("digest")?;
+    let dgst = ("sha256:".to_owned() + &mod_dgst).try_into()?;
+
+    let p = match s.metadata(&dgst) {
+        Ok(m) => s.path(&m.digest),
+        Err(_) => {
+            let data = std::fs::read(mod_path).context("could not open module file")?;
+            let compiled = engine.precompile_module(&data)?;
+            let mut w = s
+                .writer(dgst.to_string())
+                .context("could not open content store writer")?;
+
+            w.write(&compiled)?;
+            let compiled_digest = w.commit(None)?;
+            drop(w);
+
+            s.write_metadata(
+                &dgst,
+                &Metadata {
+                    digest: compiled_digest.clone(),
+                },
+            )?;
+            s.path(&compiled_digest)
+        }
+    };
 
+    let module = unsafe { Module::deserialize_file(&engine, &p) }?;
     Ok((wctx, module))
 }
 
@@ -147,6 +177,9 @@ impl Instance for Wasi {
             stderr: cfg.get_stderr().unwrap_or_default(),
             bundle: cfg.get_bundle().unwrap_or_default(),
             pidfd: Arc::new(Mutex::new(None)),
+            store: Arc::new(RwLock::new(ContentStore::new(
+                "/run/io.containerd.wasmtime.v1/content",
+            ))),
         }
     }
     fn start(&self) -> Result<u32, Error> {
@@ -163,8 +196,15 @@ impl Instance for Wasi {
 
         debug!("preparing module");
         let spec = load_spec(self.bundle.clone())?;
-        let m = prepare_module(engine.clone(), &spec, stdin, stdout, stderr)
-            .map_err(|e| Error::Others(format!("error setting up module: {}", e)))?;
+        let m = prepare_module(
+            engine.clone(),
+            self.store.clone(),
+            &spec,
+            stdin,
+            stdout,
+            stderr,
+        )
+        .map_err(|e| Error::Others(format!("error setting up module: {}", e)))?;
 
         let mut store = Store::new(&engine, m.0);