mirrors/krata
1
0
Fork 0
mirror of https://github.com/edera-dev/krata.git synced 2025-11-04 07:39:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

controller: make image downloads async

Browse Source
This commit is contained in:
Alex Zenla
2024-02-25 05:38:23 +00:00
parent f66ab25521
commit 973db87b98
8 changed files with 112 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
controller/src/ctl/launch.rs
View File

@ -37,7 +37,7 @@ impl ControllerLaunch<'_> {
pub async fn perform(&mut self, request: ControllerLaunchRequest<'_>) -> Result<(Uuid, u32)> {
let uuid = Uuid::new_v4();
let name = format!("krata-{uuid}");
let image_info = self.compile(request.image)?;
let image_info = self.compile(request.image).await?;
let mut gateway_mac = MacAddr6::random();
gateway_mac.set_local(true);
@ -220,9 +220,9 @@ impl ControllerLaunch<'_> {
Ok(found.unwrap())
}
fn compile(&self, image: &str) -> Result<ImageInfo> {
async fn compile(&self, image: &str) -> Result<ImageInfo> {
let image = ImageName::parse(image)?;
let compiler = ImageCompiler::new(&self.context.image_cache)?;
compiler.compile(&image)
compiler.compile(&image).await
}
}

89
controller/src/image/fetch.rs
View File

@ -1,56 +1,57 @@
use anyhow::{anyhow, Result};
use bytes::Bytes;
use oci_spec::image::{Arch, Descriptor, ImageIndex, ImageManifest, MediaType, Os, ToDockerV2S2};
use std::io::copy;
use std::io::{Read, Write};
use std::ops::DerefMut;
use ureq::{Agent, Request, Response};
use reqwest::{Client, RequestBuilder, Response};
use tokio::{fs::File, io::AsyncWriteExt};
use url::Url;
const MANIFEST_PICKER_PLATFORM: Os = Os::Linux;
const MANIFEST_PICKER_ARCHITECTURE: Arch = Arch::Amd64;
pub struct RegistryClient {
agent: Agent,
agent: Client,
url: Url,
}
impl RegistryClient {
pub fn new(url: Url) -> Result<RegistryClient> {
Ok(RegistryClient {
agent: Agent::new(),
agent: Client::new(),
url,
})
}
fn call(&mut self, req: Request) -> Result<Response> {
Ok(req.call()?)
async fn call(&mut self, req: RequestBuilder) -> Result<Response> {
self.agent.execute(req.build()?).await.map_err(|x| x.into())
}
pub fn get_blob(&mut self, name: &str, descriptor: &Descriptor) -> Result<Vec<u8>> {
pub async fn get_blob(&mut self, name: &str, descriptor: &Descriptor) -> Result<Bytes> {
let url = self
.url
.join(&format!("/v2/{}/blobs/{}", name, descriptor.digest()))?;
let response = self.call(self.agent.get(url.as_str()))?;
let mut buffer: Vec<u8> = Vec::new();
response.into_reader().read_to_end(&mut buffer)?;
Ok(buffer)
let response = self.call(self.agent.get(url.as_str())).await?;
Ok(response.bytes().await?)
}
pub fn write_blob(
pub async fn write_blob_to_file(
&mut self,
name: &str,
descriptor: &Descriptor,
dest: &mut dyn Write,
mut dest: File,
) -> Result<u64> {
let url = self
.url
.join(&format!("/v2/{}/blobs/{}", name, descriptor.digest()))?;
let response = self.call(self.agent.get(url.as_str()))?;
let mut reader = response.into_reader();
Ok(copy(reader.deref_mut(), dest)?)
let mut response = self.call(self.agent.get(url.as_str())).await?;
let mut size: u64 = 0;
while let Some(chunk) = response.chunk().await? {
dest.write_all(&chunk).await?;
size += chunk.len() as u64;
}
Ok(size)
}
pub fn get_manifest_with_digest(
async fn get_raw_manifest_with_digest(
&mut self,
name: &str,
reference: &str,
@ -65,24 +66,60 @@ impl RegistryClient {
MediaType::ImageIndex,
MediaType::ImageIndex.to_docker_v2s2()?,
);
let response = self.call(self.agent.get(url.as_str()).set("Accept", &accept))?;
let response = self
.call(self.agent.get(url.as_str()).header("Accept", &accept))
.await?;
let digest = response
.headers()
.get("Docker-Content-Digest")
.ok_or_else(|| anyhow!("fetching manifest did not yield a content digest"))?
.to_str()?
.to_string();
let manifest = serde_json::from_str(&response.text().await?)?;
Ok((manifest, digest))
}
pub async fn get_manifest_with_digest(
&mut self,
name: &str,
reference: &str,
) -> Result<(ImageManifest, String)> {
let url = self
.url
.join(&format!("/v2/{}/manifests/{}", name, reference))?;
let accept = format!(
"{}, {}, {}, {}",
MediaType::ImageManifest.to_docker_v2s2()?,
MediaType::ImageManifest,
MediaType::ImageIndex,
MediaType::ImageIndex.to_docker_v2s2()?,
);
let response = self
.call(self.agent.get(url.as_str()).header("Accept", &accept))
.await?;
let content_type = response
.header("Content-Type")
.ok_or_else(|| anyhow!("registry response did not have a Content-Type header"))?;
.headers()
.get("Content-Type")
.ok_or_else(|| anyhow!("registry response did not have a Content-Type header"))?
.to_str()?;
if content_type == MediaType::ImageIndex.to_string()
|| content_type == MediaType::ImageIndex.to_docker_v2s2()?
{
let index = ImageIndex::from_reader(response.into_reader())?;
let index = serde_json::from_str(&response.text().await?)?;
let descriptor = self
.pick_manifest(index)
.ok_or_else(|| anyhow!("unable to pick manifest from index"))?;
return self.get_manifest_with_digest(name, descriptor.digest());
return self
.get_raw_manifest_with_digest(name, descriptor.digest())
.await;
}
let digest = response
.header("Docker-Content-Digest")
.headers()
.get("Docker-Content-Digest")
.ok_or_else(|| anyhow!("fetching manifest did not yield a content digest"))?
.to_str()?
.to_string();
let manifest = ImageManifest::from_reader(response.into_reader())?;
let manifest = serde_json::from_str(&response.text().await?)?;
Ok((manifest, digest))
}

47
controller/src/image/mod.rs
View File

@ -94,8 +94,8 @@ impl ImageCompiler<'_> {
Ok(ImageCompiler { cache })
}
pub fn compile(&self, image: &ImageName) -> Result<ImageInfo> {
debug!("ImageCompiler compile image={image}");
pub async fn compile(&self, image: &ImageName) -> Result<ImageInfo> {
debug!("compile image={image}");
let mut tmp_dir = std::env::temp_dir().clone();
tmp_dir.push(format!("krata-compile-{}", Uuid::new_v4()));
@ -109,12 +109,14 @@ impl ImageCompiler<'_> {
let mut squash_file = tmp_dir.clone();
squash_file.push("image.squashfs");
let info = self.download_and_compile(image, &layer_dir, &image_dir, &squash_file)?;
let info = self
.download_and_compile(image, &layer_dir, &image_dir, &squash_file)
.await?;
fs::remove_dir_all(&tmp_dir)?;
Ok(info)
}
fn download_and_compile(
async fn download_and_compile(
&self,
image: &ImageName,
layer_dir: &Path,
@ -122,11 +124,13 @@ impl ImageCompiler<'_> {
squash_file: &PathBuf,
) -> Result<ImageInfo> {
debug!(
"ImageCompiler download manifest image={image}, image_dir={}",
"download manifest image={image}, image_dir={}",
image_dir.to_str().unwrap()
);
let mut client = RegistryClient::new(image.registry_url()?)?;
let (manifest, digest) = client.get_manifest_with_digest(&image.name, &image.reference)?;
let (manifest, digest) = client
.get_manifest_with_digest(&image.name, &image.reference)
.await?;
let cache_key = format!(
"manifest={}:squashfs-version={}\n",
digest, IMAGE_SQUASHFS_VERSION
@ -138,21 +142,24 @@ impl ImageCompiler<'_> {
}
debug!(
"ImageCompiler download config digest={} size={}",
"download config digest={} size={}",
manifest.config().digest(),
manifest.config().size(),
);
let config_bytes = client.get_blob(&image.name, manifest.config())?;
let config_bytes = client.get_blob(&image.name, manifest.config()).await?;
let config: ImageConfiguration = serde_json::from_slice(&config_bytes)?;
let mut layers: Vec<LayerFile> = Vec::new();
for layer in manifest.layers() {
layers.push(self.download_layer(image, layer, layer_dir, &mut client)?);
layers.push(
self.download_layer(image, layer, layer_dir, &mut client)
.await?,
);
}
for layer in layers {
debug!(
"ImageCompiler process layer digest={} compression={:?}",
"process layer digest={} compression={:?}",
&layer.digest, layer.compression
);
let mut archive = Archive::new(layer.open_reader()?);
@ -199,7 +206,7 @@ impl ImageCompiler<'_> {
}
trace!(
"ImageCompiler whiteout entry layer={} path={:?}",
"whiteout entry layer={} path={:?}",
&layer.digest,
entry.path()?
);
@ -219,7 +226,7 @@ impl ImageCompiler<'_> {
}
} else {
warn!(
"ImageCompiler whiteout entry missing locally layer={} path={:?} local={:?}",
"whiteout entry missing locally layer={} path={:?} local={:?}",
&layer.digest,
entry.path()?,
dst,
@ -231,7 +238,7 @@ impl ImageCompiler<'_> {
fs::remove_dir(&dst)?;
} else {
warn!(
"ImageCompiler whiteout entry missing locally layer={} path={:?} local={:?}",
"whiteout entry missing locally layer={} path={:?} local={:?}",
&layer.digest,
entry.path()?,
dst,
@ -247,7 +254,7 @@ impl ImageCompiler<'_> {
image_dir: &PathBuf,
) -> Result<()> {
trace!(
"ImageCompiler unpack entry layer={} path={:?} type={:?}",
"unpack entry layer={} path={:?} type={:?}",
&layer.digest,
entry.path()?,
entry.header().entry_type()
@ -285,7 +292,7 @@ impl ImageCompiler<'_> {
Ok(())
}
fn download_layer(
async fn download_layer(
&self,
image: &ImageName,
layer: &Descriptor,
@ -293,7 +300,7 @@ impl ImageCompiler<'_> {
client: &mut RegistryClient,
) -> Result<LayerFile> {
debug!(
"ImageCompiler download layer digest={} size={}",
"download layer digest={} size={}",
layer.digest(),
layer.size()
);
@ -303,8 +310,8 @@ impl ImageCompiler<'_> {
tmp_path.push(format!("{}.tmp", layer.digest()));
{
let mut file = File::create(&layer_path)?;
let size = client.write_blob(&image.name, layer, &mut file)?;
let file = tokio::fs::File::create(&layer_path).await?;
let size = client.write_blob_to_file(&image.name, layer, file).await?;
if layer.size() as u64 != size {
return Err(anyhow!(
"downloaded layer size differs from size in manifest",
@ -344,7 +351,7 @@ impl ImageCompiler<'_> {
.to_str()
.ok_or_else(|| anyhow!("failed to strip prefix of tmpdir"))?;
let rel = format!("/{}", rel);
trace!("ImageCompiler squash write {}", rel);
trace!("squash write {}", rel);
let typ = entry.file_type();
let metadata = fs::symlink_metadata(entry.path())?;
let uid = metadata.uid();
@ -400,7 +407,7 @@ impl ImageCompiler<'_> {
.ok_or_else(|| anyhow!("failed to convert squashfs string"))?;
let mut file = File::create(squash_file)?;
trace!("ImageCompiler squash generate: {}", squash_file_path);
trace!("squash generate: {}", squash_file_path);
writer.write(&mut file)?;
Ok(())
}