Skip to content

Commit

Permalink
feat: Adapt client downloading when Unknown location
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfauvel
sfauvel committed Feb 20, 2025
1 parent c42e98c commit 9d855a3
Show file tree
Hide file tree
Showing 3 changed files with 178 additions and 36 deletions.
57 changes: 55 additions & 2 deletions mithril-client/src/cardano_database_client/download_unpack.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -359,6 +359,9 @@ impl InternalArtifactDownloader {
let file_size = 0;
let file_downloader = match &location {
ImmutablesLocation::CloudStorage { .. } => self.http_file_downloader.clone(),
ImmutablesLocation::Unknown => {
return Err(anyhow!("Unknown location type to download immutable"));
}
};
let file_downloader_uris =
FileDownloaderUri::expand_immutable_files_location_to_file_downloader_uris(
Expand Down Expand Up @@ -414,8 +417,12 @@ impl InternalArtifactDownloader {
.await;
let file_downloader = match &location {
AncillaryLocation::CloudStorage { .. } => self.http_file_downloader.clone(),
AncillaryLocation::Unknown => {
continue;
}
};
let file_downloader_uri = location.into();
let file_downloader_uri = location.try_into()?;

let downloaded = file_downloader
.download_unpack(
&file_downloader_uri,
Expand Down Expand Up @@ -468,7 +475,7 @@ mod tests {

use crate::cardano_database_client::CardanoDatabaseClientDependencyInjector;
use crate::feedback::StackFeedbackReceiver;
use crate::file_downloader::MockFileDownloaderBuilder;
use crate::file_downloader::{MockFileDownloader, MockFileDownloaderBuilder};
use crate::test_utils;

use super::*;
Expand Down Expand Up @@ -876,6 +883,32 @@ mod tests {
.expect_err("download_unpack_immutable_files should fail");
}

#[tokio::test]
async fn download_unpack_immutable_files_fails_if_location_is_unknown() {
let total_immutable_files = 2;
let immutable_file_range = ImmutableFileRange::Range(1, total_immutable_files);
let target_dir = TempDir::new("cardano_database_client", "download_unpack").build();
let artifact_downloader = InternalArtifactDownloader::new(
Arc::new(MockFileDownloader::new()),
FeedbackSender::new(&[]),
test_utils::test_logger(),
);

artifact_downloader
.download_unpack_immutable_files(
&[ImmutablesLocation::Unknown {}],
immutable_file_range
.to_range_inclusive(total_immutable_files)
.unwrap(),
&CompressionAlgorithm::default(),
&target_dir,
1,
"download_id",
)
.await
.expect_err("download_unpack_immutable_files should fail");
}

#[tokio::test]
async fn download_unpack_immutable_files_succeeds_if_all_are_retrieved_with_same_location()
{
Expand Down Expand Up @@ -1091,6 +1124,26 @@ mod tests {
.expect_err("download_unpack_ancillary_file should fail");
}

#[tokio::test]
async fn download_unpack_ancillary_files_fails_if_location_is_unknown() {
let target_dir = Path::new(".");
let artifact_downloader = InternalArtifactDownloader::new(
Arc::new(MockFileDownloader::new()),
FeedbackSender::new(&[]),
test_utils::test_logger(),
);

artifact_downloader
.download_unpack_ancillary_file(
&[AncillaryLocation::Unknown {}],
&CompressionAlgorithm::default(),
target_dir,
"download_id",
)
.await
.expect_err("download_unpack_ancillary_file should fail");
}

#[tokio::test]
async fn download_unpack_ancillary_file_succeeds_if_at_least_one_location_is_retrieved() {
let target_dir = Path::new(".");
Expand Down
75 changes: 50 additions & 25 deletions mithril-client/src/cardano_database_client/proving.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ use mithril_common::{
messages::{
CardanoDatabaseDigestListItemMessage, CardanoDatabaseSnapshotMessage, CertificateMessage,
},
StdResult,
};

use crate::{
Expand Down Expand Up @@ -106,32 +107,39 @@ impl InternalArtifactProver {
DigestLocation::CloudStorage { .. } | DigestLocation::Aggregator { .. } => {
self.http_file_downloader.clone()
}
};
let file_downloader_uri: FileDownloaderUri = location.into();
let downloaded = file_downloader
.download_unpack(
&file_downloader_uri,
digest_file_target_dir,
None,
DownloadEvent::Digest {
download_id: download_id.clone(),
},
)
.await;
match downloaded {
Ok(_) => {
self.feedback_sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::DigestDownloadCompleted { download_id },
))
.await;
return Ok(());
DigestLocation::Unknown => {
return Err(anyhow!("Unknown location type to download digest"));
}
Err(e) => {
slog::error!(
self.logger,
"Failed downloading and unpacking digest for location {file_downloader_uri:?}"; "error" => ?e
);
};
let file_downloader_uri: StdResult<FileDownloaderUri> = location.try_into();
if let Ok(file_downloader_uri) = file_downloader_uri {
let downloaded = file_downloader
.download_unpack(
&file_downloader_uri,
digest_file_target_dir,
None,
DownloadEvent::Digest {
download_id: download_id.clone(),
},
)
.await;
match downloaded {
Ok(_) => {
self.feedback_sender
.send_event(MithrilEvent::CardanoDatabase(
MithrilEventCardanoDatabase::DigestDownloadCompleted {
download_id,
},
))
.await;
return Ok(());
}
Err(e) => {
slog::error!(
self.logger,
"Failed downloading and unpacking digest for location {file_downloader_uri:?}"; "error" => ?e
);
}
}
}
}
Expand Down Expand Up @@ -350,6 +358,8 @@ mod tests {

mod download_unpack_digest_file {

use crate::file_downloader::MockFileDownloader;

use super::*;

#[tokio::test]
Expand Down Expand Up @@ -383,6 +393,21 @@ mod tests {
.expect_err("download_unpack_digest_file should fail");
}

#[tokio::test]
async fn fails_if_location_is_unknown() {
let target_dir = Path::new(".");
let artifact_prover = InternalArtifactProver::new(
Arc::new(MockFileDownloader::new()),
FeedbackSender::new(&[]),
test_utils::test_logger(),
);

artifact_prover
.download_unpack_digest_file(&[DigestLocation::Unknown], target_dir)
.await
.expect_err("download_unpack_digest_file should fail");
}

#[tokio::test]
async fn succeeds_if_at_least_one_location_is_retrieved() {
let target_dir = Path::new(".");
Expand Down
82 changes: 73 additions & 9 deletions mithril-client/src/file_downloader/interface.rs
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
use std::{collections::HashMap, path::Path};

use anyhow::anyhow;
use async_trait::async_trait;

use mithril_common::{
entities::{
AncillaryLocation, CompressionAlgorithm, DigestLocation, FileUri, ImmutableFileNumber,
ImmutablesLocation,
},
StdResult,
StdError, StdResult,
};

use crate::feedback::{MithrilEvent, MithrilEventCardanoDatabase};
Expand Down Expand Up @@ -44,6 +45,9 @@ impl FileDownloaderUri {

Ok(immutable_files_range.zip(file_downloader_uris).collect())
}
ImmutablesLocation::Unknown => {
Err(anyhow!("Unknown location type to download immutable"))
}
}
}

Expand All @@ -67,20 +71,28 @@ impl From<FileUri> for FileDownloaderUri {
}
}

impl From<AncillaryLocation> for FileDownloaderUri {
fn from(digest_location: AncillaryLocation) -> Self {
match digest_location {
AncillaryLocation::CloudStorage { uri } => Self::FileUri(FileUri(uri)),
impl TryFrom<AncillaryLocation> for FileDownloaderUri {
type Error = StdError;

fn try_from(location: AncillaryLocation) -> Result<Self, Self::Error> {
match location {
AncillaryLocation::CloudStorage { uri } => Ok(Self::FileUri(FileUri(uri))),
AncillaryLocation::Unknown => {
Err(anyhow!("Unknown location type to download ancillary"))
}
}
}
}

impl From<DigestLocation> for FileDownloaderUri {
fn from(digest_location: DigestLocation) -> Self {
match digest_location {
impl TryFrom<DigestLocation> for FileDownloaderUri {
type Error = StdError;

fn try_from(location: DigestLocation) -> Result<Self, Self::Error> {
match location {
DigestLocation::CloudStorage { uri } | DigestLocation::Aggregator { uri } => {
Self::FileUri(FileUri(uri))
Ok(Self::FileUri(FileUri(uri)))
}
DigestLocation::Unknown => Err(anyhow!("Unknown location type to download digest")),
}
}
}
Expand Down Expand Up @@ -226,6 +238,18 @@ mod tests {
);
}

#[test]
fn immutable_files_location_to_file_downloader_uris_return_error_when_location_is_unknown() {
let immutable_files_location = ImmutablesLocation::Unknown;
let immutable_files_range: Vec<ImmutableFileNumber> = (1..=1).collect();

FileDownloaderUri::expand_immutable_files_location_to_file_downloader_uris(
&immutable_files_location,
&immutable_files_range,
)
.expect_err("expand_immutable_files_location_to_file_downloader_uris should fail");
}

#[test]
fn download_event_type_builds_correct_event() {
let download_event_type = DownloadEvent::Immutable {
Expand Down Expand Up @@ -282,4 +306,44 @@ mod tests {
event,
);
}

#[test]
fn file_downloader_uri_from_ancillary_location() {
let location = AncillaryLocation::CloudStorage {
uri: "http://whatever/ancillary-1".to_string(),
};
let file_downloader_uri: FileDownloaderUri = location.try_into().unwrap();

assert_eq!(
FileDownloaderUri::FileUri(FileUri("http://whatever/ancillary-1".to_string())),
file_downloader_uri
);
}
#[test]
fn file_downloader_uri_from_unknown_ancillary_location() {
let location = AncillaryLocation::Unknown;
let file_downloader_uri: StdResult<FileDownloaderUri> = location.try_into();

file_downloader_uri.expect_err("try_into should fail on Unknown ancillary location");
}

#[test]
fn file_downloader_uri_from_digest_location() {
let location = DigestLocation::CloudStorage {
uri: "http://whatever/digest-1".to_string(),
};
let file_downloader_uri: FileDownloaderUri = location.try_into().unwrap();

assert_eq!(
FileDownloaderUri::FileUri(FileUri("http://whatever/digest-1".to_string())),
file_downloader_uri
);
}
#[test]
fn file_downloader_uri_from_unknown_digest_location() {
let location = DigestLocation::Unknown;
let file_downloader_uri: StdResult<FileDownloaderUri> = location.try_into();

file_downloader_uri.expect_err("try_into should fail on Unknown digest location");
}
}

0 comments on commit 9d855a3

Please sign in to comment.