Add extraction progress output to GetFSFromImage/GetFSFromLayers (#11)

Author: mafredri

pkg/util/fs_util.go

@@ -96,8 +96,9 @@ type FileContext struct {
 type ExtractFunction func(string, *tar.Header, string, io.Reader) error
 
 type FSConfig struct {
-	includeWhiteout bool
-	extractFunc     ExtractFunction
+	includeWhiteout         bool
+	printExtractionProgress bool
+	extractFunc             ExtractFunction
 }
 
 type FSOpt func(*FSConfig)
@@ -126,6 +127,12 @@ func IncludeWhiteout() FSOpt {
 	}
 }
 
+func PrintExtractionProgress() FSOpt {
+	return func(opts *FSConfig) {
+		opts.printExtractionProgress = true
+	}
+}
+
 func ExtractFunc(extractFunc ExtractFunction) FSOpt {
 	return func(opts *FSConfig) {
 		opts.extractFunc = extractFunc
@@ -144,7 +151,7 @@ func GetFSFromImage(root string, img v1.Image, extract ExtractFunction) ([]strin
 		return nil, err
 	}
 
-	return GetFSFromLayers(root, layers, ExtractFunc(extract))
+	return GetFSFromLayers(root, layers, ExtractFunc(extract), PrintExtractionProgress())
 }
 
 func GetFSFromLayers(root string, layers []v1.Layer, opts ...FSOpt) ([]string, error) {
@@ -163,12 +170,37 @@ func GetFSFromLayers(root string, layers []v1.Layer, opts ...FSOpt) ([]string, e
 		return nil, errors.New("must supply an extract function")
 	}
 
+	var totalSize int64
+	layerSizes := make([]int64, 0, len(layers))
+	for i, l := range layers {
+		layerSize, err := l.Size()
+		if err != nil {
+			return nil, errors.Wrap(err, fmt.Sprintf("error checking layer size %d", i))
+		}
+		layerSizes = append(layerSizes, layerSize)
+		totalSize += layerSize
+	}
+	printExtractionProgress := cfg.printExtractionProgress
+	if totalSize == 0 {
+		printExtractionProgress = false
+	}
+
+	if printExtractionProgress {
+		logrus.Infof("Extracting image layers to %s", root)
+	}
+
 	extractedFiles := []string{}
+	var extractedBytes int64
 	for i, l := range layers {
 		if mediaType, err := l.MediaType(); err == nil {
-			logrus.Tracef("Extracting layer %d of media type %s", i, mediaType)
+			logrus.Tracef("Extracting layer %d/%d of media type %s", i+1, len(layers), mediaType)
 		} else {
-			logrus.Tracef("Extracting layer %d", i)
+			logrus.Tracef("Extracting layer %d/%d", i+1, len(layers))
+		}
+
+		progressPerc := float64(extractedBytes) / float64(totalSize) * 100
+		if printExtractionProgress {
+			logrus.Infof("Extracting layer %d/%d (%.1f%%)", i+1, len(layers), progressPerc)
 		}
 
 		r, err := l.Uncompressed()
@@ -177,6 +209,16 @@ func GetFSFromLayers(root string, layers []v1.Layer, opts ...FSOpt) ([]string, e
 		}
 		defer r.Close()
 
+		if printExtractionProgress {
+			r = &printAfterReader{
+				ReadCloser: r,
+				after:      time.Second,
+				print: func(n int) {
+					logrus.Infof("Extracting layer %d/%d (%.1f%%) %s", i+1, len(layers), progressPerc, strings.Repeat(".", n))
+				},
+			}
+		}
+
 		tr := tar.NewReader(r)
 		for {
 			hdr, err := tr.Next()
@@ -225,10 +267,38 @@ func GetFSFromLayers(root string, layers []v1.Layer, opts ...FSOpt) ([]string, e
 
 			extractedFiles = append(extractedFiles, filepath.Join(root, cleanedName))
 		}
+
+		extractedBytes += layerSizes[i]
 	}
+
+	if printExtractionProgress {
+		logrus.Infof("Extraction complete")
+	}
+
 	return extractedFiles, nil
 }
 
+type printAfterReader struct {
+	io.ReadCloser
+	t     time.Time
+	after time.Duration
+	count int
+	print func(int)
+}
+
+func (r *printAfterReader) Read(p []byte) (n int, err error) {
+	n, err = r.ReadCloser.Read(p)
+	if r.t.IsZero() {
+		r.t = time.Now()
+	}
+	if time.Since(r.t) >= r.after {
+		r.count++
+		r.print(r.count)
+		r.t = time.Now()
+	}
+	return
+}
+
 // DeleteFilesystem deletes the extracted image file system
 func DeleteFilesystem() error {
 	logrus.Info("Deleting filesystem...")

pkg/util/fs_util_test.go

@@ -1067,6 +1067,7 @@ func Test_GetFSFromLayers_with_whiteouts_include_whiteout_enabled(t *testing.T)
 	f(expectedFiles, tw)
 
 	mockLayer := mockv1.NewMockLayer(ctrl)
+	mockLayer.EXPECT().Size().Return(int64(buf.Len()), nil)
 	mockLayer.EXPECT().MediaType().Return(types.OCILayer, nil)
 
 	rc := io.NopCloser(buf)
@@ -1082,6 +1083,7 @@ func Test_GetFSFromLayers_with_whiteouts_include_whiteout_enabled(t *testing.T)
 	f(secondLayerFiles, tw)
 
 	mockLayer2 := mockv1.NewMockLayer(ctrl)
+	mockLayer2.EXPECT().Size().Return(int64(buf.Len()), nil)
 	mockLayer2.EXPECT().MediaType().Return(types.OCILayer, nil)
 
 	rc = io.NopCloser(buf)
@@ -1175,6 +1177,7 @@ func Test_GetFSFromLayers_with_whiteouts_include_whiteout_disabled(t *testing.T)
 	f(expectedFiles, tw)
 
 	mockLayer := mockv1.NewMockLayer(ctrl)
+	mockLayer.EXPECT().Size().Return(int64(buf.Len()), nil)
 	mockLayer.EXPECT().MediaType().Return(types.OCILayer, nil)
 	layerFiles := []string{
 		filepath.Join(root, "foobar"),
@@ -1197,6 +1200,7 @@ func Test_GetFSFromLayers_with_whiteouts_include_whiteout_disabled(t *testing.T)
 	f(secondLayerFiles, tw)
 
 	mockLayer2 := mockv1.NewMockLayer(ctrl)
+	mockLayer2.EXPECT().Size().Return(int64(buf.Len()), nil)
 	mockLayer2.EXPECT().MediaType().Return(types.OCILayer, nil)
 
 	rc = io.NopCloser(buf)
@@ -1280,6 +1284,7 @@ func Test_GetFSFromLayers_ignorelist(t *testing.T) {
 	f(expectedFiles, tw)
 
 	mockLayer := mockv1.NewMockLayer(ctrl)
+	mockLayer.EXPECT().Size().Return(int64(buf.Len()), nil)
 	mockLayer.EXPECT().MediaType().Return(types.OCILayer, nil)
 	layerFiles := []string{
 		filepath.Join(root, ".wh.testdir"),
@@ -1345,6 +1350,8 @@ func Test_GetFSFromLayers_ignorelist(t *testing.T) {
 
 	f(layerFiles, tw)
 
+	mockLayer.EXPECT().Size().Return(int64(buf.Len()), nil)
+
 	rc = io.NopCloser(buf)
 	mockLayer.EXPECT().Uncompressed().Return(rc, nil)
 
@@ -1410,6 +1417,7 @@ func Test_GetFSFromLayers(t *testing.T) {
 	}
 
 	mockLayer := mockv1.NewMockLayer(ctrl)
+	mockLayer.EXPECT().Size().Return(int64(buf.Len()), nil)
 	mockLayer.EXPECT().MediaType().Return(types.OCILayer, nil)
 
 	rc := io.NopCloser(buf)