Skip to content

Commit

Permalink
cli: Centralize progress parsing a bit
Browse files Browse the repository at this point in the history
Add a shared structure (so we only document the option once)
and add a `TryFrom` chain which cleans up the conversion from
the options into a `ProgressWriter`.

Signed-off-by: Colin Walters <[email protected]>
  • Loading branch information
cgwalters committed Dec 18, 2024
1 parent f443cd4 commit 1d5d4dd
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 25 deletions.
50 changes: 31 additions & 19 deletions lib/src/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
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;

Expand All @@ -26,12 +25,35 @@ use serde::{Deserialize, Serialize};

use crate::deploy::RequiredHostSpec;
use crate::lints;
use crate::progress_jsonl;
use crate::progress_jsonl::ProgressWriter;
use crate::progress_jsonl::{ProgressWriter, RawProgressFd};
use crate::spec::Host;
use crate::spec::ImageReference;
use crate::utils::sigpolicy_from_opts;

/// Shared progress options
#[derive(Debug, Parser, PartialEq, Eq)]
pub(crate) struct ProgressOptions {
/// File descriptor number which must refer to an open pipe (anonymous or named).
///
/// Interactive progress will be written to this file descriptor as "JSON lines"
/// format, where each value is separated by a newline.
#[clap(long)]
pub(crate) json_fd: Option<RawProgressFd>,
}

impl TryFrom<ProgressOptions> for ProgressWriter {
type Error = anyhow::Error;

fn try_from(value: ProgressOptions) -> Result<Self> {
let r = value
.json_fd
.map(TryInto::try_into)
.transpose()?
.unwrap_or_default();
Ok(r)
}
}

/// Perform an upgrade operation
#[derive(Debug, Parser, PartialEq, Eq)]
pub(crate) struct UpgradeOpts {
Expand All @@ -54,9 +76,8 @@ pub(crate) struct UpgradeOpts {
#[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>,
#[clap(flatten)]
pub(crate) progress: ProgressOptions,
}

/// Perform an switch operation
Expand Down Expand Up @@ -107,9 +128,8 @@ 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>,
#[clap(flatten)]
pub(crate) progress: ProgressOptions,
}

/// Options controlling rollback
Expand Down Expand Up @@ -653,11 +673,7 @@ 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();
let prog: ProgressWriter = opts.progress.try_into()?;

// If there's no specified image, let's be nice and check if the booted system is using rpm-ostree
if imgref.is_none() {
Expand Down Expand Up @@ -774,11 +790,7 @@ 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();
let prog: ProgressWriter = opts.progress.try_into()?;

// If we're doing an in-place mutation, we shortcut most of the rest of the work here
if opts.mutate_in_place {
Expand Down
30 changes: 24 additions & 6 deletions lib/src/progress_jsonl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,10 @@
//! see <https://jsonlines.org/>.
use anyhow::Result;
use fn_error_context::context;
use serde::Serialize;
use std::borrow::Cow;
use std::os::fd::{FromRawFd, OwnedFd, RawFd};
use std::str::FromStr;
use std::sync::Arc;
use std::time::Instant;
use tokio::io::{AsyncWriteExt, BufWriter};
Expand Down Expand Up @@ -131,6 +131,22 @@ pub enum Event<'t> {
},
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct RawProgressFd(RawFd);

impl FromStr for RawProgressFd {
type Err = anyhow::Error;

fn from_str(s: &str) -> Result<Self> {
let fd = s.parse::<u32>()?;
// Sanity check
if matches!(fd, 0..=2) {
anyhow::bail!("Cannot use fd {fd} for progress JSON")
}
Ok(Self(fd.try_into()?))
}
}

#[derive(Debug)]
struct ProgressWriterInner {
last_write: Option<std::time::Instant>,
Expand Down Expand Up @@ -163,14 +179,16 @@ impl From<Sender> for ProgressWriter {
}
}

impl ProgressWriter {
/// Given a raw file descriptor, create an instance of a json-lines writer.
impl TryFrom<RawProgressFd> for ProgressWriter {
type Error = anyhow::Error;

#[allow(unsafe_code)]
#[context("Creating progress writer")]
pub(crate) fn from_raw_fd(fd: RawFd) -> Result<Self> {
unsafe { OwnedFd::from_raw_fd(fd) }.try_into()
fn try_from(fd: RawProgressFd) -> Result<Self> {
unsafe { OwnedFd::from_raw_fd(fd.0) }.try_into()
}
}

impl ProgressWriter {
/// Serialize the target object to JSON as a single line
pub(crate) async fn send_impl<T: Serialize>(&self, v: T, required: bool) -> Result<()> {
let mut guard = self.inner.lock().await;
Expand Down

0 comments on commit 1d5d4dd

Please sign in to comment.