Show progress while downloading sandbox images

This commit is contained in:
ms2k 2026-07-29 00:48:24 +09:00
commit aaee63a6e3
6 changed files with 197 additions and 12 deletions

View file

@ -22,6 +22,7 @@ harness = false
[dependencies]
ed25519-dalek.workspace = true
getrandom.workspace = true
indicatif.workspace = true
lazaret-control-client = { version = "=0.0.0", path = "../lazaret-control-client" }
lazaret-core = { version = "=0.0.0", path = "../lazaret-core" }
lazaret-domain-state = { version = "=0.0.0", path = "../lazaret-domain-state" }

View file

@ -9,6 +9,7 @@ use std::os::unix::fs::MetadataExt;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use indicatif::{ProgressBar, ProgressStyle};
use lazaret_core::configuration::{
ConfigurationLoadRequest, ResolvedDomainConfiguration, load_trusted_configuration,
};
@ -17,6 +18,7 @@ use lazaret_core::identity::{
};
use lazaret_core::records::{Domain, Session, SessionPhase};
use lazaret_linux::filesystem::try_exclusive_flock;
use lazaret_prepared_lower::OciDownloadProgress;
use lazaret_runtime_supervisor::sandbox::{
DOMAIN_OPEN_FILE_LIMIT, RUNSC_EXECUTABLE_NAME, default_sandbox_layout,
};
@ -34,6 +36,45 @@ use crate::waiter::BootstrapWaiterFace;
pub use lazaret_core::protocol::PRIVATE_CLIENT_BOOTSTRAP_ROLE;
const MAX_CONCURRENT_CONNECTIONS: usize = 32;
#[derive(Default)]
struct ImageDownloadProgress {
bar: Option<ProgressBar>,
}
impl ImageDownloadProgress {
fn update(&mut self, progress: OciDownloadProgress) {
match progress {
OciDownloadProgress::Started { total_bytes: 0 } => {}
OciDownloadProgress::Started { total_bytes } => {
let style = ProgressStyle::with_template(
"Lazaret is downloading sandbox image [{wide_bar}] {bytes}/{total_bytes} ({percent}%)",
)
.expect("static progress template must be valid")
.progress_chars("=>-");
self.bar = Some(ProgressBar::new(total_bytes).with_style(style));
}
OciDownloadProgress::Advanced { downloaded_bytes } => {
if let Some(bar) = &self.bar {
bar.set_position(downloaded_bytes);
}
}
OciDownloadProgress::Finished => self.clear(),
}
}
fn clear(&mut self) {
if let Some(bar) = self.bar.take() {
bar.finish_and_clear();
}
}
}
impl Drop for ImageDownloadProgress {
fn drop(&mut self) {
self.clear();
}
}
#[derive(Debug)]
pub struct PrivateClientBootstrapError(String);
@ -140,6 +181,7 @@ fn run_private_client_bootstrap(arguments: &[OsString]) -> Result<(), PrivateCli
.enable_all()
.build()
.map_err(failure)?;
let mut image_progress = ImageDownloadProgress::default();
let mut selection = runtime
.block_on(select_runtime(
&configuration,
@ -149,6 +191,7 @@ fn run_private_client_bootstrap(arguments: &[OsString]) -> Result<(), PrivateCli
sandbox_init: &installed.join("lazaret-sandbox-init"),
runsc: &installed.join(RUNSC_EXECUTABLE_NAME),
},
|progress| image_progress.update(progress),
))
.map_err(failure)
.inspect_err(|_| signal(b'E'))?;
@ -552,6 +595,18 @@ mod tests {
use super::*;
#[test]
fn image_download_progress_accepts_the_exact_event_lifecycle() {
let mut progress = ImageDownloadProgress::default();
progress.update(OciDownloadProgress::Started { total_bytes: 1_024 });
assert!(progress.bar.is_some());
progress.update(OciDownloadProgress::Advanced {
downloaded_bytes: 1_024,
});
progress.update(OciDownloadProgress::Finished);
assert!(progress.bar.is_none());
}
#[test]
fn host_coordinator_plans_a_bounded_raise_from_a_low_soft_nofile_limit() {
assert_eq!(

View file

@ -16,8 +16,8 @@ use lazaret_linux::filesystem::SharedFileLock;
use lazaret_linux::platform::{PlatformObservation, check_platform};
use lazaret_linux::user_namespace::UserNamespaceMap;
use lazaret_prepared_lower::{
OciRegistryTransport, Platform, PreparedDestination, PreparedDestinationKind,
PreparedLowerError, PreparedLowerStore, RegistryCredentials,
OciDownloadProgress, OciRegistryTransport, Platform, PreparedDestination,
PreparedDestinationKind, PreparedLowerError, PreparedLowerStore, RegistryCredentials,
};
use lazaret_runtime_supervisor::sandbox::{
CREDENTIALS_DESTINATION, DEV_DESTINATION, ESCALATION_CLIENT_DESTINATION,
@ -150,6 +150,7 @@ pub async fn select_runtime(
configuration: &ResolvedDomainConfiguration,
store: Arc<HostCoordinatorDomainStore>,
product_paths: InstalledProductPaths<'_>,
progress: impl FnMut(OciDownloadProgress),
) -> Result<RuntimeSelection, RuntimeSelectionError> {
let kernel_release = std::fs::read_to_string("/proc/sys/kernel/osrelease").map_err(failure)?;
select_runtime_for_platform(
@ -161,6 +162,7 @@ pub async fn select_runtime(
architecture: std::env::consts::ARCH,
kernel_release: kernel_release.trim(),
},
progress,
)
.await
}
@ -170,6 +172,7 @@ async fn select_runtime_for_platform(
store: Arc<HostCoordinatorDomainStore>,
product_paths: InstalledProductPaths<'_>,
observation: &PlatformObservation<'_>,
progress: impl FnMut(OciDownloadProgress),
) -> Result<RuntimeSelection, RuntimeSelectionError> {
check_platform(observation).map_err(failure)?;
let user_namespace_map = UserNamespaceMap::resolve_current().map_err(failure)?;
@ -186,12 +189,13 @@ async fn select_runtime_for_platform(
PreparedLowerStore::under_storage_root(configuration.storage_root(), catalog.clone())?;
recover_before_cold_start(configuration, &store, &catalog, &prepared).map_err(failure)?;
let image = OciRegistryTransport
.pull_confined(
.pull_confined_with_progress(
&configuration.fingerprint_record().image,
&RegistryCredentials::Anonymous,
&Platform::linux("amd64"),
configuration.storage_root(),
random_id()?,
progress,
)
.await
.map_err(RuntimeSelectionError::PreparedLower)?;
@ -493,6 +497,7 @@ mod tests {
runsc: Path::new("/missing-runsc"),
},
&observation,
|_| {},
));
let mut context = Context::from_waker(Waker::noop());
let Poll::Ready(result) = future.as_mut().poll(&mut context) else {

View file

@ -8,6 +8,7 @@ license.workspace = true
[dependencies]
anyhow.workspace = true
flate2.workspace = true
futures-util.workspace = true
lazaret-core = { version = "=0.0.0", path = "../lazaret-core" }
lazaret-linux = { version = "=0.0.0", path = "../lazaret-linux" }
lazaret-storage-catalog = { version = "=0.0.0", path = "../lazaret-storage-catalog" }

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,
};

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(())
}