From d8d600b1d970f58d8462ebb1e79c03d5fb52dda9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anders=20F=20Bj=C3=B6rklund?= <anders.f.bjorklund@gmail.com>
Date: Sat, 15 Jul 2023 09:41:08 +0200
Subject: [PATCH] Add progress bar for decompress image
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

[NO NEW TESTS NEEDED]

Signed-off-by: Anders F Björklund <anders.f.bjorklund@gmail.com>
---
 pkg/machine/pull.go | 78 +++++++++++++++++++++++++++++++++++----------
 1 file changed, 61 insertions(+), 17 deletions(-)

diff --git a/pkg/machine/pull.go b/pkg/machine/pull.go
index c48be52e8c..2aadcbd4fd 100644
--- a/pkg/machine/pull.go
+++ b/pkg/machine/pull.go
@@ -215,43 +215,60 @@ func Decompress(localPath, uncompressedPath string) error {
 	if strings.HasSuffix(localPath, ".zip") {
 		isZip = true
 	}
+	prefix := "Copying uncompressed file"
 	compressionType := archive.DetectCompression(sourceFile)
 	if compressionType != archive.Uncompressed || isZip {
-		fmt.Println("Extracting compressed file")
+		prefix = "Extracting compressed file"
 	}
+	prefix += ": " + filepath.Base(uncompressedPath)
 	if compressionType == archive.Xz {
-		return decompressXZ(localPath, uncompressedFileWriter)
+		return decompressXZ(prefix, localPath, uncompressedFileWriter)
 	}
 	if isZip && runtime.GOOS == "windows" {
-		return decompressZip(localPath, uncompressedFileWriter)
+		return decompressZip(prefix, localPath, uncompressedFileWriter)
 	}
-	return decompressEverythingElse(localPath, uncompressedFileWriter)
+	return decompressEverythingElse(prefix, localPath, uncompressedFileWriter)
 }
 
 // Will error out if file without .Xz already exists
 // Maybe extracting then renaming is a good idea here..
 // depends on Xz: not pre-installed on mac, so it becomes a brew dependency
-func decompressXZ(src string, output io.WriteCloser) error {
+func decompressXZ(prefix string, src string, output io.WriteCloser) error {
 	var read io.Reader
 	var cmd *exec.Cmd
+
+	stat, err := os.Stat(src)
+	if err != nil {
+		return err
+	}
+	file, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	p, bar := progressBar(prefix, stat.Size(), prefix+": done")
+	proxyReader := bar.ProxyReader(file)
+	defer func() {
+		if err := proxyReader.Close(); err != nil {
+			logrus.Error(err)
+		}
+	}()
+
 	// Prefer Xz utils for fastest performance, fallback to go xi2 impl
 	if _, err := exec.LookPath("xz"); err == nil {
-		cmd = exec.Command("xz", "-d", "-c", "-k", src)
+		cmd = exec.Command("xz", "-d", "-c")
+		cmd.Stdin = proxyReader
 		read, err = cmd.StdoutPipe()
 		if err != nil {
 			return err
 		}
 		cmd.Stderr = os.Stderr
 	} else {
-		file, err := os.Open(src)
-		if err != nil {
-			return err
-		}
-		defer file.Close()
 		// This XZ implementation is reliant on buffering. It is also 3x+ slower than XZ utils.
 		// Consider replacing with a faster implementation (e.g. xi2) if podman machine is
 		// updated with a larger image for the distribution base.
-		buf := bufio.NewReader(file)
+		buf := bufio.NewReader(proxyReader)
 		read, err = xz.NewReader(buf)
 		if err != nil {
 			return err
@@ -268,18 +285,35 @@ func decompressXZ(src string, output io.WriteCloser) error {
 	}()
 
 	if cmd != nil {
-		return cmd.Run()
+		err := cmd.Start()
+		if err != nil {
+			return err
+		}
+		p.Wait()
+		return cmd.Wait()
 	}
 	<-done
+	p.Wait()
 	return nil
 }
 
-func decompressEverythingElse(src string, output io.WriteCloser) error {
+func decompressEverythingElse(prefix string, src string, output io.WriteCloser) error {
+	stat, err := os.Stat(src)
+	if err != nil {
+		return err
+	}
 	f, err := os.Open(src)
 	if err != nil {
 		return err
 	}
-	uncompressStream, _, err := compression.AutoDecompress(f)
+	p, bar := progressBar(prefix, stat.Size(), prefix+": done")
+	proxyReader := bar.ProxyReader(f)
+	defer func() {
+		if err := proxyReader.Close(); err != nil {
+			logrus.Error(err)
+		}
+	}()
+	uncompressStream, _, err := compression.AutoDecompress(proxyReader)
 	if err != nil {
 		return err
 	}
@@ -293,10 +327,11 @@ func decompressEverythingElse(src string, output io.WriteCloser) error {
 	}()
 
 	_, err = io.Copy(output, uncompressStream)
+	p.Wait()
 	return err
 }
 
-func decompressZip(src string, output io.WriteCloser) error {
+func decompressZip(prefix string, src string, output io.WriteCloser) error {
 	zipReader, err := zip.OpenReader(src)
 	if err != nil {
 		return err
@@ -318,7 +353,16 @@ func decompressZip(src string, output io.WriteCloser) error {
 			logrus.Error(err)
 		}
 	}()
-	_, err = io.Copy(output, f)
+	size := int64(zipReader.File[0].CompressedSize64)
+	p, bar := progressBar(prefix, size, prefix+": done")
+	proxyReader := bar.ProxyReader(f)
+	defer func() {
+		if err := proxyReader.Close(); err != nil {
+			logrus.Error(err)
+		}
+	}()
+	_, err = io.Copy(output, proxyReader)
+	p.Wait()
 	return err
 }