Merge pull request #19249 from afbjorklund/decompress-progress

Add progress bar for decompress image

pkg/machine/pull.go

@@ -136,6 +136,28 @@ func DownloadImage(d DistributionDownload) error {
 	return Decompress(d.Get().LocalPath, d.Get().LocalUncompressedFile)
 }
 
+func progressBar(prefix string, size int64, onComplete string) (*mpb.Progress, *mpb.Bar) {
+	p := mpb.New(
+		mpb.WithWidth(80), // Do not go below 80, see bug #17718
+		mpb.WithRefreshRate(180*time.Millisecond),
+	)
+
+	bar := p.AddBar(size,
+		mpb.BarFillerClearOnComplete(),
+		mpb.PrependDecorators(
+			decor.OnComplete(decor.Name(prefix), onComplete),
+		),
+		mpb.AppendDecorators(
+			decor.OnComplete(decor.CountersKibiByte("%.1f / %.1f"), ""),
+		),
+	)
+	if size == 0 {
+		bar.SetTotal(0, true)
+	}
+
+	return p, bar
+}
+
 // DownloadVMImage downloads a VM image from url to given path
 // with download status
 func DownloadVMImage(downloadURL *url2.URL, imageName string, localImagePath string) error {
@@ -166,20 +188,7 @@ func DownloadVMImage(downloadURL *url2.URL, imageName string, localImagePath str
 	prefix := "Downloading VM image: " + imageName
 	onComplete := prefix + ": done"
 
-	p := mpb.New(
+	p, bar := progressBar(prefix, size, onComplete)
-		mpb.WithWidth(80), // Do not go below 80, see bug #17718
-		mpb.WithRefreshRate(180*time.Millisecond),
-	)
-
-	bar := p.AddBar(size,
-		mpb.BarFillerClearOnComplete(),
-		mpb.PrependDecorators(
-			decor.OnComplete(decor.Name(prefix), onComplete),
-		),
-		mpb.AppendDecorators(
-			decor.OnComplete(decor.CountersKibiByte("%.1f / %.1f"), ""),
-		),
-	)
 
 	proxyReader := bar.ProxyReader(resp.Body)
 	defer func() {
@@ -209,43 +218,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
@@ -262,18 +288,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
 	}
@@ -287,10 +330,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
@@ -312,7 +356,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
 }