From 5b7a2f1cfd83e6cb9d93e011a7932547501d0aed Mon Sep 17 00:00:00 2001 From: kate Date: Wed, 14 Aug 2024 09:45:00 +0800 Subject: [PATCH 1/2] Add progress bar for zstd decompression --- src/ctx.rs | 34 +++++--------- src/error.rs | 9 ++++ src/extract.rs | 124 +++++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + 4 files changed, 145 insertions(+), 23 deletions(-) create mode 100644 src/extract.rs diff --git a/src/ctx.rs b/src/ctx.rs index 7a644fe..0be75f5 100644 --- a/src/ctx.rs +++ b/src/ctx.rs @@ -281,13 +281,16 @@ impl RunnerContext { let tar_tmp_location = helper::get_tmp_file("data.tar".to_string()); - let zstd_file = std::fs::File::open(&zstd_location)?; - let mut tar_tmp_file = std::fs::File::create_new(&tar_tmp_location)?; - zstd::stream::copy_decode(zstd_file, &tar_tmp_file)?; - tar_tmp_file = std::fs::File::open(&tar_tmp_location)?; // we do this again to make sure that the tar is properly opened. - - let mut archive = tar::Archive::new(&tar_tmp_file); - let x = archive.unpack(&out_dir); + if let Err(e) = crate::extract::decompress_zstd(zstd_location.clone(), tar_tmp_location.clone(), true) { + debug!("{:#?}", e); + error!("[RunnerContext::extract_package] Failed to decompress file {} ({:})", zstd_location, e); + return Err(e); + } + if let Err(e) = crate::extract::unpack_tarball(tar_tmp_location.clone(), out_dir, true) { + debug!("{:#?}", e); + error!("[RunnerContext::extract_package] Failed to unpack tarball {} ({:})", tar_tmp_location, e); + return Err(e); + } if helper::file_exists(tar_tmp_location.clone()) { if let Err(e) = std::fs::remove_file(tar_tmp_location.clone()) @@ -303,22 +306,7 @@ impl RunnerContext ); } } - match x - { - Err(e) => - { - let xe = BeansError::TarExtractFailure { - src_file: tar_tmp_location, - target_dir: out_dir, - error: e, - backtrace: Backtrace::capture() - }; - trace!("[RunnerContext::extract_package] {:}\n{:#?}", xe, xe); - sentry::capture_error(&xe); - Err(xe) - } - Ok(_) => Ok(()) - } + Ok(()) } #[cfg(target_os = "linux")] diff --git a/src/error.rs b/src/error.rs index e6a5d4e..c7109ca 100644 --- a/src/error.rs +++ b/src/error.rs @@ -52,6 +52,15 @@ pub enum BeansError error: std::io::Error, backtrace: Backtrace }, + #[error("Failed to extract item {link_name} to directory {target_dir} ({error:})")] + TarUnpackItemFailure + { + src_file: String, + target_dir: String, + link_name: String, + error: std::io::Error, + backtrace: Backtrace + }, #[error("Failed to send request ({error:})")] Reqwest { diff --git a/src/extract.rs b/src/extract.rs new file mode 100644 index 0000000..7ad2775 --- /dev/null +++ b/src/extract.rs @@ -0,0 +1,124 @@ +use log::info; +use std::{backtrace::Backtrace, fs::File, io::Read}; +use indicatif::{ProgressBar, ProgressStyle}; +use zstd::stream::read::Decoder as ZstdDecoder; + +use crate::BeansError; + +pub fn unpack_tarball(tarball_location: String, output_directory: String, show_progress: bool) -> Result<(), BeansError> +{ + let tarball = match File::open(&tarball_location) { + Ok(x) => x, + Err(e) => { + return Err(BeansError::TarExtractFailure { + src_file: tarball_location, + target_dir: output_directory, + error: e, + backtrace: Backtrace::capture() + }); + } + }; + let mut archive = tar::Archive::new(&tarball); + if show_progress { + let archive_entries = match archive.entries() + { + Ok(v) => v, + Err(e) => + { + return Err(BeansError::TarExtractFailure { + src_file: tarball_location, + target_dir: output_directory, + error: e, + backtrace: Backtrace::capture() + }); + } + }; + let archive_entry_count = (&archive_entries.count()).clone() as u64; + info!("Extracting {} files", archive_entry_count); + + let pb = ProgressBar::new(archive_entry_count); + pb.set_style(ProgressStyle::with_template("{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos}/{len} ({eta})") + .unwrap() + .with_key("eta", |state: &indicatif::ProgressState, w: &mut dyn std::fmt::Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) + .progress_chars("#>-")); + pb.set_message("Extracting files"); + + archive = tar::Archive::new(&tarball); + match archive.entries() { + Ok(etrs) => + { + for entry in etrs { + match entry { + Ok(mut x) => { + let ln = x.link_name(); + pb.set_message("Extracting files"); + let mut filename = String::new(); + if let Ok(n) = ln { + if let Some(p) = n { + if let Some(s) = p.to_str() { + pb.set_message(format!("{:}", s)); + filename = String::from(s); + } + } + } + if let Err(e) = x.unpack_in(&output_directory) { + return Err(BeansError::TarUnpackItemFailure { + src_file: tarball_location, + target_dir: output_directory, + link_name: filename, + error: e, + backtrace: Backtrace::capture() + }); + } + pb.inc(1); + }, + Err(e) => { + return Err(BeansError::TarExtractFailure { + src_file: tarball_location, + target_dir: output_directory, + error: e, + backtrace: Backtrace::capture() + }); + } + }; + } + }, + Err(e) => + { + return Err(BeansError::TarExtractFailure { + src_file: tarball_location, + target_dir: output_directory, + error: e, + backtrace: Backtrace::capture() + }); + } + }; + pb.finish(); + } else { + archive.unpack(output_directory); + } + return Ok(()); +} +pub fn decompress_zstd(zstd_location: String, output_file: String, show_progress: bool) -> Result<(), BeansError> +{ + let zstd_file = File::open(&zstd_location)?; + let zstd_file_length = &zstd_file.metadata()?.len(); + let mut tar_tmp_file = File::create_new(&output_file)?; + if show_progress { + let decoder = ZstdDecoder::new(zstd_file)?; + // estimate extracted size as x2 since idk how to get the decompressed size with zstd + let pb_decompress = ProgressBar::new((zstd_file_length.clone() * 2) as u64); + pb_decompress + .set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})") + .unwrap() + .with_key("eta", |state: &indicatif::ProgressState, w: &mut dyn std::fmt::Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) + .progress_chars("#>-")); + + std::io::copy(&mut pb_decompress.wrap_read(decoder), &mut tar_tmp_file).expect("Failed to decompress file"); + pb_decompress.finish(); + } else { + zstd::stream::copy_decode(zstd_file, &tar_tmp_file)?; + } + + Ok(()) +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index 36be658..468cd39 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,6 +25,7 @@ pub mod butler; pub mod flags; pub mod gui; pub mod logger; +pub mod extract; /// NOTE do not change, fetches from the version of beans-rs on build pub const VERSION: &str = env!("CARGO_PKG_VERSION"); From 79a17019f296b9f66120365642d0a7e922fe65d0 Mon Sep 17 00:00:00 2001 From: kate Date: Wed, 14 Aug 2024 09:46:52 +0800 Subject: [PATCH 2/2] Use cargo fmt --- src/ctx.rs | 17 +++++++--- src/extract.rs | 85 ++++++++++++++++++++++++++++++++++---------------- src/lib.rs | 26 ++++++++------- 3 files changed, 86 insertions(+), 42 deletions(-) diff --git a/src/ctx.rs b/src/ctx.rs index 0be75f5..5fb6ece 100644 --- a/src/ctx.rs +++ b/src/ctx.rs @@ -281,14 +281,23 @@ impl RunnerContext { let tar_tmp_location = helper::get_tmp_file("data.tar".to_string()); - if let Err(e) = crate::extract::decompress_zstd(zstd_location.clone(), tar_tmp_location.clone(), true) { + if let Err(e) = + crate::extract::decompress_zstd(zstd_location.clone(), tar_tmp_location.clone(), true) + { debug!("{:#?}", e); - error!("[RunnerContext::extract_package] Failed to decompress file {} ({:})", zstd_location, e); + error!( + "[RunnerContext::extract_package] Failed to decompress file {} ({:})", + zstd_location, e + ); return Err(e); } - if let Err(e) = crate::extract::unpack_tarball(tar_tmp_location.clone(), out_dir, true) { + if let Err(e) = crate::extract::unpack_tarball(tar_tmp_location.clone(), out_dir, true) + { debug!("{:#?}", e); - error!("[RunnerContext::extract_package] Failed to unpack tarball {} ({:})", tar_tmp_location, e); + error!( + "[RunnerContext::extract_package] Failed to unpack tarball {} ({:})", + tar_tmp_location, e + ); return Err(e); } if helper::file_exists(tar_tmp_location.clone()) diff --git a/src/extract.rs b/src/extract.rs index 7ad2775..63bb0f8 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1,15 +1,25 @@ +use std::{backtrace::Backtrace, + fs::File, + io::Read}; + +use indicatif::{ProgressBar, + ProgressStyle}; use log::info; -use std::{backtrace::Backtrace, fs::File, io::Read}; -use indicatif::{ProgressBar, ProgressStyle}; use zstd::stream::read::Decoder as ZstdDecoder; use crate::BeansError; -pub fn unpack_tarball(tarball_location: String, output_directory: String, show_progress: bool) -> Result<(), BeansError> +pub fn unpack_tarball( + tarball_location: String, + output_directory: String, + show_progress: bool +) -> Result<(), BeansError> { - let tarball = match File::open(&tarball_location) { + let tarball = match File::open(&tarball_location) + { Ok(x) => x, - Err(e) => { + Err(e) => + { return Err(BeansError::TarExtractFailure { src_file: tarball_location, target_dir: output_directory, @@ -19,7 +29,8 @@ pub fn unpack_tarball(tarball_location: String, output_directory: String, show_p } }; let mut archive = tar::Archive::new(&tarball); - if show_progress { + if show_progress + { let archive_entries = match archive.entries() { Ok(v) => v, @@ -35,33 +46,41 @@ pub fn unpack_tarball(tarball_location: String, output_directory: String, show_p }; let archive_entry_count = (&archive_entries.count()).clone() as u64; info!("Extracting {} files", archive_entry_count); - + let pb = ProgressBar::new(archive_entry_count); pb.set_style(ProgressStyle::with_template("{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos}/{len} ({eta})") .unwrap() .with_key("eta", |state: &indicatif::ProgressState, w: &mut dyn std::fmt::Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) .progress_chars("#>-")); pb.set_message("Extracting files"); - + archive = tar::Archive::new(&tarball); - match archive.entries() { + match archive.entries() + { Ok(etrs) => { - for entry in etrs { - match entry { - Ok(mut x) => { + for entry in etrs + { + match entry + { + Ok(mut x) => + { let ln = x.link_name(); pb.set_message("Extracting files"); let mut filename = String::new(); - if let Ok(n) = ln { - if let Some(p) = n { - if let Some(s) = p.to_str() { + if let Ok(n) = ln + { + if let Some(p) = n + { + if let Some(s) = p.to_str() + { pb.set_message(format!("{:}", s)); filename = String::from(s); } } } - if let Err(e) = x.unpack_in(&output_directory) { + if let Err(e) = x.unpack_in(&output_directory) + { return Err(BeansError::TarUnpackItemFailure { src_file: tarball_location, target_dir: output_directory, @@ -71,8 +90,9 @@ pub fn unpack_tarball(tarball_location: String, output_directory: String, show_p }); } pb.inc(1); - }, - Err(e) => { + } + Err(e) => + { return Err(BeansError::TarExtractFailure { src_file: tarball_location, target_dir: output_directory, @@ -82,7 +102,7 @@ pub fn unpack_tarball(tarball_location: String, output_directory: String, show_p } }; } - }, + } Err(e) => { return Err(BeansError::TarExtractFailure { @@ -94,31 +114,42 @@ pub fn unpack_tarball(tarball_location: String, output_directory: String, show_p } }; pb.finish(); - } else { + } + else + { archive.unpack(output_directory); } return Ok(()); } -pub fn decompress_zstd(zstd_location: String, output_file: String, show_progress: bool) -> Result<(), BeansError> +pub fn decompress_zstd( + zstd_location: String, + output_file: String, + show_progress: bool +) -> Result<(), BeansError> { let zstd_file = File::open(&zstd_location)?; let zstd_file_length = &zstd_file.metadata()?.len(); let mut tar_tmp_file = File::create_new(&output_file)?; - if show_progress { + if show_progress + { let decoder = ZstdDecoder::new(zstd_file)?; - // estimate extracted size as x2 since idk how to get the decompressed size with zstd + // estimate extracted size as x2 since idk how to get the decompressed size with + // zstd let pb_decompress = ProgressBar::new((zstd_file_length.clone() * 2) as u64); pb_decompress .set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})") .unwrap() .with_key("eta", |state: &indicatif::ProgressState, w: &mut dyn std::fmt::Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) .progress_chars("#>-")); - - std::io::copy(&mut pb_decompress.wrap_read(decoder), &mut tar_tmp_file).expect("Failed to decompress file"); + + std::io::copy(&mut pb_decompress.wrap_read(decoder), &mut tar_tmp_file) + .expect("Failed to decompress file"); pb_decompress.finish(); - } else { + } + else + { zstd::stream::copy_decode(zstd_file, &tar_tmp_file)?; } Ok(()) -} \ No newline at end of file +} diff --git a/src/lib.rs b/src/lib.rs index 468cd39..f2c7405 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,10 +22,10 @@ pub use error::*; pub mod appvar; pub mod butler; +pub mod extract; pub mod flags; pub mod gui; pub mod logger; -pub mod extract; /// NOTE do not change, fetches from the version of beans-rs on build pub const VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -60,33 +60,36 @@ pub fn data_dir() -> String /// Temporary directory which is specified by `ADASTRAL_TMPDIR`. /// -/// Will return `None` when the environment variable couldn't be found, or it's an empty string. +/// Will return `None` when the environment variable couldn't be found, or it's +/// an empty string. pub fn env_custom_tmpdir() -> Option { let s = helper::try_get_env_var(String::from("ADASTRAL_TMPDIR")); - match s { - Some(x) => match x.trim().is_empty() { + match s + { + Some(x) => match x.trim().is_empty() + { true => None, false => Some(x) }, None => s } } -/// Return `true` when the environment variable `BEANS_DEBUG` or `ADASTRAL_DEBUG` exists and -/// equals `1` or `true`. +/// Return `true` when the environment variable `BEANS_DEBUG` or +/// `ADASTRAL_DEBUG` exists and equals `1` or `true`. pub fn env_debug() -> bool { check_env_bool("BEANS_DEBUG") || check_env_bool("ADASTRAL_DEBUG") } -/// Return `true` when the environment variable `BEANS_HEADLESS` or `ADASTRAL_HEADLESS` exists and -/// equals `1` or `true`. +/// Return `true` when the environment variable `BEANS_HEADLESS` or +/// `ADASTRAL_HEADLESS` exists and equals `1` or `true`. pub fn env_headless() -> bool { check_env_bool("BEANS_HEADLESS") || check_env_bool("ADASTRAL_HEADLESS") } -/// Return `true` when the environment variable exists, and it's value equals `1` or `true (when -/// trimmed and made lowercase). +/// Return `true` when the environment variable exists, and it's value equals +/// `1` or `true (when trimmed and made lowercase). fn check_env_bool>(key: K) -> bool { std::env::var(key).is_ok_and(|x| { @@ -112,7 +115,8 @@ pub fn has_gui_support() -> bool } } - if env_headless() { + if env_headless() + { return true; }