Add progress report to client snapshot download in json output

This is done using a new type, `DownloadProgressReporter`, that wrap a
indicatif progress bar that give to use all the data that we need to
print, thank you indicatif !
Those print are limited to one every 333ms.

Author: Alenar

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

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<PathBuf> {
         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),

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?;
 

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,9 @@ use mithril_common::{
 
 use crate::{
     aggregator_client::{AggregatorHTTPClientError, CertificateClient, SnapshotClient},
-    utils::{SnapshotUnpacker, SnapshotUnpackerError},
+    utils::{
+        DownloadProgressReporter, ProgressOutputType, SnapshotUnpacker, SnapshotUnpackerError,
+    },
 };
 
 /// [SnapshotService] related errors.
@@ -72,7 +74,7 @@ pub trait SnapshotService: Sync + Send {
         snapshot_entity: &SignedEntity<Snapshot>,
         pathdir: &Path,
         genesis_verification_key: &str,
-        progress_target: ProgressDrawTarget,
+        progress_output_type: ProgressOutputType,
     ) -> StdResult<PathBuf>;
 }
 
@@ -210,12 +212,12 @@ impl SnapshotService for MithrilClientSnapshotService {
         snapshot_entity: &SignedEntity<Snapshot>,
         download_dir: &Path,
         genesis_verification_key: &str,
-        progress_target: ProgressDrawTarget,
+        progress_output_type: ProgressOutputType,
     ) -> StdResult<PathBuf> {
         debug!("Snapshot service: download.");
 
         let db_dir = download_dir.join("db");
-        let progress_bar = MultiProgress::with_draw_target(progress_target);
+        let progress_bar = MultiProgress::with_draw_target(progress_output_type.into());
         progress_bar.println("1/7 - Checking local disk info…")?;
         let unpacker = SnapshotUnpacker;
 
@@ -246,7 +248,11 @@ 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()))?;
 
@@ -550,7 +556,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 +596,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 +642,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 +690,7 @@ mod tests {
                 &snapshot,
                 &test_path,
                 &genesis_verification_key.to_json_hex().unwrap(),
-                ProgressDrawTarget::hidden(),
+                ProgressOutputType::Hidden,
             )
             .await
             .expect_err("Snapshot download should fail.");

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::*;

mithril-client/src/utils/progress_reporter.rs

@@ -0,0 +1,83 @@
+use indicatif::{ProgressBar, ProgressDrawTarget};
+use slog_scope::warn;
+use std::{
+    sync::RwLock,
+    time::{Duration, Instant},
+};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+/// Output type of a [ProgressReporter]
+pub enum ProgressOutputType {
+    /// Output to json
+    JsonReporter,
+    /// Output to tty
+    TTY,
+    /// No output
+    Hidden,
+}
+
+impl From<ProgressOutputType> 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 [ProgressBar] to allow reporting to json.
+pub struct DownloadProgressReporter {
+    progress_bar: ProgressBar,
+    output_type: ProgressOutputType,
+    last_json_report_instant: RwLock<Option<Instant>>,
+}
+
+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<Duration> {
+        match self.last_json_report_instant.read() {
+            Ok(instant) => (*instant).map(|instant| instant.elapsed()),
+            Err(_) => None,
+        }
+    }
+}