Skip to content

Commit

Permalink
refactor(precompile): make code a bit more readable
Browse files Browse the repository at this point in the history 
this is to resolve my comments made in PR containerd#492 in an effort to
make the code a bit more idiomatic and readable

Signed-off-by: jiaxiao zhou <[email protected]>
  • Loading branch information
@Mossaka
Mossaka committed Mar 7, 2024
1 parent 064baf3 commit 593c850
Show file tree
Hide file tree
Showing 3 changed files with 133 additions and 112 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 return 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<Vec<u8>>>> {
fn precompile(&self, _layers: &[WasmLayer]) -> Result<Vec<Option<WasmLayer>>> {
bail!("precompile not supported");
}

Expand Down
234 changes: 126 additions & 108 deletions crates/containerd-shim-wasm/src/sandbox/containerd/client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -369,12 +369,10 @@ impl Client {
})
}

fn get_image_manifest(&self, image_name: &str) -> Result<(ImageManifest, String)> {
fn get_image_manifest_and_digest(&self, image_name: &str) -> Result<(ImageManifest, String)> {
let image = self.get_image(image_name)?;
let image_digest = self.extract_image_content_sha(&image)?;
let manifest = self.read_content(&image_digest)?;
let manifest = manifest.as_slice();
let manifest = ImageManifest::from_reader(manifest)?;
let manifest = ImageManifest::from_reader(self.read_content(&image_digest)?.as_slice())?;
Ok((manifest, image_digest))
}

Expand All @@ -387,7 +385,7 @@ impl Client {
engine: &T,
) -> Result<(Vec<oci::WasmLayer>, Platform)> {
let container = self.get_container(containerd_id.to_string())?;
let (manifest, image_digest) = self.get_image_manifest(&container.image)?;
let (manifest, image_digest) = self.get_image_manifest_and_digest(&container.image)?;

let image_config_descriptor = manifest.config();
let image_config = self.read_content(image_config_descriptor.digest())?;
Expand Down Expand Up @@ -416,126 +414,140 @@ impl Client {
.iter()
.filter(|x| is_wasm_layer(x.media_type(), T::supported_layers_types()))
.map(|original_config| {
let mut digest_to_load = original_config.digest().clone();
if can_precompile {
let info = self.get_info(&digest_to_load)?;
if info.labels.contains_key(&precompile_id) {
// Safe to unwrap here since we already checked for the label's existence
digest_to_load = info.labels.get(&precompile_id).unwrap().clone();
log::info!(
"layer {} has pre-compiled content: {} ",
info.digest,
&digest_to_load
);
}
}
log::debug!("loading digest: {} ", &digest_to_load);
self.read_content(&digest_to_load)
.map(|module| WasmLayer {
config: original_config.clone(),
layer: module,
})
.or_else(|e| {
// handle content being removed from the content store out of band
if digest_to_load != *original_config.digest() {
log::error!("failed to load precompiled layer: {}", e);
log::error!("falling back to original layer and marking for recompile");
needs_precompile = can_precompile; // only mark for recompile if engine is capable
self.read_content(original_config.digest())
.map(|module| WasmLayer {
config: original_config.clone(),
layer: module,
})
} else {
Err(e)
}
})
self.read_wasm_layer(
original_config,
can_precompile,
&precompile_id,
&mut needs_precompile,
)
})
.collect::<Result<Vec<_>>>()?;

if layers.is_empty() {
log::info!("no WASM layers found in OCI image");
return Ok((vec![], platform));
}

if needs_precompile {
log::info!("precompiling layers for image: {}", container.image);
match engine.precompile(&layers) {
let compiled_layers = 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, 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 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_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_config.digest())?;
original_layer
.labels
.insert(precompile_id.clone(), precompiled_content.digest.clone());
original_layer.labels.insert(
format!("containerd.io/gc.ref.content.precompile.{}", i),
precompiled_content.digest.clone(),
);
self.update_info(original_layer)?;

// The original image is considered a root object, by adding a ref to the new compiled content
// We tell containerd to not garbage collect the new content until this image is removed from the system
// this ensures that we keep the content around after the lease is dropped
// We also save the precompiled flag here since the image labels can be mutated containerd, for example if the image is pulled twice
log::debug!(
"updating image content with precompile digest to avoid garbage collection"
);
let mut image_content = self.get_info(&image_digest)?;
image_content.labels.insert(
format!("containerd.io/gc.ref.content.precompile.{}", i),
precompiled_content.digest,
);
image_content
.labels
.insert(precompile_id.clone(), "true".to_string());
self.update_info(image_content)?;

layers_for_runtime.push(WasmLayer {
config: original_config.clone(),
layer: compiled_layer.clone(),
});
}
return Ok((layers_for_runtime, platform));
compiled_layers
}
Err(e) => {
log::error!("precompilation failed: {}", e);
return Ok((layers, platform));
}
};
}

if layers.is_empty() {
log::info!("no WASM layers found in OCI image");
return Ok((vec![], platform));
}
let mut layers_for_runtime = Vec::with_capacity(compiled_layers.len());
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 compiled_layer = compiled_layer.as_ref().unwrap();
let original_config = &compiled_layer.config;
let labels = HashMap::from([(
format!("{precompile_id}/original"),
original_config.digest().to_string(),
)]);
let precompiled_content =
self.save_content(compiled_layer.layer.clone(), &precompile_id, labels)?;

log::debug!(
"updating original layer {} with compiled layer {}",
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_config.digest())?;
original_layer
.labels
.insert(precompile_id.clone(), precompiled_content.digest.clone());
original_layer.labels.insert(
format!("containerd.io/gc.ref.content.precompile.{}", i),
precompiled_content.digest.clone(),
);
self.update_info(original_layer)?;

// The original image is considered a root object, by adding a ref to the new compiled content
// We tell containerd to not garbage collect the new content until this image is removed from the system
// this ensures that we keep the content around after the lease is dropped
// We also save the precompiled flag here since the image labels can be mutated containerd, for example if the image is pulled twice
log::debug!(
"updating image content with precompile digest to avoid garbage collection"
);
let mut image_content = self.get_info(&image_digest)?;
image_content.labels.insert(
format!("containerd.io/gc.ref.content.precompile.{}", i),
precompiled_content.digest,
);
image_content
.labels
.insert(precompile_id.clone(), "true".to_string());
self.update_info(image_content)?;

layers_for_runtime.push(compiled_layer.clone());
}
return Ok((layers_for_runtime, platform));
};

log::info!("using OCI layers");
Ok((layers, platform))
}

fn read_wasm_layer(
&self,
original_config: &oci_spec::image::Descriptor,
can_precompile: bool,
precompile_id: &String,
needs_precompile: &mut bool,
) -> std::prelude::v1::Result<WasmLayer, ShimError> {
let mut digest_to_load = original_config.digest().clone();
if can_precompile {
let info = self.get_info(&digest_to_load)?;
if let Some(label) = info.labels.get(precompile_id) {
// Safe to unwrap here since we already checked for the label's existence
digest_to_load = label.clone();
log::info!(
"layer {} has pre-compiled content: {} ",
info.digest,
&digest_to_load
);
}
}
log::debug!("loading digest: {} ", &digest_to_load);
self.read_content(&digest_to_load)
.map(|module| WasmLayer {
config: original_config.clone(),
layer: module,
})
.or_else(|e| {
// handle content being removed from the content store out of band
if digest_to_load != *original_config.digest() {
log::error!("failed to load precompiled layer: {}", e);
log::error!("falling back to original layer and marking for recompile");
*needs_precompile = can_precompile; // only mark for recompile if engine is capable
self.read_content(original_config.digest())
.map(|module| WasmLayer {
config: original_config.clone(),
layer: module,
})
} else {
Err(e)
}
})
}
}

fn precompile_label(name: &str, version: &str) -> String {
Expand Down Expand Up @@ -706,7 +718,7 @@ mod tests {
assert_eq!(layers.len(), 1);
assert_eq!(layers[0].layer, fake_precompiled_bytes.bytes);

let (manifest, _) = client.get_image_manifest(&image_name).unwrap();
let (manifest, _) = client.get_image_manifest_and_digest(&image_name).unwrap();
let original_config = manifest.layers().first().unwrap();
let info = client.get_info(original_config.digest()).unwrap();

Expand Down Expand Up @@ -847,7 +859,7 @@ mod tests {
assert_eq!(layers[0].layer, fake_precompiled_bytes.bytes);
assert_eq!(layers[1].layer, fake_precompiled_bytes2.bytes);

let (manifest, _) = client.get_image_manifest(&image_name).unwrap();
let (manifest, _) = client.get_image_manifest_and_digest(&image_name).unwrap();

let original_config1 = manifest.layers().first().unwrap();
let info1 = client.get_info(original_config1.digest()).unwrap();
Expand Down Expand Up @@ -959,7 +971,10 @@ mod tests {
&[WASM_LAYER_MEDIA_TYPE, "textfile"]
}

fn precompile(&self, layers: &[WasmLayer]) -> Result<Vec<Option<Vec<u8>>>, anyhow::Error> {
fn precompile(
&self,
layers: &[WasmLayer],
) -> Result<Vec<Option<WasmLayer>>, 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 @@ -983,7 +998,10 @@ mod tests {
true
});
let precompiled = self.precompiled_layers[&key].clone();
compiled_layers.push(Some(precompiled));
compiled_layers.push(Some(WasmLayer {
config: layer.config.clone(),
layer: precompiled,
}));
self.layers_compiled_per_call.fetch_add(1, Ordering::SeqCst);
}
Ok(compiled_layers)
Expand Down
9 changes: 6 additions & 3 deletions crates/containerd-shim-wasmtime/src/instance.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,8 +113,8 @@ impl<T: WasiConfig> Engine for WasmtimeEngine<T> {
Ok(status)
}

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

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

let compiled_layer = self.engine.precompile_module(&layer.layer)?;
compiled_layers.push(Some(compiled_layer));
compiled_layers.push(Some(WasmLayer {
config: layer.config.clone(),
layer: compiled_layer,
}));
}

Ok(compiled_layers)
Expand Down

0 comments on commit 593c850

Please sign in to comment.