Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Client snapshot download progress with json output #1172

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion mithril-client/Cargo.toml
Original file line number Diff line number Diff line change
@@ -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 }
Expand Down
12 changes: 6 additions & 6 deletions mithril-client/src/aggregator_client/http_client.rs
Original file line number Diff line number Diff line change
@@ -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};

Expand All @@ -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 {
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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| {
Expand Down Expand Up @@ -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(())
Expand Down
12 changes: 6 additions & 6 deletions mithril-client/src/aggregator_client/snapshot_client.rs
Original file line number Diff line number Diff line change
@@ -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)]
Expand Down Expand Up @@ -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)
Expand All @@ -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),
Expand Down
17 changes: 9 additions & 8 deletions mithril-client/src/commands/snapshot/download.rs
Original file line number Diff line number Diff line change
@@ -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)]
Expand Down Expand Up @@ -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?;

Expand Down
41 changes: 24 additions & 17 deletions mithril-client/src/services/snapshot.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -72,7 +75,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>;
}

Expand Down Expand Up @@ -210,20 +213,20 @@ 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);
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)
Expand All @@ -234,23 +237,27 @@ 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()
.with_key("eta", |state: &ProgressState, w: &mut dyn Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap())
.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?;

Expand All @@ -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(
Expand Down Expand Up @@ -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.");
Expand Down Expand Up @@ -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.");
Expand Down Expand Up @@ -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.");
Expand Down Expand Up @@ -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.");
Expand Down
2 changes: 2 additions & 0 deletions mithril-client/src/utils/mod.rs
Original file line number Diff line number Diff line change
@@ -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::*;
Loading
Loading