Merge pull request #1183 from input-output-hk/ensemble/1115/download-unpack-snapshots-simultaneously

Download unpack snapshots simultaneously

Author: Alenar

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-client"
-version = "0.3.36"
+version = "0.3.37"
 description = "A Mithril Client"
 authors = { workspace = true }
 edition = { workspace = true }
@@ -19,6 +19,7 @@ cli-table = "0.4.7"
 config = "0.13.3"
 directories = "5.0.1"
 flate2 = "1.0.27"
+flume = "0.11.0"
 fs2 = "0.4.3"
 futures = "0.3.28"
 hex = "0.4.3"

mithril-client/Cargo.toml

@@ -1,3 +1,4 @@
+use anyhow::anyhow;
 use async_recursion::async_recursion;
 use async_trait::async_trait;
 use futures::StreamExt;
@@ -6,14 +7,14 @@ use semver::Version;
 use slog_scope::debug;
 use std::{path::Path, sync::Arc};
 use thiserror::Error;
-use tokio::{fs, io::AsyncWriteExt, sync::RwLock};
+use tokio::sync::RwLock;
 
 #[cfg(test)]
 use mockall::automock;
 
-use mithril_common::{StdError, MITHRIL_API_VERSION_HEADER};
+use mithril_common::{StdError, StdResult, MITHRIL_API_VERSION_HEADER};
 
-use crate::utils::DownloadProgressReporter;
+use crate::utils::{DownloadProgressReporter, SnapshotUnpacker};
 
 /// Error tied with the Aggregator client
 #[derive(Error, Debug)]
@@ -51,11 +52,11 @@ pub trait AggregatorClient: Sync + Send {
     /// Get the content back from the Aggregator, the URL is a relative path for a resource
     async fn get_content(&self, url: &str) -> Result<String, AggregatorHTTPClientError>;
 
-    /// Download large files on the disk
-    async fn download(
+    /// Download and unpack large archives on the disk
+    async fn download_unpack(
         &self,
         url: &str,
-        filepath: &Path,
+        target_dir: &Path,
         progress_reporter: DownloadProgressReporter,
     ) -> Result<(), AggregatorHTTPClientError>;
 
@@ -179,45 +180,65 @@ impl AggregatorClient for AggregatorHTTPClient {
             })
     }
 
-    async fn download(
+    async fn download_unpack(
         &self,
         url: &str,
-        filepath: &Path,
+        target_dir: &Path,
         progress_reporter: DownloadProgressReporter,
     ) -> Result<(), AggregatorHTTPClientError> {
-        let response = self.get(url).await?;
-        let mut local_file = fs::File::create(filepath).await.map_err(|e| {
-            AggregatorHTTPClientError::SubsystemError {
-                message: format!(
-                    "Download: could not create download archive '{}'.",
-                    filepath.display()
+        if !target_dir.is_dir() {
+            Err(AggregatorHTTPClientError::SubsystemError {
+                message: "Download-Unpack: prerequisite error".to_string(),
+                error: anyhow!(
+                    "target path is not a directory or does not exist: `{}`",
+                    target_dir.display()
                 ),
-                error: e.into(),
-            }
-        })?;
+            })?;
+        }
+
         let mut downloaded_bytes: u64 = 0;
-        let mut remote_stream = response.bytes_stream();
+        let mut remote_stream = self.get(url).await?.bytes_stream();
+        let (sender, receiver) = flume::bounded(5);
+
+        let dest_dir = target_dir.to_path_buf();
+        let unpack_thread = tokio::task::spawn_blocking(move || -> StdResult<()> {
+            let unpacker = SnapshotUnpacker;
+            unpacker.unpack_snapshot(receiver, &dest_dir)
+        });
 
         while let Some(item) = remote_stream.next().await {
             let chunk = item.map_err(|e| {
                 AggregatorHTTPClientError::RemoteServerTechnical(format!(
                     "Download: Could not read from byte stream: {e}"
                 ))
             })?;
-            local_file.write_all(&chunk).await.map_err(|e| {
+
+            sender.send_async(chunk.to_vec()).await.map_err(|e| {
                 AggregatorHTTPClientError::SubsystemError {
-                    message: format!(
-                        "Download: could not write {} bytes to file '{}'.",
-                        chunk.len(),
-                        filepath.display()
-                    ),
+                    message: format!("Download: could not write {} bytes to stream.", chunk.len()),
                     error: e.into(),
                 }
             })?;
+
             downloaded_bytes += chunk.len() as u64;
             progress_reporter.report(downloaded_bytes);
         }
 
+        drop(sender); // Signal EOF
+        unpack_thread
+            .await
+            .map_err(|join_error| AggregatorHTTPClientError::SubsystemError {
+                message: format!(
+                    "Unpack: panic while unpacking to dir '{}'",
+                    target_dir.display()
+                ),
+                error: join_error.into(),
+            })?
+            .map_err(|unpack_error| AggregatorHTTPClientError::SubsystemError {
+                message: format!("Unpack: could not unpack to dir '{}'", target_dir.display()),
+                error: unpack_error,
+            })?;
+
         Ok(())
     }
 

mithril-client/src/aggregator_client/http_client.rs

@@ -1,10 +1,7 @@
 //! This module contains a struct to exchange snapshot information with the Aggregator
 
 use slog_scope::warn;
-use std::{
-    path::{Path, PathBuf},
-    sync::Arc,
-};
+use std::{path::Path, sync::Arc};
 use thiserror::Error;
 
 use mithril_common::{
@@ -59,28 +56,24 @@ impl SnapshotClient {
         Ok(message)
     }
 
-    /// Download the snapshot identified by the given snapshot in the given directory
-    pub async fn download(
+    /// Download and unpack the given snapshot to the given directory
+    pub async fn download_unpack(
         &self,
         snapshot: &Snapshot,
-        download_dir: &Path,
+        target_dir: &Path,
         progress_reporter: DownloadProgressReporter,
-    ) -> StdResult<PathBuf> {
-        let filepath = PathBuf::new()
-            .join(download_dir)
-            .join(format!("snapshot-{}.tar.gz", snapshot.digest));
-
+    ) -> StdResult<()> {
         for url in snapshot.locations.as_slice() {
             if self.http_client.probe(url).await.is_ok() {
-                match self
+                return match self
                     .http_client
-                    .download(url, &filepath, progress_reporter)
+                    .download_unpack(url, target_dir, progress_reporter)
                     .await
                 {
-                    Ok(()) => return Ok(filepath),
+                    Ok(()) => Ok(()),
                     Err(e) => {
                         warn!("Failed downloading snapshot from '{url}' Error: {e}.");
-                        return Err(e.into());
+                        Err(e.into())
                     }
                 };
             }

mithril-client/src/aggregator_client/snapshot_client.rs

@@ -39,13 +39,10 @@ pub enum SnapshotServiceError {
 
     /// Error raised when the certificate verification failed for the downloaded
     /// archive.
-    #[error("Certificate verification failed (snapshot digest = '{digest}'). The archive has been downloaded in '{path}'.")]
+    #[error("Certificate verification failed (snapshot digest = '{digest}').")]
     CouldNotVerifySnapshot {
         /// snapshot digest
         digest: String,
-
-        /// The path of the downloaded archive
-        path: PathBuf,
     },
 
     /// The given certificate could not be found, contains the certificate hash
@@ -218,14 +215,21 @@ impl SnapshotService for MithrilClientSnapshotService {
         debug!("Snapshot service: download.");
 
         let db_dir = download_dir.join("db");
-        let progress_bar = ProgressPrinter::new(progress_output_type, 7);
+        let progress_bar = ProgressPrinter::new(progress_output_type, 6);
         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)?;
         }
 
+        std::fs::create_dir_all(&db_dir).with_context(|| {
+            format!(
+                "Download: could not create target directory '{}'.",
+                db_dir.display()
+            )
+        })?;
+
         progress_bar.report_step(2, "Fetching the certificate's information…")?;
         let certificate = self
             .certificate_client
@@ -241,26 +245,21 @@ impl SnapshotService for MithrilClientSnapshotService {
         let verifier = self.verify_certificate_chain(genesis_verification_key, &certificate);
         self.wait_spinner(&progress_bar, verifier).await?;
 
-        progress_bar.report_step(4, "Downloading the snapshot…")?;
+        progress_bar.report_step(4, "Downloading and unpacking 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(
+        self.snapshot_client
+            .download_unpack(
                 &snapshot_entity.artifact,
-                download_dir,
+                &db_dir,
                 DownloadProgressReporter::new(pb, progress_output_type),
             )
             .await
             .with_context(|| format!("Could not download file in '{}'", download_dir.display()))?;
 
-        progress_bar.report_step(5, "Unpacking the snapshot…")?;
-        let unpacker = unpacker.unpack_snapshot(&snapshot_path, &db_dir);
-        self.wait_spinner(&progress_bar, unpacker).await?;
-
         // Append 'clean' file to speedup node bootstrap
         if let Err(error) = File::create(db_dir.join("clean")) {
             warn!(
@@ -269,14 +268,14 @@ impl SnapshotService for MithrilClientSnapshotService {
             );
         };
 
-        progress_bar.report_step(6, "Computing the snapshot digest…")?;
+        progress_bar.report_step(5, "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.report_step(7, "Verifying the snapshot signature…")?;
+        progress_bar.report_step(6, "Verifying the snapshot signature…")?;
         let expected_message = {
             let mut protocol_message = certificate.protocol_message.clone();
             protocol_message.set_message_part(
@@ -295,7 +294,6 @@ impl SnapshotService for MithrilClientSnapshotService {
 
             return Err(SnapshotServiceError::CouldNotVerifySnapshot {
                 digest: snapshot_entity.artifact.digest.clone(),
-                path: snapshot_path.canonicalize().unwrap(),
             }
             .into());
         }
@@ -400,7 +398,7 @@ mod tests {
         let mut http_client = MockAggregatorHTTPClient::new();
         http_client.expect_probe().returning(|_| Ok(()));
         http_client
-            .expect_download()
+            .expect_download_unpack()
             .returning(move |_, _, _| Ok(()))
             .times(1);
         http_client.expect_get_content().returning(|_| {
@@ -650,10 +648,7 @@ mod tests {
 
         if let Some(e) = err.downcast_ref::<SnapshotServiceError>() {
             match e {
-                SnapshotServiceError::CouldNotVerifySnapshot {
-                    digest,
-                    path: _path,
-                } => {
+                SnapshotServiceError::CouldNotVerifySnapshot { digest } => {
                     assert_eq!("digest-10", digest.as_str());
                 }
                 _ => panic!("Wrong error type when snapshot could not be verified."),

mithril-client/src/services/snapshot.rs

@@ -2,7 +2,9 @@
 //! This module contains tools needed mostly in services layers.
 
 mod progress_reporter;
+mod stream_reader;
 mod unpacker;
 
 pub use progress_reporter::*;
+pub use stream_reader::*;
 pub use unpacker::*;