diff --git a/Cargo.lock b/Cargo.lock index 3910bfecb72..b63dcdc6f5a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2153,7 +2153,7 @@ dependencies = [ [[package]] name = "mithril-client" -version = "0.3.34" +version = "0.3.35" dependencies = [ "anyhow", "async-recursion", diff --git a/mithril-client/Cargo.toml b/mithril-client/Cargo.toml index dcfd73fe53c..9e6ebe94ecc 100644 --- a/mithril-client/Cargo.toml +++ b/mithril-client/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "mithril-client" -version = "0.3.34" +version = "0.3.35" description = "A Mithril Client" authors = { workspace = true } edition = { workspace = true } diff --git a/mithril-client/src/aggregator_client/http_client.rs b/mithril-client/src/aggregator_client/http_client.rs index 4abdc9968d0..b79524e9143 100644 --- a/mithril-client/src/aggregator_client/http_client.rs +++ b/mithril-client/src/aggregator_client/http_client.rs @@ -1,12 +1,10 @@ -use std::{path::Path, sync::Arc}; - use async_recursion::async_recursion; use async_trait::async_trait; use futures::StreamExt; -use indicatif::ProgressBar; use reqwest::{Client, Response, StatusCode}; use semver::Version; use slog_scope::debug; +use std::{path::Path, sync::Arc}; use thiserror::Error; use tokio::{fs, io::AsyncWriteExt, sync::RwLock}; @@ -15,6 +13,8 @@ use mockall::automock; use mithril_common::{StdError, MITHRIL_API_VERSION_HEADER}; +use crate::utils::DownloadProgressReporter; + /// Error tied with the Aggregator client #[derive(Error, Debug)] pub enum AggregatorHTTPClientError { @@ -56,7 +56,7 @@ pub trait AggregatorClient: Sync + Send { &self, url: &str, filepath: &Path, - progress_bar: ProgressBar, + progress_reporter: DownloadProgressReporter, ) -> Result<(), AggregatorHTTPClientError>; /// Test if the given URL points to a valid location & existing content. @@ -183,7 +183,7 @@ impl AggregatorClient for AggregatorHTTPClient { &self, url: &str, filepath: &Path, - progress_bar: ProgressBar, + progress_reporter: DownloadProgressReporter, ) -> Result<(), AggregatorHTTPClientError> { let response = self.get(url).await?; let mut local_file = fs::File::create(filepath).await.map_err(|e| { @@ -215,7 +215,7 @@ impl AggregatorClient for AggregatorHTTPClient { } })?; downloaded_bytes += chunk.len() as u64; - progress_bar.set_position(downloaded_bytes); + progress_reporter.report(downloaded_bytes); } Ok(()) diff --git a/mithril-client/src/aggregator_client/snapshot_client.rs b/mithril-client/src/aggregator_client/snapshot_client.rs index a3e01913350..6d77b057e6e 100644 --- a/mithril-client/src/aggregator_client/snapshot_client.rs +++ b/mithril-client/src/aggregator_client/snapshot_client.rs @@ -1,20 +1,20 @@ //! This module contains a struct to exchange snapshot information with the Aggregator +use slog_scope::warn; use std::{ path::{Path, PathBuf}, sync::Arc, }; +use thiserror::Error; -use indicatif::ProgressBar; use mithril_common::{ entities::Snapshot, messages::{SnapshotListItemMessage, SnapshotListMessage, SnapshotMessage}, StdResult, }; -use slog_scope::warn; -use thiserror::Error; -use super::AggregatorClient; +use crate::aggregator_client::AggregatorClient; +use crate::utils::DownloadProgressReporter; /// Error for the Snapshot client #[derive(Error, Debug)] @@ -64,7 +64,7 @@ impl SnapshotClient { &self, snapshot: &Snapshot, download_dir: &Path, - progress_bar: ProgressBar, + progress_reporter: DownloadProgressReporter, ) -> StdResult { let filepath = PathBuf::new() .join(download_dir) @@ -74,7 +74,7 @@ impl SnapshotClient { if self.http_client.probe(url).await.is_ok() { match self .http_client - .download(url, &filepath, progress_bar) + .download(url, &filepath, progress_reporter) .await { Ok(()) => return Ok(filepath), diff --git a/mithril-client/src/commands/snapshot/download.rs b/mithril-client/src/commands/snapshot/download.rs index 7cb88f4df80..83195e57c68 100644 --- a/mithril-client/src/commands/snapshot/download.rs +++ b/mithril-client/src/commands/snapshot/download.rs @@ -1,11 +1,12 @@ -use std::{path::PathBuf, sync::Arc}; - use clap::Parser; use config::{builder::DefaultState, Config, ConfigBuilder}; -use indicatif::ProgressDrawTarget; +use std::{path::PathBuf, sync::Arc}; + use mithril_common::{messages::FromMessageAdapter, StdResult}; -use crate::{dependencies::DependenciesBuilder, FromSnapshotMessageAdapter}; +use crate::{ + dependencies::DependenciesBuilder, utils::ProgressOutputType, FromSnapshotMessageAdapter, +}; /// Clap command to download the snapshot and verify the certificate. #[derive(Parser, Debug, Clone)] @@ -36,17 +37,17 @@ impl SnapshotDownloadCommand { let snapshot_service = dependencies_builder.get_snapshot_service().await?; let snapshot_entity = FromSnapshotMessageAdapter::adapt(snapshot_service.show(&self.digest).await?); - let progress_target = if self.json { - ProgressDrawTarget::hidden() + let progress_output_type = if self.json { + ProgressOutputType::JsonReporter } else { - ProgressDrawTarget::stdout() + ProgressOutputType::TTY }; let filepath = snapshot_service .download( &snapshot_entity, &self.download_dir, &config.get_string("genesis_verification_key")?, - progress_target, + progress_output_type, ) .await?; diff --git a/mithril-client/src/services/snapshot.rs b/mithril-client/src/services/snapshot.rs index 00f65585c6c..afb17166bd9 100644 --- a/mithril-client/src/services/snapshot.rs +++ b/mithril-client/src/services/snapshot.rs @@ -1,7 +1,7 @@ use anyhow::Context; use async_trait::async_trait; use futures::Future; -use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressState, ProgressStyle}; +use indicatif::{MultiProgress, ProgressBar, ProgressState, ProgressStyle}; use slog_scope::{debug, warn}; use std::{ fmt::Write, @@ -24,7 +24,10 @@ use mithril_common::{ use crate::{ aggregator_client::{AggregatorHTTPClientError, CertificateClient, SnapshotClient}, - utils::{SnapshotUnpacker, SnapshotUnpackerError}, + utils::{ + DownloadProgressReporter, ProgressOutputType, ProgressPrinter, SnapshotUnpacker, + SnapshotUnpackerError, + }, }; /// [SnapshotService] related errors. @@ -72,7 +75,7 @@ pub trait SnapshotService: Sync + Send { snapshot_entity: &SignedEntity, pathdir: &Path, genesis_verification_key: &str, - progress_target: ProgressDrawTarget, + progress_output_type: ProgressOutputType, ) -> StdResult; } @@ -210,20 +213,20 @@ impl SnapshotService for MithrilClientSnapshotService { snapshot_entity: &SignedEntity, download_dir: &Path, genesis_verification_key: &str, - progress_target: ProgressDrawTarget, + progress_output_type: ProgressOutputType, ) -> StdResult { debug!("Snapshot service: download."); let db_dir = download_dir.join("db"); - let progress_bar = MultiProgress::with_draw_target(progress_target); - progress_bar.println("1/7 - Checking local disk info…")?; + let progress_bar = ProgressPrinter::new(progress_output_type, 7); + progress_bar.report_step(1, "Checking local disk info…")?; let unpacker = SnapshotUnpacker; if let Err(e) = unpacker.check_prerequisites(&db_dir, snapshot_entity.artifact.size) { self.check_disk_space_error(e)?; } - progress_bar.println("2/7 - Fetching the certificate's information…")?; + progress_bar.report_step(2, "Fetching the certificate's information…")?; let certificate = self .certificate_client .get(&snapshot_entity.certificate_id) @@ -234,11 +237,11 @@ impl SnapshotService for MithrilClientSnapshotService { ) })?; - progress_bar.println("3/7 - Verifying the certificate chain…")?; + progress_bar.report_step(3, "Verifying the certificate chain…")?; let verifier = self.verify_certificate_chain(genesis_verification_key, &certificate); self.wait_spinner(&progress_bar, verifier).await?; - progress_bar.println("4/7 - Downloading the snapshot…")?; + progress_bar.report_step(4, "Downloading the snapshot…")?; let pb = progress_bar.add(ProgressBar::new(snapshot_entity.artifact.size)); pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})") .unwrap() @@ -246,11 +249,15 @@ impl SnapshotService for MithrilClientSnapshotService { .progress_chars("#>-")); let snapshot_path = self .snapshot_client - .download(&snapshot_entity.artifact, download_dir, pb) + .download( + &snapshot_entity.artifact, + download_dir, + DownloadProgressReporter::new(pb, progress_output_type), + ) .await .with_context(|| format!("Could not download file in '{}'", download_dir.display()))?; - progress_bar.println("5/7 - Unpacking the snapshot…")?; + progress_bar.report_step(5, "Unpacking the snapshot…")?; let unpacker = unpacker.unpack_snapshot(&snapshot_path, &db_dir); self.wait_spinner(&progress_bar, unpacker).await?; @@ -262,14 +269,14 @@ impl SnapshotService for MithrilClientSnapshotService { ); }; - progress_bar.println("6/7 - Computing the snapshot digest…")?; + progress_bar.report_step(6, "Computing the snapshot digest…")?; let unpacked_snapshot_digest = self .immutable_digester .compute_digest(&db_dir, &certificate.beacon) .await .with_context(|| format!("Could not compute digest in '{}'", db_dir.display()))?; - progress_bar.println("7/7 - Verifying the snapshot signature…")?; + progress_bar.report_step(7, "Verifying the snapshot signature…")?; let expected_message = { let mut protocol_message = certificate.protocol_message.clone(); protocol_message.set_message_part( @@ -550,7 +557,7 @@ mod tests { &snapshot, &test_path, &genesis_verification_key.to_json_hex().unwrap(), - ProgressDrawTarget::hidden(), + ProgressOutputType::Hidden, ) .await .expect("Snapshot download should succeed."); @@ -590,7 +597,7 @@ mod tests { &snapshot, &test_path, &genesis_verification_key.to_json_hex().unwrap(), - ProgressDrawTarget::hidden(), + ProgressOutputType::Hidden, ) .await .expect("Snapshot download should succeed."); @@ -636,7 +643,7 @@ mod tests { &signed_entity, &test_path, &genesis_verification_key.to_json_hex().unwrap(), - ProgressDrawTarget::hidden(), + ProgressOutputType::Hidden, ) .await .expect_err("Snapshot digest comparison should fail."); @@ -684,7 +691,7 @@ mod tests { &snapshot, &test_path, &genesis_verification_key.to_json_hex().unwrap(), - ProgressDrawTarget::hidden(), + ProgressOutputType::Hidden, ) .await .expect_err("Snapshot download should fail."); diff --git a/mithril-client/src/utils/mod.rs b/mithril-client/src/utils/mod.rs index c4521fbd4f1..ff94c448724 100644 --- a/mithril-client/src/utils/mod.rs +++ b/mithril-client/src/utils/mod.rs @@ -1,6 +1,8 @@ //! Utilities module //! This module contains tools needed mostly in services layers. +mod progress_reporter; mod unpacker; +pub use progress_reporter::*; pub use unpacker::*; diff --git a/mithril-client/src/utils/progress_reporter.rs b/mithril-client/src/utils/progress_reporter.rs new file mode 100644 index 00000000000..2b7b7f887d8 --- /dev/null +++ b/mithril-client/src/utils/progress_reporter.rs @@ -0,0 +1,127 @@ +use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget}; +use mithril_common::StdResult; +use slog_scope::warn; +use std::{ + ops::Deref, + sync::RwLock, + time::{Duration, Instant}, +}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// Output type of a [ProgressPrinter] or a [DownloadProgressReporter] +pub enum ProgressOutputType { + /// Output to json + JsonReporter, + /// Output to tty + TTY, + /// No output + Hidden, +} + +impl From for ProgressDrawTarget { + fn from(value: ProgressOutputType) -> Self { + match value { + ProgressOutputType::JsonReporter => ProgressDrawTarget::hidden(), + ProgressOutputType::TTY => ProgressDrawTarget::stdout(), + ProgressOutputType::Hidden => ProgressDrawTarget::hidden(), + } + } +} + +/// Wrapper of a indicatif [MultiProgress] to allow reporting to json. +pub struct ProgressPrinter { + multi_progress: MultiProgress, + output_type: ProgressOutputType, + number_of_steps: u16, +} + +impl ProgressPrinter { + /// Instanciate a new progress printer + pub fn new(output_type: ProgressOutputType, number_of_steps: u16) -> Self { + Self { + multi_progress: MultiProgress::with_draw_target(output_type.into()), + output_type, + number_of_steps, + } + } + + /// Report the current step + pub fn report_step(&self, step_number: u16, text: &str) -> StdResult<()> { + match self.output_type { + ProgressOutputType::JsonReporter => println!( + r#"{{"step_num": {step_number}, "total_steps": {}, "message": "{text}"}}"#, + self.number_of_steps + ), + ProgressOutputType::TTY => self + .multi_progress + .println(format!("{step_number}/{} - {text}", self.number_of_steps))?, + ProgressOutputType::Hidden => (), + }; + + Ok(()) + } +} + +impl Deref for ProgressPrinter { + type Target = MultiProgress; + + fn deref(&self) -> &Self::Target { + &self.multi_progress + } +} + +/// Wrapper of a indicatif [ProgressBar] to allow reporting to json. +pub struct DownloadProgressReporter { + progress_bar: ProgressBar, + output_type: ProgressOutputType, + last_json_report_instant: RwLock>, +} + +impl DownloadProgressReporter { + /// Instanciate a new progress reporter + pub fn new(progress_bar: ProgressBar, output_type: ProgressOutputType) -> Self { + Self { + progress_bar, + output_type, + last_json_report_instant: RwLock::new(None), + } + } + + /// Report the current progress + pub fn report(&self, actual_position: u64) { + self.progress_bar.set_position(actual_position); + + if let ProgressOutputType::JsonReporter = self.output_type { + let should_report = match self.get_remaining_time_since_last_json_report() { + Some(remaining_time) => remaining_time > Duration::from_millis(333), + None => true, + }; + + if should_report { + println!( + r#"{{ "bytesDownloaded": {}, "bytesTotal": {}, "secondsLeft": {}.{}, "secondsElapsed": {}.{} }}"#, + self.progress_bar.position(), + self.progress_bar.length().unwrap_or(0), + self.progress_bar.eta().as_secs(), + self.progress_bar.eta().subsec_millis(), + self.progress_bar.elapsed().as_secs(), + self.progress_bar.elapsed().subsec_millis(), + ); + + match self.last_json_report_instant.write() { + Ok(mut instant) => *instant = Some(Instant::now()), + Err(error) => { + warn!("failed to update last json report instant, error: {error:?}") + } + }; + } + }; + } + + fn get_remaining_time_since_last_json_report(&self) -> Option { + match self.last_json_report_instant.read() { + Ok(instant) => (*instant).map(|instant| instant.elapsed()), + Err(_) => None, + } + } +}