diff options
Diffstat (limited to 'crates/common/download')
-rw-r--r-- | crates/common/download/Cargo.toml | 27 | ||||
-rw-r--r-- | crates/common/download/examples/simple_download.rs | 23 | ||||
-rw-r--r-- | crates/common/download/src/download.rs | 422 | ||||
-rw-r--r-- | crates/common/download/src/error.rs | 39 | ||||
-rw-r--r-- | crates/common/download/src/lib.rs | 5 |
5 files changed, 516 insertions, 0 deletions
diff --git a/crates/common/download/Cargo.toml b/crates/common/download/Cargo.toml new file mode 100644 index 00000000..5cb353bc --- /dev/null +++ b/crates/common/download/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "download" +version = "0.4.3" +authors = ["thin-edge.io team <info@thin-edge.io>"] +edition = "2018" +license = "Apache-2.0" +description = "download_manager" + +[dependencies] +backoff = { version = "0.3", features = ["tokio"] } +json_sm = { path = "../../core/json_sm" } +log = "0.4" +nix = "0.23" +reqwest = { version = "0.11", default-features = false, features = ["rustls-tls"] } +tedge_utils = { path = "../tedge_utils" } +tempfile = "3.2" +thiserror = "1.0" +tokio = "1.12" +url = "2.2" + +[dev-dependencies] +anyhow = "1.0" +mockito = "0.30" +regex = "1.5" +tempfile = "3.2" +test-case = "1.2" +tokio-test = "0.4" diff --git a/crates/common/download/examples/simple_download.rs b/crates/common/download/examples/simple_download.rs new file mode 100644 index 00000000..821ac0e3 --- /dev/null +++ b/crates/common/download/examples/simple_download.rs @@ -0,0 +1,23 @@ +use anyhow::Result; +use download::Downloader; +use json_sm::DownloadInfo; + +/// This example shows how to use the `downlaoder`. +#[tokio::main] +async fn main() -> Result<()> { + // Create Download metadata. + let url_data = DownloadInfo::new( + "https://file-examples-com.github.io/uploads/2017/02/file_example_CSV_5000.csv", + ); + + // Create downloader instance with desired file path and target directory. + let downloader = Downloader::new("test_download", &None, "/tmp"); + + // Call `download` method to get data from url. + let () = downloader.download(&url_data).await?; + + // Call cleanup method to remove downloaded file if no longer necessary. + let () = downloader.cleanup().await?; + + Ok(()) +} diff --git a/crates/common/download/src/download.rs b/crates/common/download/src/download.rs new file mode 100644 index 00000000..3f88e72e --- /dev/null +++ b/crates/common/download/src/download.rs @@ -0,0 +1,422 @@ +use crate::error::DownloadError; +use backoff::{future::retry, ExponentialBackoff}; +use json_sm::DownloadInfo; +use nix::{ + fcntl::{fallocate, FallocateFlags}, + sys::statvfs, +}; +use std::{ + fs::File, + io::Write, + os::unix::prelude::AsRawFd, + path::{Path, PathBuf}, + time::Duration, +}; + +#[derive(Debug)] +pub struct Downloader { + target_filename: PathBuf, + download_target: PathBuf, +} + +impl Downloader { + pub fn new(name: &str, version: &Option<String>, target_dir_path: impl AsRef<Path>) -> Self { + let mut filename = name.to_string(); + if let Some(version) = version { + filename.push('_'); + filename.push_str(version.as_str()); + } + + let mut download_target = PathBuf::new().join(&target_dir_path).join(&filename); + download_target.set_extension("tmp"); + + let target_filename = PathBuf::new().join(target_dir_path).join(filename); + + Self { + target_filename, + download_target, + } + } + + pub async fn download(&self, url: &DownloadInfo) -> Result<(), DownloadError> { + // Default retry is an exponential retry with a limit of 15 minutes total. + // Let's set some more reasonable retry policy so we don't block the downloads for too long. + + let backoff = ExponentialBackoff { + initial_interval: Duration::from_secs(30), + max_elapsed_time: Some(Duration::from_secs(300)), + ..Default::default() + }; + + let mut response = retry(backoff, || async { + let client = if let Some(json_sm::Auth::Bearer(token)) = &url.auth { + reqwest::Client::new().get(url.url()).bearer_auth(token) + } else { + reqwest::Client::new().get(url.url()) + }; + + match client + .send() + .await + .map_err(|err| { + if err.is_connect() || err.is_builder() { + backoff::Error::Permanent(err) + } else { + log::warn!("Failed to Download. {:?}\nRetrying.", &err); + backoff::Error::Transient(err) + } + })? + .error_for_status() + { + Ok(response) => Ok(response), + + Err(err) => match err.status() { + Some(status_error) if status_error.is_client_error() => { + Err(backoff::Error::Permanent(err)) + } + _ => Err(backoff::Error::Transient(err)), + }, + } + }) + .await?; + + let file_len = match response.content_length() { + Some(len) => len as usize, + None => 0, + }; + let mut file = + create_file_and_try_pre_allocate_space(self.target_filename.as_path(), file_len)?; + + while let Some(chunk) = response.chunk().await? { + if let Err(err) = file.write_all(&chunk) { + drop(file); + std::fs::remove_file(self.target_filename.as_path())?; + return Err(DownloadError::FromIo { + reason: format!("Failed to download the file with an error {}", err), + }); + } + } + + Ok(()) + } + + pub fn filename(&self) -> &Path { + self.target_filename.as_path() + } + + pub async fn cleanup(&self) -> Result<(), DownloadError> { + let _res = tokio::fs::remove_file(&self.target_filename).await; + Ok(()) + } +} + +fn create_file_and_try_pre_allocate_space( + file_path: &Path, + file_len: usize, +) -> Result<File, DownloadError> { + let file = File::create(file_path)?; + if file_len > 0 { + if let Some(root) = file_path.parent() { + let tmpstats = statvfs::statvfs(root)?; + // Reserve 5% of total disk space + let five_percent_disk_space = (tmpstats.blocks() * tmpstats.block_size()) * 5 / 100; + let usable_disk_space = + tmpstats.blocks_free() * tmpstats.block_size() - five_percent_disk_space; + + if file_len >= usable_disk_space as usize { + return Err(DownloadError::InsufficientSpace); + } + // Reserve diskspace + let _ = fallocate( + file.as_raw_fd(), + FallocateFlags::empty(), + 0, + file_len as nix::libc::off_t, + ); + } + } + Ok(file) +} + +#[cfg(test)] +mod tests { + use crate::DownloadError; + + use super::Downloader; + use anyhow::bail; + use json_sm::{Auth, DownloadInfo}; + use mockito::mock; + use nix::sys::statvfs; + use std::io::Write; + use std::path::{Path, PathBuf}; + use tempfile::{NamedTempFile, TempDir}; + use test_case::test_case; + + #[test] + fn construct_downloader_filename() { + let name = "test_download"; + let version = Some("test1".to_string()); + let target_dir_path = PathBuf::from("/tmp"); + + let downloader = Downloader::new(name, &version, &target_dir_path); + + let expected_path = Path::new("/tmp/test_download_test1"); + assert_eq!(downloader.filename(), expected_path); + } + + #[tokio::test] + async fn downloader_download_content_no_auth() -> anyhow::Result<()> { + let _mock1 = mock("GET", "/some_file.txt") + .with_status(200) + .with_body(b"hello") + .create(); + + let name = "test_download"; + let version = Some("test1".to_string()); + let target_dir_path = TempDir::new()?; + + let mut target_url = mockito::server_url(); + target_url.push_str("/some_file.txt"); + + let url = DownloadInfo::new(&target_url); + + let downloader = Downloader::new(&name, &version, target_dir_path.path()); + let () = downloader.download(&url).await?; + + let log_content = std::fs::read(downloader.filename())?; + + assert_eq!("hello".as_bytes(), log_content); + + Ok(()) + } + + #[tokio::test] + async fn downloader_download_with_content_length_larger_than_usable_disk_space( + ) -> anyhow::Result<()> { + let tmpstats = statvfs::statvfs("/tmp")?; + let usable_disk_space = tmpstats.blocks_free() * tmpstats.block_size(); + let _mock1 = mock("GET", "/some_file.txt") + .with_header("content-length", &(usable_disk_space.to_string())) + .create(); + + let name = "test_download_with_length"; + let version = Some("test1".to_string()); + let target_dir_path = TempDir::new()?; + + let mut target_url = mockito::server_url(); + target_url.push_str("/some_file.txt"); + + let url = DownloadInfo::new(&target_url); + + let downloader = Downloader::new(&name, &version, target_dir_path.path()); + match downloader.download(&url).await { + Err(DownloadError::InsufficientSpace) => return Ok(()), + _ => bail!("failed"), + } + } + + #[tokio::test] + async fn downloader_download_with_reasonable_content_length() -> anyhow::Result<()> { + let file = create_file_with_size(10 * 1024 * 1024)?; + let file_path = file.into_temp_path(); + + let _mock1 = mock("GET", "/some_file.txt") + .with_body_from_file(&file_path) + .create(); + + let name = "test_download_with_length"; + let version = Some("test1".to_string()); + let target_dir_path = TempDir::new()?; + + let mut target_url = mockito::server_url(); + target_url.push_str("/some_file.txt"); + + let url = DownloadInfo::new(&target_url); + + let downloader = Downloader::new(&name, &version, target_dir_path.path()); + + match downloader.download(&url).await { + Ok(()) => { + let log_content = std::fs::read(downloader.filename())?; + let expected_content = std::fs::read(file_path)?; + assert_eq!(log_content, expected_content); + return Ok(()); + } + _ => bail!("failed"), + } + } + + #[tokio::test] + async fn downloader_download_verify_file_content() -> anyhow::Result<()> { + let file = create_file_with_size(10)?; + + let _mock1 = mock("GET", "/some_file.txt") + .with_body_from_file(file.into_temp_path()) + .create(); + + let name = "test_download_with_length"; + let version = Some("test1".to_string()); + let target_dir_path = TempDir::new()?; + + let mut target_url = mockito::server_url(); + target_url.push_str("/some_file.txt"); + + let url = DownloadInfo::new(&target_url); + + let downloader = Downloader::new(&name, &version, target_dir_path.path()); + downloader.download(&url).await?; + + let log_content = std::fs::read(downloader.filename())?; + + assert_eq!("Some data!".as_bytes(), log_content); + + Ok(()) + } + + #[tokio::test] + async fn downloader_download_without_content_length() -> anyhow::Result<()> { + let _mock1 = mock("GET", "/some_file.txt").create(); + + let name = "test_download_without_length"; + let version = Some("test1".to_string()); + let target_dir_path = TempDir::new()?; + + let mut target_url = mockito::server_url(); + target_url.push_str("/some_file.txt"); + + let url = DownloadInfo::new(&target_url); + + let downloader = Downloader::new(&name, &version, target_dir_path.path()); + match downloader.download(&url).await { + Ok(()) => { + assert_eq!("".as_bytes(), std::fs::read(downloader.filename())?); + return Ok(()); + } + _ => { + bail!("failed") + } + } + } + + // Parameters: + // + // status code + // + // bearer token boolean + // + // maybe url + // + // expected std error + // + // description + #[test_case( + 200, + false, + Some("not_a_url"), + "builder error: relative URL without a base" + ; "builder error" + )] + #[test_case( + 200, + true, + Some("not_a_url"), + "builder error: relative URL without a base" + ; "builder error with auth" + )] + #[test_case( + 200, + false, + Some("http://not_a_url"), + "dns error: failed to lookup address information" + ; "dns error" + )] + #[test_case( + 200, + true, + Some("http://not_a_url"), + "dns error: failed to lookup address information" + ; "dns error with auth" + )] + #[test_case( + 404, + false, + None, + "404 Not Found" + ; "client error" + )] + #[test_case( + 404, + true, + None, + "404 Not Found" + ; "client error with auth" + )] + #[tokio::test] + async fn downloader_download_processing_error( + status_code: usize, + with_token: bool, + url: Option<&str>, + expected_err: &str, + ) -> anyhow::Result<()> { + let name = "test_download"; + let version = Some("test1".to_string()); + let target_dir_path = TempDir::new()?; + + // bearer/no bearer setup + let _mock1 = { + if with_token { + mock("GET", "/some_file.txt") + .match_header("authorization", "Bearer token") + .with_status(status_code) + .create() + } else { + mock("GET", "/some_file.txt") + .with_status(status_code) + .create() + } + }; + + // url/no url setup + let url = { + if let Some(url) = url { + DownloadInfo::new(&url) + } else { + let mut target_url = mockito::server_url(); + target_url.push_str("/some_file.txt"); + DownloadInfo::new(&target_url) + } + }; + + // applying token if `with_token` = true + let url = { + if with_token { + url.with_auth(Auth::Bearer(String::from("token"))) + } else { + url + } + }; + + let downloader = Downloader::new(&name, &version, target_dir_path.path()); + match downloader.download(&url).await { + Ok(_success) => anyhow::bail!("Expected client error."), + Err(err) => { + assert!(err.to_string().contains(expected_err)); + } + }; + Ok(()) + } + + fn create_file_with_size(size: usize) -> Result<NamedTempFile, anyhow::Error> { + let mut file = NamedTempFile::new()?; + let data: String = "Some data!".into(); + let loops = size / data.len(); + let mut buffer = String::with_capacity(size); + for _ in 0..loops { + buffer.push_str("Some data!"); + } + + file.write_all(buffer.as_bytes())?; + + Ok(file) + } +} diff --git a/crates/common/download/src/error.rs b/crates/common/download/src/error.rs new file mode 100644 index 00000000..4647dfef --- /dev/null +++ b/crates/common/download/src/error.rs @@ -0,0 +1,39 @@ +#[derive(Debug, thiserror::Error)] +pub enum DownloadError { + #[error(transparent)] + FromBackoff(#[from] backoff::Error<reqwest::Error>), + + #[error(transparent)] + FromElapsed(#[from] tokio::time::error::Elapsed), + + #[error("I/O error: {reason:?}")] + FromIo { reason: String }, + + #[error("JSON parse error: {reason:?}")] + FromReqwest { reason: String }, + + #[error(transparent)] + FromUrlParse(#[from] url::ParseError), + + #[error(transparent)] + FromNix(#[from] nix::Error), + + #[error("Not enough disk space")] + InsufficientSpace, +} + +impl From<reqwest::Error> for DownloadError { + fn from(err: reqwest::Error) -> Self { + DownloadError::FromReqwest { + reason: format!("{}", err), + } + } +} + +impl From<std::io::Error> for DownloadError { + fn from(err: std::io::Error) -> Self { + DownloadError::FromIo { + reason: format!("{}", err), + } + } +} diff --git a/crates/common/download/src/lib.rs b/crates/common/download/src/lib.rs new file mode 100644 index 00000000..205dd8b8 --- /dev/null +++ b/crates/common/download/src/lib.rs @@ -0,0 +1,5 @@ +mod download; +mod error; + +pub use crate::download::Downloader; +pub use crate::error::DownloadError; |