From 26edf9ca9faa973a374af136eadde0c6a1bccd24 Mon Sep 17 00:00:00 2001
From: Connor Peet <connor@peet.io>
Date: Wed, 3 Apr 2024 17:08:49 -0700
Subject: [PATCH] cli: add progress for extraction (#209499)

For https://github.com/microsoft/vscode-remote-tunnels/issues/724
---
 cli/src/tunnels/code_server.rs |  6 ++++-
 cli/src/util/tar.rs            | 46 ++++++++++++++++++++++------------
 cli/src/util/zipper.rs         |  6 ++++-
 3 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/cli/src/tunnels/code_server.rs b/cli/src/tunnels/code_server.rs
index bb5a9f8d08f95..ca7649adabb87 100644
--- a/cli/src/tunnels/code_server.rs
+++ b/cli/src/tunnels/code_server.rs
@@ -433,7 +433,11 @@ impl<'a> ServerBuilder<'a> {
 				.await?;
 
 				let server_dir = target_dir.join(SERVER_FOLDER_NAME);
-				unzip_downloaded_release(&archive_path, &server_dir, SilentCopyProgress())?;
+				unzip_downloaded_release(
+					&archive_path,
+					&server_dir,
+					self.logger.get_download_logger("server inflate progress:"),
+				)?;
 
 				if !skip_requirements_check().await {
 					let output = capture_command_and_check_status(
diff --git a/cli/src/util/tar.rs b/cli/src/util/tar.rs
index 0a5496411f7f3..fe4d426970021 100644
--- a/cli/src/util/tar.rs
+++ b/cli/src/util/tar.rs
@@ -13,7 +13,7 @@ use tar::Archive;
 use super::errors::wrapdbg;
 use super::io::ReportCopyProgress;
 
-fn should_skip_first_segment(file: &fs::File) -> Result<bool, WrappedError> {
+fn should_skip_first_segment(file: &fs::File) -> Result<(bool, u64), WrappedError> {
 	// unfortunately, we need to re-read the archive here since you cannot reuse
 	// `.entries()`. But this will generally only look at one or two files, so this
 	// should be acceptably speedy... If not, we could hardcode behavior for
@@ -39,17 +39,21 @@ fn should_skip_first_segment(file: &fs::File) -> Result<bool, WrappedError> {
 			.to_owned()
 	};
 
-	let mut had_multiple = false;
+	let mut num_entries = 1;
+	let mut had_different_prefixes = false;
 	for file in entries.flatten() {
-		had_multiple = true;
-		if let Ok(name) = file.path() {
-			if name.iter().next() != Some(&first_name) {
-				return Ok(false);
+		if !had_different_prefixes {
+			if let Ok(name) = file.path() {
+				if name.iter().next() != Some(&first_name) {
+					had_different_prefixes = true;
+				}
 			}
 		}
+
+		num_entries += 1;
 	}
 
-	Ok(had_multiple) // prefix removal is invalid if there's only a single file
+	Ok((!had_different_prefixes && num_entries > 1, num_entries)) // prefix removal is invalid if there's only a single file
 }
 
 pub fn decompress_tarball<T>(
@@ -62,7 +66,11 @@ where
 {
 	let mut tar_gz = fs::File::open(path)
 		.map_err(|e| wrap(e, format!("error opening file {}", path.display())))?;
-	let skip_first = should_skip_first_segment(&tar_gz)?;
+
+	let (skip_first, num_entries) = should_skip_first_segment(&tar_gz)?;
+	let report_progress_every = num_entries / 20;
+	let mut entries_so_far = 0;
+	let mut last_reported_at = 0;
 
 	// reset since skip logic read the tar already:
 	tar_gz
@@ -71,12 +79,19 @@ where
 
 	let tar = GzDecoder::new(tar_gz);
 	let mut archive = Archive::new(tar);
-
-	let results = archive
+	archive
 		.entries()
 		.map_err(|e| wrap(e, format!("error opening archive {}", path.display())))?
 		.filter_map(|e| e.ok())
-		.map(|mut entry| {
+		.try_for_each::<_, Result<_, WrappedError>>(|mut entry| {
+			// approximate progress based on where we are in the archive:
+			entries_so_far += 1;
+			if entries_so_far - last_reported_at > report_progress_every {
+				reporter.report_progress(entries_so_far, num_entries);
+				entries_so_far += 1;
+				last_reported_at = entries_so_far;
+			}
+
 			let entry_path = entry
 				.path()
 				.map_err(|e| wrap(e, "error reading entry path"))?;
@@ -95,12 +110,11 @@ where
 			entry
 				.unpack(&path)
 				.map_err(|e| wrapdbg(e, format!("error unpacking {}", path.display())))?;
-			Ok(path)
-		})
-		.collect::<Result<Vec<PathBuf>, WrappedError>>()?;
 
-	// Tarballs don't have a way to get the number of entries ahead of time
-	reporter.report_progress(results.len() as u64, results.len() as u64);
+			Ok(())
+		})?;
+
+	reporter.report_progress(num_entries, num_entries);
 
 	Ok(())
 }
diff --git a/cli/src/util/zipper.rs b/cli/src/util/zipper.rs
index 0e9939d4db516..69bcf2d23f6d3 100644
--- a/cli/src/util/zipper.rs
+++ b/cli/src/util/zipper.rs
@@ -55,8 +55,12 @@ where
 		.map_err(|e| wrap(e, format!("failed to open zip archive {}", path.display())))?;
 
 	let skip_segments_no = usize::from(should_skip_first_segment(&mut archive));
+	let report_progress_every = archive.len() / 20;
+
 	for i in 0..archive.len() {
-		reporter.report_progress(i as u64, archive.len() as u64);
+		if i % report_progress_every == 0 {
+			reporter.report_progress(i as u64, archive.len() as u64);
+		}
 		let mut file = archive
 			.by_index(i)
 			.map_err(|e| wrap(e, format!("could not open zip entry {}", i)))?;