Skip to content

Commit

Permalink
use Result<Vec<Option<Vec<u8>>>> and make labels clear
Browse files Browse the repository at this point in the history 
Signed-off-by: James Sturtevant <[email protected]>
  • Loading branch information
@jsturtevant
jsturtevant committed Mar 5, 2024
1 parent 5f127ed commit 03bc3d2
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 61 deletions.
2 changes: 1 addition & 1 deletion crates/containerd-shim-wasm/src/container/engine.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ pub trait Engine: Clone + Send + Sync + 'static {
/// The cached, precompiled layers will be reloaded on subsequent runs.
/// The runtime is expected to runtime the same number of layers passed in, if the layer cannot be precompiled it should return None for that layer.
/// In some edge cases it is possible that the layers may already be precompiled and None should be returned in this case.
fn precompile(&self, _layers: &[WasmLayer]) -> Result<Vec<Option<WasmLayer>>> {
fn precompile(&self, _layers: &[WasmLayer]) -> Result<Vec<Option<Vec<u8>>>> {
bail!("precompile not supported");
}

Expand Down
87 changes: 40 additions & 47 deletions crates/containerd-shim-wasm/src/sandbox/containerd/client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,11 +137,11 @@ impl Client {
fn save_content(
&self,
data: Vec<u8>,
original_digest: &str,
label: &str,
unique_id: &str,
labels: HashMap<String, String>,
) -> Result<WriteContent> {
let expected = format!("sha256:{}", digest(data.clone()));
let reference = format!("precompile-{}", label);
let reference = format!("precompile-{}", unique_id);
let lease = self.lease(reference.clone())?;

let digest = self.rt.block_on(async {
Expand Down Expand Up @@ -193,10 +193,6 @@ impl Client {
// we don't need to copy the content again. Container tells us it found the blob
// by returning the offset of the content that was found.
let data_to_write = data[response.offset as usize..].to_vec();

// Write and commit at same time
let mut labels = HashMap::new();
labels.insert(format!("{}/original", label), original_digest.to_string());
let commit_request = WriteContentRequest {
action: WriteAction::Commit.into(),
total: len,
Expand Down Expand Up @@ -439,31 +435,38 @@ impl Client {
log::info!("precompiling layers for image: {}", container.image);
match engine.precompile(&layers) {
Ok(compiled_layers) => {
if compiled_layers.len() != layers.len() {
return Err(ShimError::FailedPrecondition(
"precompile returned wrong number of layers".to_string(),
));
}

let mut layers_for_runtime = Vec::with_capacity(compiled_layers.len());
for (i, precompiled_layer) in compiled_layers.iter().enumerate() {
if precompiled_layer.is_none() {
log::debug!("no precompiled layer using original");
for (i, compiled_layer) in compiled_layers.iter().enumerate() {
if compiled_layer.is_none() {
log::debug!("no compiled layer using original");
layers_for_runtime.push(layers[i].clone());
continue;
}

let precompiled_layer = precompiled_layer.as_ref().unwrap();
let original_layer = &layers[i];
let precompiled_content = self.save_content(
precompiled_layer.layer.clone(),
original_layer.config.digest(),
&precompile_id,
)?;
let compiled_layer = compiled_layer.as_ref().unwrap();
let original_config = &layers[i].config;
let labels = HashMap::from([(
format!("{precompile_id}/original"),
original_config.digest().to_string(),
)]);
let precompiled_content =
self.save_content(compiled_layer.clone(), &precompile_id, labels)?;

log::debug!(
"updating original layer {} with compiled layer {}",
original_layer.config.digest(),
original_config.digest(),
precompiled_content.digest
);
// We add two labels here:
// - one with cache key per engine instance
// - one with a gc ref flag so it doesn't get cleaned up as long as the original layer exists
let mut original_layer = self.get_info(original_layer.config.digest())?;
let mut original_layer = self.get_info(original_config.digest())?;
original_layer
.labels
.insert(precompile_id.clone(), precompiled_content.digest.clone());
Expand All @@ -489,7 +492,15 @@ impl Client {
.labels
.insert(precompile_id.clone(), "true".to_string());
self.update_info(image_content)?;
layers_for_runtime.push(precompiled_layer.clone());

// use the original digest information but replace the digest
let mut config = original_config.clone();
config.set_digest(digest(compiled_layer));
config.set_size(compiled_layer.len() as i64);
layers_for_runtime.push(WasmLayer {
config,
layer: compiled_layer.clone(),
});
}
return Ok((layers_for_runtime, platform));
}
Expand Down Expand Up @@ -529,7 +540,6 @@ mod tests {
use std::sync::atomic::{AtomicI32, Ordering};
use std::sync::Arc;

use oci_spec::image::Descriptor;
use oci_tar_builder::WASM_LAYER_MEDIA_TYPE;
use rand::prelude::*;

Expand All @@ -549,23 +559,23 @@ 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", &label).unwrap();
let label = HashMap::from([(precompile_label("test", "hasdfh"), "original".to_string())]);
let returned = client.save_content(data, "test", label.clone()).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", &label)
.save_content(data.clone(), "test", label.clone())
.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
// a second call should be successful since it already exists
drop(returned);

// a second call should be successful since it already exists
let returned = client.save_content(data, "original", &label).unwrap();
let returned = client.save_content(data, "test", label.clone()).unwrap();
assert_eq!(expected, returned.digest);

client.delete_content(expected.clone()).unwrap();
Expand Down Expand Up @@ -856,7 +866,7 @@ mod tests {
#[derive(Clone)]
struct FakePrecomiplerEngine {
precompile_id: Option<String>,
precompiled_layers: HashMap<String, WasmLayer>,
precompiled_layers: HashMap<String, Vec<u8>>,
precompile_called: Arc<AtomicI32>,
layers_compiled_per_call: Arc<AtomicI32>,
}
Expand Down Expand Up @@ -884,18 +894,8 @@ mod tests {
precompiled_content: &oci_helpers::ImageContent,
) {
let key = digest(original);

self.precompiled_layers.insert(
key,
WasmLayer {
config: Descriptor::new(
MediaType::Other(precompiled_content.media_type.clone()),
precompiled_content.bytes.len() as i64,
digest(precompiled_content.bytes.clone()),
),
layer: precompiled_content.bytes.clone(),
},
);
self.precompiled_layers
.insert(key, precompiled_content.bytes.clone());
}
}

Expand All @@ -920,10 +920,7 @@ mod tests {
&[WASM_LAYER_MEDIA_TYPE, "textfile"]
}

fn precompile(
&self,
layers: &[WasmLayer],
) -> Result<Vec<Option<WasmLayer>>, anyhow::Error> {
fn precompile(&self, layers: &[WasmLayer]) -> Result<Vec<Option<Vec<u8>>>, anyhow::Error> {
self.layers_compiled_per_call.store(0, Ordering::SeqCst);
self.precompile_called.fetch_add(1, Ordering::SeqCst);
let mut compiled_layers = vec![];
Expand All @@ -935,11 +932,7 @@ mod tests {
}

let key = digest(layer.layer.clone());
if self
.precompiled_layers
.values()
.any(|l| l.config.digest() == &key)
{
if self.precompiled_layers.values().any(|l| digest(l) == key) {
// simulate scenario were one of the layers is already compiled
compiled_layers.push(None);
continue;
Expand Down
17 changes: 4 additions & 13 deletions crates/containerd-shim-wasmtime/src/instance.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ use containerd_shim_wasm::container::{
Engine, Entrypoint, Instance, RuntimeContext, Stdio, WasmBinaryType,
};
use containerd_shim_wasm::sandbox::WasmLayer;
use sha256::digest;
use wasi_common::I32Exit;
use wasmtime::component::{self as wasmtime_component, Component, ResourceTable};
use wasmtime::{Config, Module, Precompiled, Store};
Expand Down Expand Up @@ -114,8 +113,8 @@ impl<T: WasiConfig> Engine for WasmtimeEngine<T> {
Ok(status)
}

fn precompile(&self, layers: &[WasmLayer]) -> Result<Vec<Option<WasmLayer>>> {
let mut compiled_layers = Vec::<Option<WasmLayer>>::with_capacity(layers.len());
fn precompile(&self, layers: &[WasmLayer]) -> Result<Vec<Option<Vec<u8>>>> {
let mut compiled_layers = Vec::<Option<Vec<u8>>>::with_capacity(layers.len());

for layer in layers {
if self.engine.detect_precompiled(&layer.layer).is_some() {
Expand All @@ -124,16 +123,8 @@ impl<T: WasiConfig> Engine for WasmtimeEngine<T> {
continue;
}

let compiled_layer = match self.engine.precompile_module(&layer.layer) {
Ok(compiled_layer) => compiled_layer,
Err(err) => return Err(err),
};
let mut config = layer.config.clone();
config.set_digest(digest(&compiled_layer));
compiled_layers.push(Some(WasmLayer {
layer: compiled_layer,
config,
}));
let compiled_layer = self.engine.precompile_module(&layer.layer)?;
compiled_layers.push(Some(compiled_layer));
}

Ok(compiled_layers)
Expand Down

0 comments on commit 03bc3d2

Please sign in to comment.