feat: concurrent image pulls

This commit is contained in:
Alex Zenla
2024-04-16 08:20:09 +00:00
parent 3d29a68645
commit a99541460d
13 changed files with 331 additions and 73 deletions

1
Cargo.lock generated
View File

@ -1452,6 +1452,7 @@ dependencies = [
"log", "log",
"prost", "prost",
"redb", "redb",
"scopeguard",
"signal-hook", "signal-hook",
"tokio", "tokio",
"tokio-stream", "tokio-stream",

View File

@ -23,6 +23,7 @@ krata-runtime = { path = "../runtime", version = "^0.0.9" }
log = { workspace = true } log = { workspace = true }
prost = { workspace = true } prost = { workspace = true }
redb = { workspace = true } redb = { workspace = true }
scopeguard = { workspace = true }
signal-hook = { workspace = true } signal-hook = { workspace = true }
tokio = { workspace = true } tokio = { workspace = true }
tokio-stream = { workspace = true } tokio-stream = { workspace = true }

View File

@ -18,7 +18,7 @@ use krata::{
}; };
use krataoci::{ use krataoci::{
name::ImageName, name::ImageName,
packer::{service::OciPackerService, OciImagePacked, OciPackedFormat}, packer::{service::OciPackerService, OciPackedFormat, OciPackedImage},
progress::{OciProgress, OciProgressContext}, progress::{OciProgress, OciProgressContext},
}; };
use std::{pin::Pin, str::FromStr}; use std::{pin::Pin, str::FromStr};
@ -90,8 +90,8 @@ enum ConsoleDataSelect {
} }
enum PullImageSelect { enum PullImageSelect {
Progress(usize), Progress(Option<OciProgress>),
Completed(Result<Result<OciImagePacked, anyhow::Error>, JoinError>), Completed(Result<Result<OciPackedImage, anyhow::Error>, JoinError>),
} }
#[tonic::async_trait] #[tonic::async_trait]
@ -367,32 +367,46 @@ impl ControlService for DaemonControlService {
OciImageFormat::Erofs => OciPackedFormat::Erofs, OciImageFormat::Erofs => OciPackedFormat::Erofs,
OciImageFormat::Tar => OciPackedFormat::Tar, OciImageFormat::Tar => OciPackedFormat::Tar,
}; };
let (sender, mut receiver) = channel::<OciProgress>(100); let (context, mut receiver) = OciProgressContext::create();
let context = OciProgressContext::new(sender);
let our_packer = self.packer.clone(); let our_packer = self.packer.clone();
let output = try_stream! { let output = try_stream! {
let mut task = tokio::task::spawn(async move { let mut task = tokio::task::spawn(async move {
our_packer.request(name, format, context).await our_packer.request(name, format, context).await
}); });
let abort_handle = task.abort_handle();
let _task_cancel_guard = scopeguard::guard(abort_handle, |handle| {
handle.abort();
});
loop { loop {
let mut progresses = Vec::new();
let what = select! { let what = select! {
x = receiver.recv_many(&mut progresses, 10) => PullImageSelect::Progress(x), x = receiver.recv() => PullImageSelect::Progress(x.ok()),
x = &mut task => PullImageSelect::Completed(x), x = &mut task => PullImageSelect::Completed(x),
}; };
match what { match what {
PullImageSelect::Progress(count) => { PullImageSelect::Progress(Some(mut progress)) => {
if count > 0 { let mut drain = 0;
let progress = progresses.remove(progresses.len() - 1); loop {
let reply = PullImageReply { if drain >= 10 {
progress: Some(convert_oci_progress(progress)), break;
digest: String::new(), }
format: OciImageFormat::Unknown.into(),
}; if let Ok(latest) = receiver.try_recv() {
yield reply; progress = latest;
} else {
break;
}
drain += 1;
} }
let reply = PullImageReply {
progress: Some(convert_oci_progress(progress)),
digest: String::new(),
format: OciImageFormat::Unknown.into(),
};
yield reply;
}, },
PullImageSelect::Completed(result) => { PullImageSelect::Completed(result) => {
@ -414,6 +428,10 @@ impl ControlService for DaemonControlService {
yield reply; yield reply;
break; break;
}, },
_ => {
continue;
}
} }
} }
}; };

View File

@ -5,10 +5,10 @@ use env_logger::Env;
use krataoci::{ use krataoci::{
name::ImageName, name::ImageName,
packer::{service::OciPackerService, OciPackedFormat}, packer::{service::OciPackerService, OciPackedFormat},
progress::{OciProgress, OciProgressContext}, progress::OciProgressContext,
registry::OciPlatform, registry::OciPlatform,
}; };
use tokio::{fs, sync::mpsc::channel}; use tokio::fs;
#[tokio::main] #[tokio::main]
async fn main() -> Result<()> { async fn main() -> Result<()> {
@ -22,14 +22,28 @@ async fn main() -> Result<()> {
fs::create_dir(&cache_dir).await?; fs::create_dir(&cache_dir).await?;
} }
let (sender, mut receiver) = channel::<OciProgress>(100); let (context, mut receiver) = OciProgressContext::create();
tokio::task::spawn(async move { tokio::task::spawn(async move {
loop { loop {
let mut progresses = Vec::new(); let Ok(mut progress) = receiver.recv().await else {
let _ = receiver.recv_many(&mut progresses, 100).await; return;
let Some(progress) = progresses.last() else {
continue;
}; };
let mut drain = 0;
loop {
if drain >= 10 {
break;
}
if let Ok(latest) = receiver.try_recv() {
progress = latest;
} else {
break;
}
drain += 1;
}
println!("phase {:?}", progress.phase); println!("phase {:?}", progress.phase);
for (id, layer) in &progress.layers { for (id, layer) in &progress.layers {
println!( println!(
@ -39,7 +53,6 @@ async fn main() -> Result<()> {
} }
} }
}); });
let context = OciProgressContext::new(sender);
let service = OciPackerService::new(seed, &cache_dir, OciPlatform::current())?; let service = OciPackerService::new(seed, &cache_dir, OciPlatform::current())?;
let packed = service let packed = service
.request(image.clone(), OciPackedFormat::Squashfs, context) .request(image.clone(), OciPackedFormat::Squashfs, context)

View File

@ -5,8 +5,10 @@ use crate::vfs::{VfsNode, VfsTree};
use anyhow::{anyhow, Result}; use anyhow::{anyhow, Result};
use log::{debug, trace, warn}; use log::{debug, trace, warn};
use oci_spec::image::{ImageConfiguration, ImageManifest}; use oci_spec::image::{ImageConfiguration, ImageManifest};
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::pin::Pin; use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc; use std::sync::Arc;
use tokio::fs; use tokio::fs;
use tokio::io::AsyncRead; use tokio::io::AsyncRead;
@ -34,11 +36,12 @@ impl Drop for OciImageAssembled {
pub struct OciImageAssembler { pub struct OciImageAssembler {
downloader: OciImageFetcher, downloader: OciImageFetcher,
resolved: OciResolvedImage, resolved: Option<OciResolvedImage>,
progress: OciBoundProgress, progress: OciBoundProgress,
work_dir: PathBuf, work_dir: PathBuf,
disk_dir: PathBuf, disk_dir: PathBuf,
tmp_dir: Option<PathBuf>, tmp_dir: Option<PathBuf>,
success: AtomicBool,
} }
impl OciImageAssembler { impl OciImageAssembler {
@ -82,11 +85,12 @@ impl OciImageAssembler {
Ok(OciImageAssembler { Ok(OciImageAssembler {
downloader, downloader,
resolved, resolved: Some(resolved),
progress, progress,
work_dir, work_dir,
disk_dir: target_dir, disk_dir: target_dir,
tmp_dir, tmp_dir,
success: AtomicBool::new(false),
}) })
} }
@ -98,11 +102,11 @@ impl OciImageAssembler {
self.assemble_with(&layer_dir).await self.assemble_with(&layer_dir).await
} }
async fn assemble_with(self, layer_dir: &Path) -> Result<OciImageAssembled> { async fn assemble_with(mut self, layer_dir: &Path) -> Result<OciImageAssembled> {
let local = self let Some(ref resolved) = self.resolved else {
.downloader return Err(anyhow!("resolved image was not available when expected"));
.download(self.resolved.clone(), layer_dir) };
.await?; let local = self.downloader.download(resolved, layer_dir).await?;
let mut vfs = VfsTree::new(); let mut vfs = VfsTree::new();
for layer in &local.layers { for layer in &local.layers {
debug!( debug!(
@ -146,13 +150,20 @@ impl OciImageAssembler {
fs::remove_file(&layer.path).await?; fs::remove_file(&layer.path).await?;
} }
} }
Ok(OciImageAssembled {
let Some(resolved) = self.resolved.take() else {
return Err(anyhow!("resolved image was not available when expected"));
};
let assembled = OciImageAssembled {
vfs: Arc::new(vfs), vfs: Arc::new(vfs),
digest: self.resolved.digest, digest: resolved.digest,
manifest: self.resolved.manifest, manifest: resolved.manifest,
config: local.config, config: local.config,
tmp_dir: self.tmp_dir, tmp_dir: self.tmp_dir.clone(),
}) };
self.success.store(true, Ordering::Release);
Ok(assembled)
} }
async fn process_whiteout_entry( async fn process_whiteout_entry(
@ -223,6 +234,18 @@ impl OciImageAssembler {
} }
} }
impl Drop for OciImageAssembler {
fn drop(&mut self) {
if !self.success.load(Ordering::Acquire) {
if let Some(tmp_dir) = self.tmp_dir.clone() {
tokio::task::spawn(async move {
let _ = fs::remove_dir_all(tmp_dir).await;
});
}
}
}
}
async fn delete_disk_paths(node: &VfsNode) -> Result<()> { async fn delete_disk_paths(node: &VfsNode) -> Result<()> {
let mut queue = vec![node]; let mut queue = vec![node];
while !queue.is_empty() { while !queue.is_empty() {

View File

@ -219,7 +219,7 @@ impl OciImageFetcher {
pub async fn download( pub async fn download(
&self, &self,
image: OciResolvedImage, image: &OciResolvedImage,
layer_dir: &Path, layer_dir: &Path,
) -> Result<OciLocalImage> { ) -> Result<OciLocalImage> {
let config: OciSchema<ImageConfiguration>; let config: OciSchema<ImageConfiguration>;
@ -270,7 +270,7 @@ impl OciImageFetcher {
.await; .await;
} }
Ok(OciLocalImage { Ok(OciLocalImage {
image, image: image.clone(),
config, config,
layers, layers,
}) })

View File

@ -7,7 +7,12 @@ use crate::{
}; };
use anyhow::{anyhow, Result}; use anyhow::{anyhow, Result};
use log::warn; use log::warn;
use tokio::{fs::File, pin, process::Command, select}; use tokio::{
fs::File,
pin,
process::{Child, Command},
select,
};
#[derive(Debug, Clone, Copy)] #[derive(Debug, Clone, Copy)]
pub enum OciPackerBackendType { pub enum OciPackerBackendType {
@ -56,7 +61,7 @@ impl OciPackerBackend for OciPackerMkSquashfs {
}) })
.await; .await;
let mut child = Command::new("mksquashfs") let child = Command::new("mksquashfs")
.arg("-") .arg("-")
.arg(file) .arg(file)
.arg("-comp") .arg("-comp")
@ -66,7 +71,9 @@ impl OciPackerBackend for OciPackerMkSquashfs {
.stderr(Stdio::null()) .stderr(Stdio::null())
.stdout(Stdio::null()) .stdout(Stdio::null())
.spawn()?; .spawn()?;
let mut child = ChildProcessKillGuard(child);
let stdin = child let stdin = child
.0
.stdin .stdin
.take() .take()
.ok_or(anyhow!("unable to acquire stdin stream"))?; .ok_or(anyhow!("unable to acquire stdin stream"))?;
@ -77,7 +84,7 @@ impl OciPackerBackend for OciPackerMkSquashfs {
} }
Ok(()) Ok(())
})); }));
let wait = child.wait(); let wait = child.0.wait();
pin!(wait); pin!(wait);
let status_result = loop { let status_result = loop {
if let Some(inner) = writer.as_mut() { if let Some(inner) = writer.as_mut() {
@ -138,7 +145,7 @@ impl OciPackerBackend for OciPackerMkfsErofs {
}) })
.await; .await;
let mut child = Command::new("mkfs.erofs") let child = Command::new("mkfs.erofs")
.arg("-L") .arg("-L")
.arg("root") .arg("root")
.arg("--tar=-") .arg("--tar=-")
@ -147,14 +154,16 @@ impl OciPackerBackend for OciPackerMkfsErofs {
.stderr(Stdio::null()) .stderr(Stdio::null())
.stdout(Stdio::null()) .stdout(Stdio::null())
.spawn()?; .spawn()?;
let mut child = ChildProcessKillGuard(child);
let stdin = child let stdin = child
.0
.stdin .stdin
.take() .take()
.ok_or(anyhow!("unable to acquire stdin stream"))?; .ok_or(anyhow!("unable to acquire stdin stream"))?;
let mut writer = Some(tokio::task::spawn( let mut writer = Some(tokio::task::spawn(
async move { vfs.write_to_tar(stdin).await }, async move { vfs.write_to_tar(stdin).await },
)); ));
let wait = child.wait(); let wait = child.0.wait();
pin!(wait); pin!(wait);
let status_result = loop { let status_result = loop {
if let Some(inner) = writer.as_mut() { if let Some(inner) = writer.as_mut() {
@ -229,3 +238,11 @@ impl OciPackerBackend for OciPackerTar {
Ok(()) Ok(())
} }
} }
struct ChildProcessKillGuard(Child);
impl Drop for ChildProcessKillGuard {
fn drop(&mut self) {
let _ = self.0.start_kill();
}
}

View File

@ -1,5 +1,5 @@
use crate::{ use crate::{
packer::{OciImagePacked, OciPackedFormat}, packer::{OciPackedFormat, OciPackedImage},
schema::OciSchema, schema::OciSchema,
}; };
@ -25,7 +25,7 @@ impl OciPackerCache {
&self, &self,
digest: &str, digest: &str,
format: OciPackedFormat, format: OciPackedFormat,
) -> Result<Option<OciImagePacked>> { ) -> Result<Option<OciPackedImage>> {
let mut fs_path = self.cache_dir.clone(); let mut fs_path = self.cache_dir.clone();
let mut config_path = self.cache_dir.clone(); let mut config_path = self.cache_dir.clone();
let mut manifest_path = self.cache_dir.clone(); let mut manifest_path = self.cache_dir.clone();
@ -46,7 +46,7 @@ impl OciPackerCache {
let config_bytes = fs::read(&config_path).await?; let config_bytes = fs::read(&config_path).await?;
let config: ImageConfiguration = serde_json::from_slice(&config_bytes)?; let config: ImageConfiguration = serde_json::from_slice(&config_bytes)?;
debug!("cache hit digest={}", digest); debug!("cache hit digest={}", digest);
Some(OciImagePacked::new( Some(OciPackedImage::new(
digest.to_string(), digest.to_string(),
fs_path.clone(), fs_path.clone(),
format, format,
@ -63,7 +63,7 @@ impl OciPackerCache {
) )
} }
pub async fn store(&self, packed: OciImagePacked) -> Result<OciImagePacked> { pub async fn store(&self, packed: OciPackedImage) -> Result<OciPackedImage> {
debug!("cache store digest={}", packed.digest); debug!("cache store digest={}", packed.digest);
let mut fs_path = self.cache_dir.clone(); let mut fs_path = self.cache_dir.clone();
let mut manifest_path = self.cache_dir.clone(); let mut manifest_path = self.cache_dir.clone();
@ -74,7 +74,7 @@ impl OciPackerCache {
fs::rename(&packed.path, &fs_path).await?; fs::rename(&packed.path, &fs_path).await?;
fs::write(&config_path, packed.config.raw()).await?; fs::write(&config_path, packed.config.raw()).await?;
fs::write(&manifest_path, packed.manifest.raw()).await?; fs::write(&manifest_path, packed.manifest.raw()).await?;
Ok(OciImagePacked::new( Ok(OciPackedImage::new(
packed.digest, packed.digest,
fs_path.clone(), fs_path.clone(),
packed.format, packed.format,

View File

@ -9,7 +9,7 @@ pub mod backend;
pub mod cache; pub mod cache;
pub mod service; pub mod service;
#[derive(Debug, Default, Clone, Copy)] #[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash)]
pub enum OciPackedFormat { pub enum OciPackedFormat {
#[default] #[default]
Squashfs, Squashfs,
@ -36,7 +36,7 @@ impl OciPackedFormat {
} }
#[derive(Clone)] #[derive(Clone)]
pub struct OciImagePacked { pub struct OciPackedImage {
pub digest: String, pub digest: String,
pub path: PathBuf, pub path: PathBuf,
pub format: OciPackedFormat, pub format: OciPackedFormat,
@ -44,15 +44,15 @@ pub struct OciImagePacked {
pub manifest: OciSchema<ImageManifest>, pub manifest: OciSchema<ImageManifest>,
} }
impl OciImagePacked { impl OciPackedImage {
pub fn new( pub fn new(
digest: String, digest: String,
path: PathBuf, path: PathBuf,
format: OciPackedFormat, format: OciPackedFormat,
config: OciSchema<ImageConfiguration>, config: OciSchema<ImageConfiguration>,
manifest: OciSchema<ImageManifest>, manifest: OciSchema<ImageManifest>,
) -> OciImagePacked { ) -> OciPackedImage {
OciImagePacked { OciPackedImage {
digest, digest,
path, path,
format, format,

View File

@ -1,22 +1,40 @@
use std::path::{Path, PathBuf}; use std::{
collections::{hash_map::Entry, HashMap},
fmt::Display,
path::{Path, PathBuf},
sync::Arc,
};
use anyhow::{anyhow, Result}; use anyhow::{anyhow, Result};
use tokio::{
sync::{watch, Mutex},
task::JoinHandle,
};
use crate::{ use crate::{
assemble::OciImageAssembler, assemble::OciImageAssembler,
fetch::OciImageFetcher, fetch::{OciImageFetcher, OciResolvedImage},
name::ImageName, name::ImageName,
progress::{OciBoundProgress, OciProgress, OciProgressContext}, progress::{OciBoundProgress, OciProgress, OciProgressContext},
registry::OciPlatform, registry::OciPlatform,
}; };
use super::{cache::OciPackerCache, OciImagePacked, OciPackedFormat}; use log::{error, info, warn};
use super::{cache::OciPackerCache, OciPackedFormat, OciPackedImage};
pub struct OciPackerTask {
progress: OciBoundProgress,
watch: watch::Sender<Option<Result<OciPackedImage>>>,
task: JoinHandle<()>,
}
#[derive(Clone)] #[derive(Clone)]
pub struct OciPackerService { pub struct OciPackerService {
seed: Option<PathBuf>, seed: Option<PathBuf>,
platform: OciPlatform, platform: OciPlatform,
cache: OciPackerCache, cache: OciPackerCache,
tasks: Arc<Mutex<HashMap<OciPackerTaskKey, OciPackerTask>>>,
} }
impl OciPackerService { impl OciPackerService {
@ -29,6 +47,7 @@ impl OciPackerService {
seed, seed,
cache: OciPackerCache::new(cache_dir)?, cache: OciPackerCache::new(cache_dir)?,
platform, platform,
tasks: Arc::new(Mutex::new(HashMap::new())),
}) })
} }
@ -36,7 +55,7 @@ impl OciPackerService {
&self, &self,
digest: &str, digest: &str,
format: OciPackedFormat, format: OciPackedFormat,
) -> Result<Option<OciImagePacked>> { ) -> Result<Option<OciPackedImage>> {
self.cache.recall(digest, format).await self.cache.recall(digest, format).await
} }
@ -45,14 +64,98 @@ impl OciPackerService {
name: ImageName, name: ImageName,
format: OciPackedFormat, format: OciPackedFormat,
progress_context: OciProgressContext, progress_context: OciProgressContext,
) -> Result<OciImagePacked> { ) -> Result<OciPackedImage> {
let progress = OciProgress::new(); let progress = OciProgress::new();
let progress = OciBoundProgress::new(progress_context.clone(), progress); let progress = OciBoundProgress::new(progress_context.clone(), progress);
let fetcher = let fetcher =
OciImageFetcher::new(self.seed.clone(), self.platform.clone(), progress.clone()); OciImageFetcher::new(self.seed.clone(), self.platform.clone(), progress.clone());
let resolved = fetcher.resolve(name).await?; let resolved = fetcher.resolve(name).await?;
let key = OciPackerTaskKey {
digest: resolved.digest.clone(),
format,
};
let (progress_copy_task, mut receiver) = match self.tasks.lock().await.entry(key.clone()) {
Entry::Occupied(entry) => {
let entry = entry.get();
(
Some(entry.progress.also_update(progress_context).await),
entry.watch.subscribe(),
)
}
Entry::Vacant(entry) => {
let task = self
.clone()
.launch(key.clone(), format, resolved, fetcher, progress.clone())
.await;
let (watch, receiver) = watch::channel(None);
let task = OciPackerTask {
progress: progress.clone(),
task,
watch,
};
entry.insert(task);
(None, receiver)
}
};
let _progress_task_guard = scopeguard::guard(progress_copy_task, |task| {
if let Some(task) = task {
task.abort();
}
});
let _task_cancel_guard = scopeguard::guard(self.clone(), |service| {
service.maybe_cancel_task(key);
});
loop {
receiver.changed().await?;
let current = receiver.borrow_and_update();
if current.is_some() {
return current
.as_ref()
.map(|x| x.as_ref().map_err(|err| anyhow!("{}", err)).cloned())
.unwrap();
}
}
}
async fn launch(
self,
key: OciPackerTaskKey,
format: OciPackedFormat,
resolved: OciResolvedImage,
fetcher: OciImageFetcher,
progress: OciBoundProgress,
) -> JoinHandle<()> {
info!("packer task {} started", 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)
.await
{
self.finish(&key, Err(error)).await;
}
})
}
async fn task(
&self,
key: OciPackerTaskKey,
format: OciPackedFormat,
resolved: OciResolvedImage,
fetcher: OciImageFetcher,
progress: OciBoundProgress,
) -> Result<()> {
if let Some(cached) = self.cache.recall(&resolved.digest, format).await? { if let Some(cached) = self.cache.recall(&resolved.digest, format).await? {
return Ok(cached); self.finish(&key, Ok(cached)).await;
return Ok(());
} }
let assembler = let assembler =
OciImageAssembler::new(fetcher, resolved, progress.clone(), None, None).await?; OciImageAssembler::new(fetcher, resolved, progress.clone(), None, None).await?;
@ -67,7 +170,7 @@ impl OciPackerService {
packer packer
.pack(progress, assembled.vfs.clone(), &target) .pack(progress, assembled.vfs.clone(), &target)
.await?; .await?;
let packed = OciImagePacked::new( let packed = OciPackedImage::new(
assembled.digest.clone(), assembled.digest.clone(),
file, file,
format, format,
@ -75,6 +178,59 @@ impl OciPackerService {
assembled.manifest.clone(), assembled.manifest.clone(),
); );
let packed = self.cache.store(packed).await?; let packed = self.cache.store(packed).await?;
Ok(packed) self.finish(&key, Ok(packed)).await;
Ok(())
}
async fn finish(&self, key: &OciPackerTaskKey, result: Result<OciPackedImage>) {
let Some(task) = self.tasks.lock().await.remove(key) else {
error!("packer task {} was not found when task completed", key);
return;
};
match result.as_ref() {
Ok(_) => {
info!("packer task {} completed", key);
}
Err(err) => {
warn!("packer task {} failed: {}", key, err);
}
}
task.watch.send_replace(Some(result));
}
fn maybe_cancel_task(self, key: OciPackerTaskKey) {
tokio::task::spawn(async move {
let tasks = self.tasks.lock().await;
if let Some(task) = tasks.get(&key) {
if task.watch.is_closed() {
task.task.abort();
}
}
});
}
fn ensure_task_gone(self, key: OciPackerTaskKey) {
tokio::task::spawn(async move {
let mut tasks = self.tasks.lock().await;
if let Some(task) = tasks.remove(&key) {
warn!("packer task {} aborted", key);
task.watch.send_replace(Some(Err(anyhow!("task aborted"))));
}
});
}
}
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
struct OciPackerTaskKey {
digest: String,
format: OciPackedFormat,
}
impl Display for OciPackerTaskKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_fmt(format_args!("{}:{}", self.digest, self.format.extension()))
} }
} }

View File

@ -1,7 +1,11 @@
use std::sync::Arc;
use indexmap::IndexMap; use indexmap::IndexMap;
use tokio::sync::{mpsc::Sender, Mutex}; use std::sync::Arc;
use tokio::{
sync::{broadcast, Mutex},
task::JoinHandle,
};
const OCI_PROGRESS_QUEUE_LEN: usize = 100;
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct OciProgress { pub struct OciProgress {
@ -99,16 +103,25 @@ pub enum OciProgressLayerPhase {
#[derive(Clone)] #[derive(Clone)]
pub struct OciProgressContext { pub struct OciProgressContext {
sender: Sender<OciProgress>, sender: broadcast::Sender<OciProgress>,
} }
impl OciProgressContext { impl OciProgressContext {
pub fn new(sender: Sender<OciProgress>) -> OciProgressContext { pub fn create() -> (OciProgressContext, broadcast::Receiver<OciProgress>) {
let (sender, receiver) = broadcast::channel(OCI_PROGRESS_QUEUE_LEN);
(OciProgressContext::new(sender), receiver)
}
pub fn new(sender: broadcast::Sender<OciProgress>) -> OciProgressContext {
OciProgressContext { sender } OciProgressContext { sender }
} }
pub fn update(&self, progress: &OciProgress) { pub fn update(&self, progress: &OciProgress) {
let _ = self.sender.try_send(progress.clone()); let _ = self.sender.send(progress.clone());
}
pub fn subscribe(&self) -> broadcast::Receiver<OciProgress> {
self.sender.subscribe()
} }
} }
@ -137,4 +150,20 @@ impl OciBoundProgress {
function(&mut progress); function(&mut progress);
self.context.update(&progress); self.context.update(&progress);
} }
pub async fn also_update(&self, context: OciProgressContext) -> JoinHandle<()> {
let progress = self.instance.lock().await.clone();
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;
}
}
}
})
}
} }

View File

@ -1,7 +1,7 @@
use anyhow::Result; use anyhow::Result;
use backhand::{FilesystemWriter, NodeHeader}; use backhand::{FilesystemWriter, NodeHeader};
use krata::launchcfg::LaunchInfo; use krata::launchcfg::LaunchInfo;
use krataoci::packer::OciImagePacked; use krataoci::packer::OciPackedImage;
use log::trace; use log::trace;
use std::fs; use std::fs;
use std::fs::File; use std::fs::File;
@ -9,13 +9,13 @@ use std::path::PathBuf;
use uuid::Uuid; use uuid::Uuid;
pub struct ConfigBlock<'a> { pub struct ConfigBlock<'a> {
pub image: &'a OciImagePacked, pub image: &'a OciPackedImage,
pub file: PathBuf, pub file: PathBuf,
pub dir: PathBuf, pub dir: PathBuf,
} }
impl ConfigBlock<'_> { impl ConfigBlock<'_> {
pub fn new<'a>(uuid: &Uuid, image: &'a OciImagePacked) -> Result<ConfigBlock<'a>> { pub fn new<'a>(uuid: &Uuid, image: &'a OciPackedImage) -> Result<ConfigBlock<'a>> {
let mut dir = std::env::temp_dir().clone(); let mut dir = std::env::temp_dir().clone();
dir.push(format!("krata-cfg-{}", uuid)); dir.push(format!("krata-cfg-{}", uuid));
fs::create_dir_all(&dir)?; fs::create_dir_all(&dir)?;

View File

@ -10,7 +10,7 @@ use krata::launchcfg::{
LaunchInfo, LaunchNetwork, LaunchNetworkIpv4, LaunchNetworkIpv6, LaunchNetworkResolver, LaunchInfo, LaunchNetwork, LaunchNetworkIpv4, LaunchNetworkIpv6, LaunchNetworkResolver,
LaunchPackedFormat, LaunchRoot, LaunchPackedFormat, LaunchRoot,
}; };
use krataoci::packer::OciImagePacked; use krataoci::packer::OciPackedImage;
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
use uuid::Uuid; use uuid::Uuid;
use xenclient::{DomainChannel, DomainConfig, DomainDisk, DomainNetworkInterface}; use xenclient::{DomainChannel, DomainConfig, DomainDisk, DomainNetworkInterface};
@ -30,7 +30,7 @@ pub struct GuestLaunchRequest {
pub env: HashMap<String, String>, pub env: HashMap<String, String>,
pub run: Option<Vec<String>>, pub run: Option<Vec<String>>,
pub debug: bool, pub debug: bool,
pub image: OciImagePacked, pub image: OciPackedImage,
} }
pub struct GuestLauncher { pub struct GuestLauncher {