use std::collections::HashSet;
use std::io::{self, BufWriter, Read, Seek};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use rayon::prelude::*;
use crate::archive::{MappedArchive, sanitize_relative_path};
use crate::codec::EncodingKind;
use crate::error::XzipError;
use crate::filter::PathFilter;
use crate::fs_util;
use crate::verbosity::Verbosity;
const WRITE_BUFFER_SIZE: usize = 128 * 1024;
const PROGRESS_INTERVAL: usize = 250;
const PARALLEL_MIN_FILES: usize = 256;
const PARALLEL_MIN_BYTES: u64 = 4 * 1024 * 1024;
#[cfg(windows)]
const PARALLEL_MAX_WRITERS: usize = 4;
#[derive(Debug, Clone)]
struct ExtractPlan {
index: usize,
outpath: PathBuf,
display_name: String,
is_dir: bool,
uncompressed_size: u64,
compressed_size: u64,
}
#[derive(Debug, Default)]
struct PhaseTimings {
scan: Duration,
mkdir: Duration,
extract: Duration,
}
fn ensure_dir(path: &Path, created: &mut HashSet<PathBuf>) -> Result<bool, XzipError> {
let path = path.to_path_buf();
if created.contains(&path) {
return Ok(false);
}
fs_util::create_dir_all(&path)?;
created.insert(path);
Ok(true)
}
fn available_cpus() -> usize {
std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1)
}
fn total_uncompressed_bytes(plans: &[ExtractPlan]) -> u64 {
plans
.iter()
.filter(|plan| !plan.is_dir)
.map(|plan| plan.uncompressed_size)
.sum()
}
fn resolve_jobs(jobs: usize, file_count: usize, total_bytes: u64) -> usize {
let jobs = match jobs {
0 => {
if file_count >= PARALLEL_MIN_FILES && total_bytes >= PARALLEL_MIN_BYTES {
available_cpus()
} else {
1
}
}
1 => 1,
n => n,
};
#[cfg(windows)]
if jobs > 1 {
return jobs.min(PARALLEL_MAX_WRITERS);
}
jobs
}
fn scan_entries<R: Read + Seek>(
archive: &mut zip::read::ZipArchive<R>,
output_dir: &Path,
encoding: EncodingKind,
filter: &PathFilter,
) -> Result<Vec<ExtractPlan>, XzipError> {
let mut plans = Vec::new();
for i in 0..archive.len() {
let entry = archive.by_index(i)?;
let raw_name = entry.name_raw();
let decoded_name = encoding.decode(raw_name)?;
if !filter.allows(&decoded_name) {
continue;
}
let safe_rel = sanitize_relative_path(&decoded_name)?;
let outpath = output_dir.join(&safe_rel);
plans.push(ExtractPlan {
index: i,
outpath,
display_name: decoded_name,
is_dir: entry.is_dir(),
uncompressed_size: entry.size(),
compressed_size: entry.compressed_size(),
});
}
Ok(plans)
}
fn create_output_dirs(output_dir: &Path, plans: &[ExtractPlan]) -> Result<usize, XzipError> {
let mut created = HashSet::new();
ensure_dir(output_dir, &mut created)?;
let mut dirs_created = 0usize;
for plan in plans {
if plan.is_dir {
if ensure_dir(&plan.outpath, &mut created)? {
dirs_created += 1;
}
} else if let Some(parent) = plan.outpath.parent() {
if ensure_dir(parent, &mut created)? {
dirs_created += 1;
}
}
}
Ok(dirs_created)
}
fn print_file_progress(verbosity: Verbosity, plan: &ExtractPlan) {
if verbosity.is_extra() {
eprintln!(
"extracting: {} ({} -> {} bytes)",
plan.display_name, plan.uncompressed_size, plan.compressed_size
);
}
}
fn maybe_print_throttled_progress(processed: usize, total: usize) {
if processed == total || processed % PROGRESS_INTERVAL == 0 {
eprintln!("extracting: {processed}/{total} files...");
}
}
fn extract_file<R: Read + Seek>(
archive: &mut zip::read::ZipArchive<R>,
plan: &ExtractPlan,
) -> Result<(), XzipError> {
let mut entry = archive.by_index(plan.index)?;
let mut outfile = BufWriter::with_capacity(
WRITE_BUFFER_SIZE,
fs_util::create_output_file(&plan.outpath)?,
);
io::copy(&mut entry, &mut outfile)?;
Ok(())
}
fn extract_sequential<R: Read + Seek>(
archive: &mut zip::read::ZipArchive<R>,
file_plans: &[ExtractPlan],
verbosity: Verbosity,
) -> Result<(), XzipError> {
let total = file_plans.len();
for (processed, plan) in file_plans.iter().enumerate() {
if verbosity.is_verbose() {
if verbosity.is_extra() {
print_file_progress(verbosity, plan);
} else {
maybe_print_throttled_progress(processed + 1, total);
}
}
extract_file(archive, plan)?;
}
Ok(())
}
fn extract_parallel(
mapped: &MappedArchive,
file_plans: &[ExtractPlan],
jobs: usize,
verbosity: Verbosity,
) -> Result<(), XzipError> {
let total = file_plans.len();
let processed = AtomicUsize::new(0);
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(jobs)
.build()
.map_err(|err| XzipError::Io(io::Error::other(err.to_string())))?;
pool.install(|| {
file_plans.par_iter().try_for_each_init(
|| {
mapped
.open_archive()
.expect("archive was readable during scan and should remain readable")
},
|archive, plan| {
extract_file(archive, plan)?;
if verbosity.is_verbose() && !verbosity.is_extra() {
let count = processed.fetch_add(1, Ordering::Relaxed) + 1;
maybe_print_throttled_progress(count, total);
}
Ok::<(), XzipError>(())
},
)
})?;
Ok(())
}
fn format_duration(duration: Duration) -> String {
format!("{:.2}s", duration.as_secs_f64())
}
fn print_summary(
files_extracted: usize,
dirs_created: usize,
elapsed: Duration,
effective_jobs: usize,
profile: bool,
timings: &PhaseTimings,
) {
let file_label = if files_extracted == 1 {
"file"
} else {
"files"
};
eprintln!(
"{files_extracted} {file_label} extracted ({dirs_created} directories) in {:.2}s [jobs={effective_jobs}]",
elapsed.as_secs_f64()
);
if profile {
eprintln!(
"profile: scan={} | mkdir={} | extract={} | jobs={effective_jobs}",
format_duration(timings.scan),
format_duration(timings.mkdir),
format_duration(timings.extract),
);
}
}
#[allow(clippy::too_many_arguments)]
pub fn unpack_archive(
input_zip: &Path,
output_dir: &Path,
encoding: EncodingKind,
include: &[String],
exclude: &[String],
verbosity: Verbosity,
jobs: usize,
profile: bool,
) -> Result<(), XzipError> {
let started = Instant::now();
let mut timings = PhaseTimings::default();
let mapped = MappedArchive::open(input_zip)?;
let scan_started = Instant::now();
let mut archive = mapped.open_archive()?;
let filter = PathFilter::new(include, exclude)?;
let plans = scan_entries(&mut archive, output_dir, encoding, &filter)?;
timings.scan = scan_started.elapsed();
let mkdir_started = Instant::now();
let dirs_created = create_output_dirs(output_dir, &plans)?;
timings.mkdir = mkdir_started.elapsed();
let file_plans: Vec<_> = plans.iter().filter(|plan| !plan.is_dir).cloned().collect();
let files_extracted = file_plans.len();
let total_bytes = total_uncompressed_bytes(&plans);
let effective_jobs = if verbosity.is_extra() {
1
} else {
resolve_jobs(jobs, files_extracted, total_bytes)
};
let extract_started = Instant::now();
if effective_jobs == 1 {
extract_sequential(&mut archive, &file_plans, verbosity)?;
} else {
extract_parallel(&mapped, &file_plans, effective_jobs, verbosity)?;
}
timings.extract = extract_started.elapsed();
if verbosity.is_verbose() || profile {
print_summary(
files_extracted,
dirs_created,
started.elapsed(),
effective_jobs,
profile,
&timings,
);
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{PARALLEL_MIN_BYTES, PARALLEL_MIN_FILES, resolve_jobs};
#[test]
fn auto_jobs_stays_sequential_for_small_archives() {
assert_eq!(
resolve_jobs(0, PARALLEL_MIN_FILES - 1, PARALLEL_MIN_BYTES),
1
);
assert_eq!(
resolve_jobs(0, PARALLEL_MIN_FILES, PARALLEL_MIN_BYTES - 1),
1
);
}
#[test]
fn auto_jobs_parallelizes_large_archives() {
#[cfg(not(windows))]
assert!(resolve_jobs(0, PARALLEL_MIN_FILES, PARALLEL_MIN_BYTES) > 1);
#[cfg(windows)]
assert_eq!(
resolve_jobs(0, PARALLEL_MIN_FILES, PARALLEL_MIN_BYTES),
super::PARALLEL_MAX_WRITERS.min(
std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1)
)
);
}
#[test]
fn explicit_jobs_are_honored() {
assert_eq!(resolve_jobs(1, 10_000, PARALLEL_MIN_BYTES * 10), 1);
assert_eq!(resolve_jobs(4, 10_000, PARALLEL_MIN_BYTES * 10), {
#[cfg(windows)]
{
4.min(super::PARALLEL_MAX_WRITERS)
}
#[cfg(not(windows))]
{
4
}
});
}
}