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

feat: oci concurrency improvements (#95)

Browse Source 
* feat: implement improved and detailed oci progress indication

* feat: implement on-disk indexes of images

* oci: utilize rw-lock for increased cache performance
This commit is contained in:
Alex Zenla
2024-04-16 09:29:54 -07:00
committed by GitHub
parent e450ebd2a2
commit 8135307283
18 changed files with 834 additions and 346 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
crates/oci/examples/squashify.rs
View File

@ -25,37 +25,19 @@ async fn main() -> Result<()> {
let (context, mut receiver) = OciProgressContext::create();
tokio::task::spawn(async move {
loop {
let Ok(mut progress) = receiver.recv().await else {
return;
};
let mut drain = 0;
loop {
if drain >= 10 {
break;
}
if let Ok(latest) = receiver.try_recv() {
progress = latest;
} else {
break;
}
drain += 1;
if (receiver.changed().await).is_err() {
break;
}
let progress = receiver.borrow_and_update();
println!("phase {:?}", progress.phase);
for (id, layer) in &progress.layers {
println!(
"{} {:?} {} of {}",
id, layer.phase, layer.value, layer.total
)
println!("{} {:?} {:?}", id, layer.phase, layer.indication,)
}
}
});
let service = OciPackerService::new(seed, &cache_dir, OciPlatform::current())?;
let service = OciPackerService::new(seed, &cache_dir, OciPlatform::current()).await?;
let packed = service
.request(image.clone(), OciPackedFormat::Squashfs, context)
.request(image.clone(), OciPackedFormat::Squashfs, false, context)
.await?;
println!(
"generated squashfs of {} to {}",

34
crates/oci/src/assemble.rs
View File

@ -1,23 +1,23 @@
use crate::fetch::{OciImageFetcher, OciImageLayer, OciResolvedImage};
use crate::fetch::{OciImageFetcher, OciImageLayer, OciImageLayerReader, OciResolvedImage};
use crate::progress::OciBoundProgress;
use crate::schema::OciSchema;
use crate::vfs::{VfsNode, VfsTree};
use anyhow::{anyhow, Result};
use log::{debug, trace, warn};
use oci_spec::image::{ImageConfiguration, ImageManifest};
use oci_spec::image::{Descriptor, ImageConfiguration, ImageManifest};
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::fs;
use tokio::io::AsyncRead;
use tokio_stream::StreamExt;
use tokio_tar::{Archive, Entry};
use uuid::Uuid;
pub struct OciImageAssembled {
pub digest: String,
pub descriptor: Descriptor,
pub manifest: OciSchema<ImageManifest>,
pub config: OciSchema<ImageConfiguration>,
pub vfs: Arc<VfsTree>,
@ -115,12 +115,14 @@ impl OciImageAssembler {
);
self.progress
.update(|progress| {
progress.extracting_layer(&layer.digest, 0, 1);
progress.start_extracting_layer(&layer.digest);
})
.await;
debug!("process layer digest={}", &layer.digest,);
let mut archive = layer.archive().await?;
let mut entries = archive.entries()?;
let mut count = 0u64;
let mut size = 0u64;
while let Some(entry) = entries.next().await {
let mut entry = entry?;
let path = entry.path()?;
@ -134,14 +136,21 @@ impl OciImageAssembler {
self.process_whiteout_entry(&mut vfs, &entry, name, layer)
.await?;
} else {
vfs.insert_tar_entry(&entry)?;
self.process_write_entry(&mut vfs, &mut entry, layer)
let reference = vfs.insert_tar_entry(&entry)?;
self.progress
.update(|progress| {
progress.extracting_layer(&layer.digest, &reference.name);
})
.await;
size += self
.process_write_entry(&mut vfs, &mut entry, layer)
.await?;
count += 1;
}
}
self.progress
.update(|progress| {
progress.extracted_layer(&layer.digest);
progress.extracted_layer(&layer.digest, count, size);
})
.await;
}
@ -157,6 +166,7 @@ impl OciImageAssembler {
let assembled = OciImageAssembled {
vfs: Arc::new(vfs),
descriptor: resolved.descriptor,
digest: resolved.digest,
manifest: resolved.manifest,
config: local.config,
@ -169,7 +179,7 @@ impl OciImageAssembler {
async fn process_whiteout_entry(
&self,
vfs: &mut VfsTree,
entry: &Entry<Archive<Pin<Box<dyn AsyncRead + Send>>>>,
entry: &Entry<Archive<Pin<Box<dyn OciImageLayerReader + Send>>>>,
name: &str,
layer: &OciImageLayer,
) -> Result<()> {
@ -210,11 +220,11 @@ impl OciImageAssembler {
async fn process_write_entry(
&self,
vfs: &mut VfsTree,
entry: &mut Entry<Archive<Pin<Box<dyn AsyncRead + Send>>>>,
entry: &mut Entry<Archive<Pin<Box<dyn OciImageLayerReader + Send>>>>,
layer: &OciImageLayer,
) -> Result<()> {
) -> Result<u64> {
if !entry.header().entry_type().is_file() {
return Ok(());
return Ok(0);
}
trace!(
"unpack entry layer={} path={:?} type={:?}",
@ -230,7 +240,7 @@ impl OciImageAssembler {
.await?
.ok_or(anyhow!("unpack did not return a path"))?;
vfs.set_disk_path(&entry.path()?, &path)?;
Ok(())
Ok(entry.header().size()?)
}
}

72
crates/oci/src/fetch.rs
View File

@ -10,6 +10,8 @@ use super::{
use std::{
fmt::Debug,
io::SeekFrom,
os::unix::fs::MetadataExt,
path::{Path, PathBuf},
pin::Pin,
};
@ -18,12 +20,13 @@ use anyhow::{anyhow, Result};
use async_compression::tokio::bufread::{GzipDecoder, ZstdDecoder};
use log::debug;
use oci_spec::image::{
Descriptor, ImageConfiguration, ImageIndex, ImageManifest, MediaType, ToDockerV2S2,
Descriptor, DescriptorBuilder, ImageConfiguration, ImageIndex, ImageManifest, MediaType,
ToDockerV2S2,
};
use serde::de::DeserializeOwned;
use tokio::{
fs::File,
io::{AsyncRead, AsyncReadExt, BufReader, BufWriter},
fs::{self, File},
io::{AsyncRead, AsyncReadExt, AsyncSeekExt, BufReader, BufWriter},
};
use tokio_stream::StreamExt;
use tokio_tar::Archive;
@ -43,16 +46,43 @@ pub enum OciImageLayerCompression {
#[derive(Clone, Debug)]
pub struct OciImageLayer {
pub metadata: Descriptor,
pub path: PathBuf,
pub digest: String,
pub compression: OciImageLayerCompression,
}
#[async_trait::async_trait]
pub trait OciImageLayerReader: AsyncRead + Sync {
async fn position(&mut self) -> Result<u64>;
}
#[async_trait::async_trait]
impl OciImageLayerReader for BufReader<File> {
async fn position(&mut self) -> Result<u64> {
Ok(self.seek(SeekFrom::Current(0)).await?)
}
}
#[async_trait::async_trait]
impl OciImageLayerReader for GzipDecoder<BufReader<File>> {
async fn position(&mut self) -> Result<u64> {
self.get_mut().position().await
}
}
#[async_trait::async_trait]
impl OciImageLayerReader for ZstdDecoder<BufReader<File>> {
async fn position(&mut self) -> Result<u64> {
self.get_mut().position().await
}
}
impl OciImageLayer {
pub async fn decompress(&self) -> Result<Pin<Box<dyn AsyncRead + Send>>> {
pub async fn decompress(&self) -> Result<Pin<Box<dyn OciImageLayerReader + Send>>> {
let file = File::open(&self.path).await?;
let reader = BufReader::new(file);
let reader: Pin<Box<dyn AsyncRead + Send>> = match self.compression {
let reader: Pin<Box<dyn OciImageLayerReader + Send>> = match self.compression {
OciImageLayerCompression::None => Box::pin(reader),
OciImageLayerCompression::Gzip => Box::pin(GzipDecoder::new(reader)),
OciImageLayerCompression::Zstd => Box::pin(ZstdDecoder::new(reader)),
@ -60,7 +90,7 @@ impl OciImageLayer {
Ok(reader)
}
pub async fn archive(&self) -> Result<Archive<Pin<Box<dyn AsyncRead + Send>>>> {
pub async fn archive(&self) -> Result<Archive<Pin<Box<dyn OciImageLayerReader + Send>>>> {
let decompress = self.decompress().await?;
Ok(Archive::new(decompress))
}
@ -70,6 +100,7 @@ impl OciImageLayer {
pub struct OciResolvedImage {
pub name: ImageName,
pub digest: String,
pub descriptor: Descriptor,
pub manifest: OciSchema<ImageManifest>,
}
@ -199,6 +230,7 @@ impl OciImageFetcher {
);
return Ok(OciResolvedImage {
name: image,
descriptor: found.clone(),
digest: found.digest().clone(),
manifest,
});
@ -207,11 +239,20 @@ impl OciImageFetcher {
}
let mut client = OciRegistryClient::new(image.registry_url()?, self.platform.clone())?;
let (manifest, digest) = client
let (manifest, descriptor, digest) = client
.get_manifest_with_digest(&image.name, &image.reference)
.await?;
let descriptor = descriptor.unwrap_or_else(|| {
DescriptorBuilder::default()
.media_type(MediaType::ImageManifest)
.size(manifest.raw().len() as i64)
.digest(digest.clone())
.build()
.unwrap()
});
Ok(OciResolvedImage {
name: image,
descriptor,
digest,
manifest,
})
@ -225,7 +266,7 @@ impl OciImageFetcher {
let config: OciSchema<ImageConfiguration>;
self.progress
.update(|progress| {
progress.phase = OciProgressPhase::ConfigAcquire;
progress.phase = OciProgressPhase::ConfigDownload;
})
.await;
let mut client = OciRegistryClient::new(image.name.registry_url()?, self.platform.clone())?;
@ -245,10 +286,10 @@ impl OciImageFetcher {
}
self.progress
.update(|progress| {
progress.phase = OciProgressPhase::LayerAcquire;
progress.phase = OciProgressPhase::LayerDownload;
for layer in image.manifest.item().layers() {
progress.add_layer(layer.digest(), layer.size() as usize);
progress.add_layer(layer.digest());
}
})
.await;
@ -256,7 +297,7 @@ impl OciImageFetcher {
for layer in image.manifest.item().layers() {
self.progress
.update(|progress| {
progress.downloading_layer(layer.digest(), 0, layer.size() as usize);
progress.downloading_layer(layer.digest(), 0, layer.size() as u64);
})
.await;
layers.push(
@ -265,7 +306,7 @@ impl OciImageFetcher {
);
self.progress
.update(|progress| {
progress.downloaded_layer(layer.digest());
progress.downloaded_layer(layer.digest(), layer.size() as u64);
})
.await;
}
@ -304,6 +345,12 @@ impl OciImageFetcher {
}
}
let metadata = fs::metadata(&layer_path).await?;
if layer.size() as u64 != metadata.size() {
return Err(anyhow!("layer size differs from size in manifest",));
}
let mut media_type = layer.media_type().clone();
// docker layer compatibility
@ -318,6 +365,7 @@ impl OciImageFetcher {
other => return Err(anyhow!("found layer with unknown media type: {}", other)),
};
Ok(OciImageLayer {
metadata: layer.clone(),
path: layer_path,
digest: layer.digest().clone(),
compression,

8
crates/oci/src/name.rs
View File

@ -15,7 +15,13 @@ pub struct ImageName {
impl fmt::Display for ImageName {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if let Some(port) = self.port {
if DOCKER_HUB_MIRROR == self.hostname && self.port.is_none() {
if self.name.starts_with("library/") {
write!(f, "{}:{}", &self.name[8..], self.reference)
} else {
write!(f, "{}:{}", self.name, self.reference)
}
} else if let Some(port) = self.port {
write!(
f,
"{}:{}/{}:{}",

48
crates/oci/src/packer/backend.rs
View File

@ -1,14 +1,12 @@
use std::{path::Path, process::Stdio, sync::Arc};
use std::{os::unix::fs::MetadataExt, path::Path, process::Stdio, sync::Arc};
use super::OciPackedFormat;
use crate::{
progress::{OciBoundProgress, OciProgressPhase},
vfs::VfsTree,
};
use crate::{progress::OciBoundProgress, vfs::VfsTree};
use anyhow::{anyhow, Result};
use log::warn;
use tokio::{
fs::File,
fs::{self, File},
io::BufWriter,
pin,
process::{Child, Command},
select,
@ -55,9 +53,7 @@ impl OciPackerBackend for OciPackerMkSquashfs {
async fn pack(&self, progress: OciBoundProgress, vfs: Arc<VfsTree>, file: &Path) -> Result<()> {
progress
.update(|progress| {
progress.phase = OciProgressPhase::Packing;
progress.total = 1;
progress.value = 0;
progress.start_packing();
})
.await;
@ -120,12 +116,9 @@ impl OciPackerBackend for OciPackerMkSquashfs {
status.code().unwrap()
))
} else {
let metadata = fs::metadata(&file).await?;
progress
.update(|progress| {
progress.phase = OciProgressPhase::Packing;
progress.total = 1;
progress.value = 1;
})
.update(|progress| progress.complete(metadata.size()))
.await;
Ok(())
}
@ -136,12 +129,10 @@ pub struct OciPackerMkfsErofs {}
#[async_trait::async_trait]
impl OciPackerBackend for OciPackerMkfsErofs {
async fn pack(&self, progress: OciBoundProgress, vfs: Arc<VfsTree>, path: &Path) -> Result<()> {
async fn pack(&self, progress: OciBoundProgress, vfs: Arc<VfsTree>, file: &Path) -> Result<()> {
progress
.update(|progress| {
progress.phase = OciProgressPhase::Packing;
progress.total = 1;
progress.value = 0;
progress.start_packing();
})
.await;
@ -149,7 +140,7 @@ impl OciPackerBackend for OciPackerMkfsErofs {
.arg("-L")
.arg("root")
.arg("--tar=-")
.arg(path)
.arg(file)
.stdin(Stdio::piped())
.stderr(Stdio::null())
.stdout(Stdio::null())
@ -200,11 +191,10 @@ impl OciPackerBackend for OciPackerMkfsErofs {
status.code().unwrap()
))
} else {
let metadata = fs::metadata(&file).await?;
progress
.update(|progress| {
progress.phase = OciProgressPhase::Packing;
progress.total = 1;
progress.value = 1;
progress.complete(metadata.size());
})
.await;
Ok(())
@ -219,20 +209,18 @@ impl OciPackerBackend for OciPackerTar {
async fn pack(&self, progress: OciBoundProgress, vfs: Arc<VfsTree>, file: &Path) -> Result<()> {
progress
.update(|progress| {
progress.phase = OciProgressPhase::Packing;
progress.total = 1;
progress.value = 0;
progress.start_packing();
})
.await;
let file = File::create(file).await?;
vfs.write_to_tar(file).await?;
let output = File::create(file).await?;
let output = BufWriter::new(output);
vfs.write_to_tar(output).await?;
let metadata = fs::metadata(file).await?;
progress
.update(|progress| {
progress.phase = OciProgressPhase::Packing;
progress.total = 1;
progress.value = 1;
progress.complete(metadata.size());
})
.await;
Ok(())

204
crates/oci/src/packer/cache.rs
View File

@ -1,69 +1,121 @@
use crate::{
name::ImageName,
packer::{OciPackedFormat, OciPackedImage},
schema::OciSchema,
};
use anyhow::Result;
use log::debug;
use oci_spec::image::{ImageConfiguration, ImageManifest};
use std::path::{Path, PathBuf};
use tokio::fs;
use log::{debug, error};
use oci_spec::image::{
Descriptor, ImageConfiguration, ImageIndex, ImageIndexBuilder, ImageManifest, MediaType,
ANNOTATION_REF_NAME,
};
use std::{
path::{Path, PathBuf},
sync::Arc,
};
use tokio::{fs, sync::RwLock};
#[derive(Clone)]
pub struct OciPackerCache {
cache_dir: PathBuf,
index: Arc<RwLock<ImageIndex>>,
}
const ANNOTATION_IMAGE_NAME: &str = "io.containerd.image.name";
const ANNOTATION_OCI_PACKER_FORMAT: &str = "dev.krata.oci.packer.format";
impl OciPackerCache {
pub fn new(cache_dir: &Path) -> Result<OciPackerCache> {
Ok(OciPackerCache {
pub async fn new(cache_dir: &Path) -> Result<OciPackerCache> {
let index = ImageIndexBuilder::default()
.schema_version(2u32)
.media_type(MediaType::ImageIndex)
.manifests(Vec::new())
.build()?;
let cache = OciPackerCache {
cache_dir: cache_dir.to_path_buf(),
})
index: Arc::new(RwLock::new(index)),
};
{
let mut mutex = cache.index.write().await;
*mutex = cache.load_index().await?;
}
Ok(cache)
}
pub async fn list(&self) -> Result<Vec<Descriptor>> {
let index = self.index.read().await;
Ok(index.manifests().clone())
}
pub async fn recall(
&self,
name: ImageName,
digest: &str,
format: OciPackedFormat,
) -> Result<Option<OciPackedImage>> {
let index = self.index.read().await;
let mut descriptor: Option<Descriptor> = None;
for manifest in index.manifests() {
if manifest.digest() == digest
&& manifest
.annotations()
.as_ref()
.and_then(|x| x.get(ANNOTATION_OCI_PACKER_FORMAT))
.map(|x| x.as_str())
== Some(format.extension())
{
descriptor = Some(manifest.clone());
break;
}
}
let Some(descriptor) = descriptor else {
return Ok(None);
};
let mut fs_path = self.cache_dir.clone();
let mut config_path = self.cache_dir.clone();
let mut manifest_path = self.cache_dir.clone();
fs_path.push(format!("{}.{}", digest, format.extension()));
manifest_path.push(format!("{}.manifest.json", digest));
config_path.push(format!("{}.config.json", digest));
Ok(
if fs_path.exists() && manifest_path.exists() && config_path.exists() {
let image_metadata = fs::metadata(&fs_path).await?;
let manifest_metadata = fs::metadata(&manifest_path).await?;
let config_metadata = fs::metadata(&config_path).await?;
if image_metadata.is_file()
&& manifest_metadata.is_file()
&& config_metadata.is_file()
{
let manifest_bytes = fs::read(&manifest_path).await?;
let manifest: ImageManifest = serde_json::from_slice(&manifest_bytes)?;
let config_bytes = fs::read(&config_path).await?;
let config: ImageConfiguration = serde_json::from_slice(&config_bytes)?;
debug!("cache hit digest={}", digest);
Some(OciPackedImage::new(
digest.to_string(),
fs_path.clone(),
format,
OciSchema::new(config_bytes, config),
OciSchema::new(manifest_bytes, manifest),
))
} else {
None
}
if fs_path.exists() && manifest_path.exists() && config_path.exists() {
let image_metadata = fs::metadata(&fs_path).await?;
let manifest_metadata = fs::metadata(&manifest_path).await?;
let config_metadata = fs::metadata(&config_path).await?;
if image_metadata.is_file() && manifest_metadata.is_file() && config_metadata.is_file()
{
let manifest_bytes = fs::read(&manifest_path).await?;
let manifest: ImageManifest = serde_json::from_slice(&manifest_bytes)?;
let config_bytes = fs::read(&config_path).await?;
let config: ImageConfiguration = serde_json::from_slice(&config_bytes)?;
debug!("cache hit digest={}", digest);
Ok(Some(OciPackedImage::new(
name,
digest.to_string(),
fs_path.clone(),
format,
descriptor,
OciSchema::new(config_bytes, config),
OciSchema::new(manifest_bytes, manifest),
)))
} else {
debug!("cache miss digest={}", digest);
None
},
)
Ok(None)
}
} else {
debug!("cache miss digest={}", digest);
Ok(None)
}
}
pub async fn store(&self, packed: OciPackedImage) -> Result<OciPackedImage> {
let mut index = self.index.write().await;
let mut manifests = index.manifests().clone();
debug!("cache store digest={}", packed.digest);
let mut fs_path = self.cache_dir.clone();
let mut manifest_path = self.cache_dir.clone();
@ -74,12 +126,90 @@ impl OciPackerCache {
fs::rename(&packed.path, &fs_path).await?;
fs::write(&config_path, packed.config.raw()).await?;
fs::write(&manifest_path, packed.manifest.raw()).await?;
Ok(OciPackedImage::new(
manifests.retain(|item| {
if item.digest() != &packed.digest {
return true;
}
let Some(format) = item
.annotations()
.as_ref()
.and_then(|x| x.get(ANNOTATION_OCI_PACKER_FORMAT))
.map(|x| x.as_str())
else {
return true;
};
if format != packed.format.extension() {
return true;
}
false
});
let mut descriptor = packed.descriptor.clone();
let mut annotations = descriptor.annotations().clone().unwrap_or_default();
annotations.insert(
ANNOTATION_OCI_PACKER_FORMAT.to_string(),
packed.format.extension().to_string(),
);
let image_name = packed.name.to_string();
annotations.insert(ANNOTATION_IMAGE_NAME.to_string(), image_name);
let image_ref = packed.name.reference.clone();
annotations.insert(ANNOTATION_REF_NAME.to_string(), image_ref);
descriptor.set_annotations(Some(annotations));
manifests.push(descriptor.clone());
index.set_manifests(manifests);
self.save_index(&index).await?;
let packed = OciPackedImage::new(
packed.name,
packed.digest,
fs_path.clone(),
packed.format,
descriptor,
packed.config,
packed.manifest,
))
);
Ok(packed)
}
async fn save_empty_index(&self) -> Result<ImageIndex> {
let index = ImageIndexBuilder::default()
.schema_version(2u32)
.media_type(MediaType::ImageIndex)
.manifests(Vec::new())
.build()?;
self.save_index(&index).await?;
Ok(index)
}
async fn load_index(&self) -> Result<ImageIndex> {
let mut index_path = self.cache_dir.clone();
index_path.push("index.json");
if !index_path.exists() {
self.save_empty_index().await?;
}
let content = fs::read_to_string(&index_path).await?;
let index = match serde_json::from_str::<ImageIndex>(&content) {
Ok(index) => index,
Err(error) => {
error!("image index was corrupted, creating a new one: {}", error);
self.save_empty_index().await?
}
};
Ok(index)
}
async fn save_index(&self, index: &ImageIndex) -> Result<()> {
let mut encoded = serde_json::to_string_pretty(index)?;
encoded.push('\n');
let mut index_path = self.cache_dir.clone();
index_path.push("index.json");
fs::write(&index_path, encoded).await?;
Ok(())
}
}

10
crates/oci/src/packer/mod.rs
View File

@ -1,9 +1,9 @@
use std::path::PathBuf;
use crate::schema::OciSchema;
use crate::{name::ImageName, schema::OciSchema};
use self::backend::OciPackerBackendType;
use oci_spec::image::{ImageConfiguration, ImageManifest};
use oci_spec::image::{Descriptor, ImageConfiguration, ImageManifest};
pub mod backend;
pub mod cache;
@ -37,25 +37,31 @@ impl OciPackedFormat {
#[derive(Clone)]
pub struct OciPackedImage {
pub name: ImageName,
pub digest: String,
pub path: PathBuf,
pub format: OciPackedFormat,
pub descriptor: Descriptor,
pub config: OciSchema<ImageConfiguration>,
pub manifest: OciSchema<ImageManifest>,
}
impl OciPackedImage {
pub fn new(
name: ImageName,
digest: String,
path: PathBuf,
format: OciPackedFormat,
descriptor: Descriptor,
config: OciSchema<ImageConfiguration>,
manifest: OciSchema<ImageManifest>,
) -> OciPackedImage {
OciPackedImage {
name,
digest,
path,
format,
descriptor,
config,
manifest,
}

62
crates/oci/src/packer/service.rs
View File

@ -6,6 +6,7 @@ use std::{
};
use anyhow::{anyhow, Result};
use oci_spec::image::Descriptor;
use tokio::{
sync::{watch, Mutex},
task::JoinHandle,
@ -38,38 +39,45 @@ pub struct OciPackerService {
}
impl OciPackerService {
pub fn new(
pub async fn new(
seed: Option<PathBuf>,
cache_dir: &Path,
platform: OciPlatform,
) -> Result<OciPackerService> {
Ok(OciPackerService {
seed,
cache: OciPackerCache::new(cache_dir)?,
cache: OciPackerCache::new(cache_dir).await?,
platform,
tasks: Arc::new(Mutex::new(HashMap::new())),
})
}
pub async fn list(&self) -> Result<Vec<Descriptor>> {
self.cache.list().await
}
pub async fn recall(
&self,
digest: &str,
format: OciPackedFormat,
) -> Result<Option<OciPackedImage>> {
self.cache.recall(digest, format).await
self.cache
.recall(ImageName::parse("cached:latest")?, digest, format)
.await
}
pub async fn request(
&self,
name: ImageName,
format: OciPackedFormat,
overwrite: bool,
progress_context: OciProgressContext,
) -> Result<OciPackedImage> {
let progress = OciProgress::new();
let progress = OciBoundProgress::new(progress_context.clone(), progress);
let fetcher =
OciImageFetcher::new(self.seed.clone(), self.platform.clone(), progress.clone());
let resolved = fetcher.resolve(name).await?;
let resolved = fetcher.resolve(name.clone()).await?;
let key = OciPackerTaskKey {
digest: resolved.digest.clone(),
format,
@ -86,7 +94,15 @@ impl OciPackerService {
Entry::Vacant(entry) => {
let task = self
.clone()
.launch(key.clone(), format, resolved, fetcher, progress.clone())
.launch(
name,
key.clone(),
format,
overwrite,
resolved,
fetcher,
progress.clone(),
)
.await;
let (watch, receiver) = watch::channel(None);
@ -122,22 +138,33 @@ impl OciPackerService {
}
}
#[allow(clippy::too_many_arguments)]
async fn launch(
self,
name: ImageName,
key: OciPackerTaskKey,
format: OciPackedFormat,
overwrite: bool,
resolved: OciResolvedImage,
fetcher: OciImageFetcher,
progress: OciBoundProgress,
) -> JoinHandle<()> {
info!("packer task {} started", key);
info!("started packer task {}", key);
tokio::task::spawn(async move {
let _task_drop_guard =
scopeguard::guard((key.clone(), self.clone()), |(key, service)| {
service.ensure_task_gone(key);
});
if let Err(error) = self
.task(key.clone(), format, resolved, fetcher, progress)
.task(
name,
key.clone(),
format,
overwrite,
resolved,
fetcher,
progress,
)
.await
{
self.finish(&key, Err(error)).await;
@ -145,17 +172,26 @@ impl OciPackerService {
})
}
#[allow(clippy::too_many_arguments)]
async fn task(
&self,
name: ImageName,
key: OciPackerTaskKey,
format: OciPackedFormat,
overwrite: bool,
resolved: OciResolvedImage,
fetcher: OciImageFetcher,
progress: OciBoundProgress,
) -> Result<()> {
if let Some(cached) = self.cache.recall(&resolved.digest, format).await? {
self.finish(&key, Ok(cached)).await;
return Ok(());
if !overwrite {
if let Some(cached) = self
.cache
.recall(name.clone(), &resolved.digest, format)
.await?
{
self.finish(&key, Ok(cached)).await;
return Ok(());
}
}
let assembler =
OciImageAssembler::new(fetcher, resolved, progress.clone(), None, None).await?;
@ -171,9 +207,11 @@ impl OciPackerService {
.pack(progress, assembled.vfs.clone(), &target)
.await?;
let packed = OciPackedImage::new(
name,
assembled.digest.clone(),
file,
format,
assembled.descriptor.clone(),
assembled.config.clone(),
assembled.manifest.clone(),
);
@ -190,7 +228,7 @@ impl OciPackerService {
match result.as_ref() {
Ok(_) => {
info!("packer task {} completed", key);
info!("completed packer task {}", key);
}
Err(err) => {
@ -216,7 +254,7 @@ impl OciPackerService {
tokio::task::spawn(async move {
let mut tasks = self.tasks.lock().await;
if let Some(task) = tasks.remove(&key) {
warn!("packer task {} aborted", key);
warn!("aborted packer task {}", key);
task.watch.send_replace(Some(Err(anyhow!("task aborted"))));
}
});

145
crates/oci/src/progress.rs
View File

@ -1,18 +1,16 @@
use indexmap::IndexMap;
use std::sync::Arc;
use tokio::{
sync::{broadcast, Mutex},
sync::{watch, Mutex},
task::JoinHandle,
};
const OCI_PROGRESS_QUEUE_LEN: usize = 100;
#[derive(Clone, Debug)]
pub struct OciProgress {
pub phase: OciProgressPhase,
pub digest: Option<String>,
pub layers: IndexMap<String, OciProgressLayer>,
pub value: u64,
pub total: u64,
pub indication: OciProgressIndication,
}
impl Default for OciProgress {
@ -24,72 +22,146 @@ impl Default for OciProgress {
impl OciProgress {
pub fn new() -> Self {
OciProgress {
phase: OciProgressPhase::Resolving,
phase: OciProgressPhase::Started,
digest: None,
layers: IndexMap::new(),
value: 0,
total: 1,
indication: OciProgressIndication::Hidden,
}
}
pub fn add_layer(&mut self, id: &str, size: usize) {
pub fn start_resolving(&mut self) {
self.phase = OciProgressPhase::Resolving;
self.indication = OciProgressIndication::Spinner { message: None };
}
pub fn resolved(&mut self, digest: &str) {
self.digest = Some(digest.to_string());
self.indication = OciProgressIndication::Hidden;
}
pub fn add_layer(&mut self, id: &str) {
self.layers.insert(
id.to_string(),
OciProgressLayer {
id: id.to_string(),
phase: OciProgressLayerPhase::Waiting,
value: 0,
total: size as u64,
indication: OciProgressIndication::Spinner { message: None },
},
);
}
pub fn downloading_layer(&mut self, id: &str, downloaded: usize, total: usize) {
pub fn downloading_layer(&mut self, id: &str, downloaded: u64, total: u64) {
if let Some(entry) = self.layers.get_mut(id) {
entry.phase = OciProgressLayerPhase::Downloading;
entry.value = downloaded as u64;
entry.total = total as u64;
entry.indication = OciProgressIndication::ProgressBar {
message: None,
current: downloaded,
total,
bytes: true,
};
}
}
pub fn downloaded_layer(&mut self, id: &str) {
pub fn downloaded_layer(&mut self, id: &str, total: u64) {
if let Some(entry) = self.layers.get_mut(id) {
entry.phase = OciProgressLayerPhase::Downloaded;
entry.value = entry.total;
entry.indication = OciProgressIndication::Completed {
message: None,
total: Some(total),
bytes: true,
};
}
}
pub fn extracting_layer(&mut self, id: &str, extracted: usize, total: usize) {
pub fn start_assemble(&mut self) {
self.phase = OciProgressPhase::Assemble;
self.indication = OciProgressIndication::Hidden;
}
pub fn start_extracting_layer(&mut self, id: &str) {
if let Some(entry) = self.layers.get_mut(id) {
entry.phase = OciProgressLayerPhase::Extracting;
entry.value = extracted as u64;
entry.total = total as u64;
entry.indication = OciProgressIndication::Spinner { message: None };
}
}
pub fn extracted_layer(&mut self, id: &str) {
pub fn extracting_layer(&mut self, id: &str, file: &str) {
if let Some(entry) = self.layers.get_mut(id) {
entry.phase = OciProgressLayerPhase::Extracting;
entry.indication = OciProgressIndication::Spinner {
message: Some(file.to_string()),
};
}
}
pub fn extracted_layer(&mut self, id: &str, count: u64, total_size: u64) {
if let Some(entry) = self.layers.get_mut(id) {
entry.phase = OciProgressLayerPhase::Extracted;
entry.value = entry.total;
entry.indication = OciProgressIndication::Completed {
message: Some(format!("{} files", count)),
total: Some(total_size),
bytes: true,
};
}
}
pub fn start_packing(&mut self) {
self.phase = OciProgressPhase::Pack;
for layer in self.layers.values_mut() {
layer.indication = OciProgressIndication::Hidden;
}
self.indication = OciProgressIndication::Spinner { message: None };
}
pub fn complete(&mut self, size: u64) {
self.phase = OciProgressPhase::Complete;
self.indication = OciProgressIndication::Completed {
message: None,
total: Some(size),
bytes: true,
}
}
}
#[derive(Clone, Debug)]
pub enum OciProgressPhase {
Started,
Resolving,
Resolved,
ConfigAcquire,
LayerAcquire,
Packing,
ConfigDownload,
LayerDownload,
Assemble,
Pack,
Complete,
}
#[derive(Clone, Debug)]
pub enum OciProgressIndication {
Hidden,
ProgressBar {
message: Option<String>,
current: u64,
total: u64,
bytes: bool,
},
Spinner {
message: Option<String>,
},
Completed {
message: Option<String>,
total: Option<u64>,
bytes: bool,
},
}
#[derive(Clone, Debug)]
pub struct OciProgressLayer {
pub id: String,
pub phase: OciProgressLayerPhase,
pub value: u64,
pub total: u64,
pub indication: OciProgressIndication,
}
#[derive(Clone, Debug)]
@ -103,16 +175,16 @@ pub enum OciProgressLayerPhase {
#[derive(Clone)]
pub struct OciProgressContext {
sender: broadcast::Sender<OciProgress>,
sender: watch::Sender<OciProgress>,
}
impl OciProgressContext {
pub fn create() -> (OciProgressContext, broadcast::Receiver<OciProgress>) {
let (sender, receiver) = broadcast::channel(OCI_PROGRESS_QUEUE_LEN);
pub fn create() -> (OciProgressContext, watch::Receiver<OciProgress>) {
let (sender, receiver) = watch::channel(OciProgress::new());
(OciProgressContext::new(sender), receiver)
}
pub fn new(sender: broadcast::Sender<OciProgress>) -> OciProgressContext {
pub fn new(sender: watch::Sender<OciProgress>) -> OciProgressContext {
OciProgressContext { sender }
}
@ -120,7 +192,7 @@ impl OciProgressContext {
let _ = self.sender.send(progress.clone());
}
pub fn subscribe(&self) -> broadcast::Receiver<OciProgress> {
pub fn subscribe(&self) -> watch::Receiver<OciProgress> {
self.sender.subscribe()
}
}
@ -156,13 +228,10 @@ impl OciBoundProgress {
context.update(&progress);
let mut receiver = self.context.subscribe();
tokio::task::spawn(async move {
while let Ok(progress) = receiver.recv().await {
match context.sender.send(progress) {
Ok(_) => {}
Err(_) => {
break;
}
}
while (receiver.changed().await).is_ok() {
context
.sender
.send_replace(receiver.borrow_and_update().clone());
}
})
}

33
crates/oci/src/registry.rs
View File

@ -149,24 +149,20 @@ impl OciRegistryClient {
))?;
let mut response = self.call(self.agent.get(url.as_str())).await?;
let mut size: u64 = 0;
let mut last_progress_size: u64 = 0;
while let Some(chunk) = response.chunk().await? {
dest.write_all(&chunk).await?;
size += chunk.len() as u64;
if (size - last_progress_size) > (5 * 1024 * 1024) {
last_progress_size = size;
if let Some(ref progress) = progress {
progress
.update(|progress| {
progress.downloading_layer(
descriptor.digest(),
size as usize,
descriptor.size() as usize,
);
})
.await;
}
if let Some(ref progress) = progress {
progress
.update(|progress| {
progress.downloading_layer(
descriptor.digest(),
size,
descriptor.size() as u64,
);
})
.await;
}
}
Ok(size)
@ -207,7 +203,7 @@ impl OciRegistryClient {
&mut self,
name: N,
reference: R,
) -> Result<(OciSchema<ImageManifest>, String)> {
) -> Result<(OciSchema<ImageManifest>, Option<Descriptor>, String)> {
let url = self.url.join(&format!(
"/v2/{}/manifests/{}",
name.as_ref(),
@ -235,9 +231,10 @@ impl OciRegistryClient {
let descriptor = self
.pick_manifest(index)
.ok_or_else(|| anyhow!("unable to pick manifest from index"))?;
return self
let (manifest, digest) = self
.get_raw_manifest_with_digest(name, descriptor.digest())
.await;
.await?;
return Ok((manifest, Some(descriptor), digest));
}
let digest = response
.headers()
@ -247,7 +244,7 @@ impl OciRegistryClient {
.to_string();
let bytes = response.bytes().await?;
let manifest = serde_json::from_slice(&bytes)?;
Ok((OciSchema::new(bytes.to_vec(), manifest), digest))
Ok((OciSchema::new(bytes.to_vec(), manifest), None, digest))
}
fn pick_manifest(&mut self, index: ImageIndex) -> Option<Descriptor> {

9
crates/oci/src/vfs.rs
View File

@ -194,7 +194,7 @@ impl VfsTree {
}
}
pub fn insert_tar_entry<X: AsyncRead + Unpin>(&mut self, entry: &Entry<X>) -> Result<()> {
pub fn insert_tar_entry<X: AsyncRead + Unpin>(&mut self, entry: &Entry<X>) -> Result<&VfsNode> {
let mut meta = VfsNode::from(entry)?;
let path = entry.path()?.to_path_buf();
let parent = if let Some(parent) = path.parent() {
@ -218,8 +218,11 @@ impl VfsTree {
meta.children = old.children;
}
}
parent.children.push(meta);
Ok(())
parent.children.push(meta.clone());
let Some(reference) = parent.children.iter().find(|child| child.name == meta.name) else {
return Err(anyhow!("unable to find inserted child in vfs"));
};
Ok(reference)
}
pub fn set_disk_path(&mut self, path: &Path, disk_path: &Path) -> Result<()> {