Skip to content

Commit

Permalink
Add extraction progress output to GetFSFromImage/GetFSFromLayers (#11)
Browse files Browse the repository at this point in the history
  • Loading branch information
mafredri authored May 24, 2024
1 parent e4a3f0b commit 9d0d559
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 5 deletions.
80 changes: 75 additions & 5 deletions pkg/util/fs_util.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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) {
Expand All @@ -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()
Expand All @@ -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()
Expand Down Expand Up @@ -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...")
Expand Down
8 changes: 8 additions & 0 deletions pkg/util/fs_util_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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"),
Expand All @@ -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)
Expand Down Expand Up @@ -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"),
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 9d0d559

Please sign in to comment.