Merge pull request #921 from antheas/main

feat(upgrade): Add total progress bar and JSON output

Author: cgwalters

lib/src/cli.rs

@@ -4,6 +4,7 @@
 
 use std::ffi::{CString, OsStr, OsString};
 use std::io::Seek;
+use std::os::fd::RawFd;
 use std::os::unix::process::CommandExt;
 use std::process::Command;
 
@@ -25,6 +26,8 @@ use serde::{Deserialize, Serialize};
 
 use crate::deploy::RequiredHostSpec;
 use crate::lints;
+use crate::progress_jsonl;
+use crate::progress_jsonl::ProgressWriter;
 use crate::spec::Host;
 use crate::spec::ImageReference;
 use crate::utils::sigpolicy_from_opts;
@@ -52,6 +55,10 @@ pub(crate) struct UpgradeOpts {
     /// a userspace-only restart.
     #[clap(long, conflicts_with = "check")]
     pub(crate) apply: bool,
+
+    /// Pipe download progress to this fd in a jsonl format.
+    #[clap(long)]
+    pub(crate) json_fd: Option<RawFd>,
 }
 
 /// Perform an switch operation
@@ -101,6 +108,10 @@ pub(crate) struct SwitchOpts {
 
     /// Target image to use for the next boot.
     pub(crate) target: String,
+
+    /// Pipe download progress to this fd in a jsonl format.
+    #[clap(long)]
+    pub(crate) json_fd: Option<RawFd>,
 }
 
 /// Options controlling rollback
@@ -644,6 +655,12 @@ async fn upgrade(opts: UpgradeOpts) -> Result<()> {
     let (booted_deployment, _deployments, host) =
         crate::status::get_status_require_booted(sysroot)?;
     let imgref = host.spec.image.as_ref();
+    let prog = opts
+        .json_fd
+        .map(progress_jsonl::ProgressWriter::from_raw_fd)
+        .transpose()?
+        .unwrap_or_default();
+
     // If there's no specified image, let's be nice and check if the booted system is using rpm-ostree
     if imgref.is_none() {
         let booted_incompatible = host
@@ -700,7 +717,7 @@ async fn upgrade(opts: UpgradeOpts) -> Result<()> {
             }
         }
     } else {
-        let fetched = crate::deploy::pull(repo, imgref, None, opts.quiet).await?;
+        let fetched = crate::deploy::pull(repo, imgref, None, opts.quiet, prog.clone()).await?;
         let staged_digest = staged_image.map(|s| s.digest().expect("valid digest in status"));
         let fetched_digest = &fetched.manifest_digest;
         tracing::debug!("staged: {staged_digest:?}");
@@ -723,7 +740,7 @@ async fn upgrade(opts: UpgradeOpts) -> Result<()> {
             println!("No update available.")
         } else {
             let osname = booted_deployment.osname();
-            crate::deploy::stage(sysroot, &osname, &fetched, &spec).await?;
+            crate::deploy::stage(sysroot, &osname, &fetched, &spec, prog.clone()).await?;
             changed = true;
             if let Some(prev) = booted_image.as_ref() {
                 if let Some(fetched_manifest) = fetched.get_manifest(repo)? {
@@ -759,6 +776,11 @@ async fn switch(opts: SwitchOpts) -> Result<()> {
     );
     let target = ostree_container::OstreeImageReference { sigverify, imgref };
     let target = ImageReference::from(target);
+    let prog = opts
+        .json_fd
+        .map(progress_jsonl::ProgressWriter::from_raw_fd)
+        .transpose()?
+        .unwrap_or_default();
 
     // If we're doing an in-place mutation, we shortcut most of the rest of the work here
     if opts.mutate_in_place {
@@ -794,7 +816,7 @@ async fn switch(opts: SwitchOpts) -> Result<()> {
     }
     let new_spec = RequiredHostSpec::from_spec(&new_spec)?;
 
-    let fetched = crate::deploy::pull(repo, &target, None, opts.quiet).await?;
+    let fetched = crate::deploy::pull(repo, &target, None, opts.quiet, prog.clone()).await?;
 
     if !opts.retain {
         // By default, we prune the previous ostree ref so it will go away after later upgrades
@@ -808,7 +830,7 @@ async fn switch(opts: SwitchOpts) -> Result<()> {
     }
 
     let stateroot = booted_deployment.osname();
-    crate::deploy::stage(sysroot, &stateroot, &fetched, &new_spec).await?;
+    crate::deploy::stage(sysroot, &stateroot, &fetched, &new_spec, prog.clone()).await?;
 
     if opts.apply {
         crate::reboot::reboot()?;
@@ -850,18 +872,20 @@ async fn edit(opts: EditOpts) -> Result<()> {
     host.spec.verify_transition(&new_host.spec)?;
     let new_spec = RequiredHostSpec::from_spec(&new_host.spec)?;
 
+    let prog = ProgressWriter::default();
+
     // We only support two state transitions right now; switching the image,
     // or flipping the bootloader ordering.
     if host.spec.boot_order != new_host.spec.boot_order {
         return crate::deploy::rollback(sysroot).await;
     }
 
-    let fetched = crate::deploy::pull(repo, new_spec.image, None, opts.quiet).await?;
+    let fetched = crate::deploy::pull(repo, new_spec.image, None, opts.quiet, prog.clone()).await?;
 
     // TODO gc old layers here
 
     let stateroot = booted_deployment.osname();
-    crate::deploy::stage(sysroot, &stateroot, &fetched, &new_spec).await?;
+    crate::deploy::stage(sysroot, &stateroot, &fetched, &new_spec, prog.clone()).await?;
 
     Ok(())
 }