ms2k/lazaret
SHA256
Watch
0
Code Releases 3 Activity

Show progress while downloading sandbox images

Browse source
This commit is contained in:
ms2k 2026-07-29 00:48:24 +09:00
commit aaee63a6e3
2 changed files with 197 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

4
crates/lazaret-prepared-lower/src/lib.rs
View file

@ -13,4 +13,6 @@ pub use store::{
PrepareRecovery, PrepareRequest, PreparedArtifactGcResult, PreparedArtifactSelection,
PreparedLowerStore, ReclaimOutcome,
};
pub use transport::{OciRegistryTransport, Platform, RegistryCredentials, TransportedImage};
pub use transport::{
OciDownloadProgress, OciRegistryTransport, Platform, RegistryCredentials, TransportedImage,
};

83
crates/lazaret-prepared-lower/src/transport.rs
View file

@ -7,6 +7,7 @@ use std::str::FromStr;
use std::sync::Arc;
use anyhow::{Context, Result, bail};
use futures_util::TryStreamExt;
use lazaret_core::configuration::ConfinedPathPlan;
use oci_client::Reference;
use oci_client::client::{Client, ClientConfig};
@ -19,6 +20,7 @@ use rustix::fs::{
};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tokio::io::AsyncWriteExt;
use crate::{PreparedLowerError, PreparedLowerResult};
@ -138,6 +140,13 @@ impl Drop for RetainedLayout {
#[derive(Clone, Debug, Default)]
pub struct OciRegistryTransport;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OciDownloadProgress {
Started { total_bytes: u64 },
Advanced { downloaded_bytes: u64 },
Finished,
}
impl OciRegistryTransport {
/// Pulls one explicitly selected platform into a digest-verified private OCI layout.
///
@ -174,6 +183,30 @@ impl OciRegistryTransport {
platform: &Platform,
storage_root: &ConfinedPathPlan,
staging_id: [u8; 32],
) -> PreparedLowerResult<TransportedImage> {
self.pull_confined_with_progress(
reference,
credentials,
platform,
storage_root,
staging_id,
|_| {},
)
.await
}
/// Pulls an image into Storage Root and reports exact missing-layer download bytes.
///
/// # Errors
/// Refuses the same conditions as [`Self::pull_confined`].
pub async fn pull_confined_with_progress(
&self,
reference: &str,
credentials: &RegistryCredentials,
platform: &Platform,
storage_root: &ConfinedPathPlan,
staging_id: [u8; 32],
mut progress: impl FnMut(OciDownloadProgress),
) -> PreparedLowerResult<TransportedImage> {
self.pull_to(
reference,
@ -184,6 +217,7 @@ impl OciRegistryTransport {
staging_id,
},
ClientConfig::default(),
&mut progress,
)
.await
}
@ -196,12 +230,14 @@ impl OciRegistryTransport {
layout_path: &Path,
config: ClientConfig,
) -> PreparedLowerResult<TransportedImage> {
let mut progress = |_| {};
self.pull_to(
reference,
credentials,
platform,
LayoutDestination::Ambient(layout_path),
config,
&mut progress,
)
.await
}
@ -213,6 +249,7 @@ impl OciRegistryTransport {
platform: &Platform,
destination: LayoutDestination<'_>,
mut config: ClientConfig,
progress: &mut dyn FnMut(OciDownloadProgress),
) -> PreparedLowerResult<TransportedImage> {
validate_credentials(credentials).map_err(PreparedLowerError::authentication)?;
let reference = Reference::from_str(reference)
@ -260,7 +297,7 @@ impl OciRegistryTransport {
.write_verified_blob(&manifest.config.digest, config_json.as_bytes())
.context("writing verified OCI config")
.map_err(PreparedLowerError::publication)?;
pull_layers(&client, &reference, &manifest, &layout).await?;
pull_layers(&client, &reference, &manifest, &layout, progress).await?;
let local_manifest =
serde_json::to_vec(&manifest).map_err(PreparedLowerError::manifest_selection)?;
let local_digest = format!("sha256:{}", hex(&Sha256::digest(&local_manifest)));
@ -294,7 +331,10 @@ async fn pull_layers(
reference: &Reference,
manifest: &OciImageManifest,
layout: &LayoutWriter,
progress: &mut dyn FnMut(OciDownloadProgress),
) -> PreparedLowerResult<()> {
let mut missing = Vec::new();
let mut total_bytes = 0_u64;
for descriptor in &manifest.layers {
let size = u64::try_from(descriptor.size)
.context("negative OCI layer size")
@ -315,18 +355,48 @@ async fn pull_layers(
.map_err(PreparedLowerError::layer_validation)?;
continue;
}
total_bytes = total_bytes.checked_add(size).ok_or_else(|| {
PreparedLowerError::manifest_selection(anyhow::anyhow!(
"total OCI layer size exceeds u64"
))
})?;
missing.push((descriptor, size, digest_component));
}
progress(OciDownloadProgress::Started { total_bytes });
let mut downloaded_bytes = 0_u64;
for (descriptor, size, digest_component) in missing {
let (file, temporary) = layout
.create_blob(digest_component)
.with_context(|| format!("creating OCI layer {}", descriptor.digest))
.map_err(PreparedLowerError::materialization)?;
let file = tokio::fs::File::from_std(file);
if let Err(error) = client.pull_blob(reference, descriptor, file).await {
let mut file = tokio::fs::File::from_std(file);
let mut stream = match client.pull_blob_stream(reference, descriptor).await {
Ok(stream) => stream,
Err(error) => {
layout.remove_temporary_blob(&temporary);
return Err(distribution_error(
error,
format!("pulling OCI layer {}", descriptor.digest),
));
}
};
let download = async {
while let Some(bytes) = stream.try_next().await? {
file.write_all(&bytes).await?;
downloaded_bytes = downloaded_bytes.saturating_add(bytes.len() as u64);
progress(OciDownloadProgress::Advanced { downloaded_bytes });
}
file.flush().await
}
.await;
if let Err(error) = download {
layout.remove_temporary_blob(&temporary);
return Err(distribution_error(
error,
format!("pulling OCI layer {}", descriptor.digest),
return Err(PreparedLowerError::transport(
anyhow::Error::new(error)
.context(format!("streaming OCI layer {}", descriptor.digest)),
));
}
drop(file);
let observed = layout
.temporary_blob_size(&temporary)
.with_context(|| format!("measuring OCI layer {}", descriptor.digest))
@ -342,6 +412,7 @@ async fn pull_layers(
.with_context(|| format!("publishing OCI layer {}", descriptor.digest))
.map_err(PreparedLowerError::publication)?;
}
progress(OciDownloadProgress::Finished);
Ok(())
}