mutate: add uncompressed blob size annotation

Sometimes it is useful to have a reliable estimate of the space that
would be taken up by a layer when it is unpacked. Since gzip's headers
are only accurate up to 4GiB, larger layers can not be easily measured
without uncompressing them.

This commit adds a field to the gzip and zstd compressors to store
the bytes read from the stream during compression and uses that to set
an annotation on newly generated layers.

If the noop compressor is used, the annotation is not added, as the
existing "compressed" layer size would be the same.

Signed-off-by: Michael McCracken <[email protected]>

mutate/compress.go

@@ -22,6 +22,10 @@ type Compressor interface {
 	// indicate what compression type is used, e.g. "gzip", or "" for no
 	// compression.
 	MediaTypeSuffix() string
+
+	// BytesRead returns the number of bytes read from the uncompressed input
+	// stream at the current time, no guarantee of completion.
+	BytesRead() int64
 }
 
 type noopCompressor struct{}
@@ -34,28 +38,36 @@ func (nc noopCompressor) MediaTypeSuffix() string {
 	return ""
 }
 
+func (nc noopCompressor) BytesRead() int64 {
+	return -1
+}
+
 // NoopCompressor provides no compression.
 var NoopCompressor Compressor = noopCompressor{}
 
 // GzipCompressor provides gzip compression.
-var GzipCompressor Compressor = gzipCompressor{}
+var GzipCompressor Compressor = &gzipCompressor{}
 
-type gzipCompressor struct{}
+type gzipCompressor struct {
+	bytesRead int64
+}
 
-func (gz gzipCompressor) Compress(reader io.Reader) (io.ReadCloser, error) {
+func (gz *gzipCompressor) Compress(reader io.Reader) (io.ReadCloser, error) {
 	pipeReader, pipeWriter := io.Pipe()
 
 	gzw := gzip.NewWriter(pipeWriter)
 	if err := gzw.SetConcurrency(256<<10, 2*runtime.NumCPU()); err != nil {
 		return nil, errors.Wrapf(err, "set concurrency level to %v blocks", 2*runtime.NumCPU())
 	}
 	go func() {
-		if _, err := system.Copy(gzw, reader); err != nil {
+		bytesRead, err := system.Copy(gzw, reader)
+		if err != nil {
 			log.Warnf("gzip compress: could not compress layer: %v", err)
 			// #nosec G104
 			_ = pipeWriter.CloseWithError(errors.Wrap(err, "compressing layer"))
 			return
 		}
+		gz.bytesRead = bytesRead
 		if err := gzw.Close(); err != nil {
 			log.Warnf("gzip compress: could not close gzip writer: %v", err)
 			// #nosec G104
@@ -76,25 +88,33 @@ func (gz gzipCompressor) MediaTypeSuffix() string {
 	return "gzip"
 }
 
+func (gz gzipCompressor) BytesRead() int64 {
+	return gz.bytesRead
+}
+
 // ZstdCompressor provides zstd compression.
-var ZstdCompressor Compressor = zstdCompressor{}
+var ZstdCompressor Compressor = &zstdCompressor{}
 
-type zstdCompressor struct{}
+type zstdCompressor struct {
+	bytesRead int64
+}
 
-func (zs zstdCompressor) Compress(reader io.Reader) (io.ReadCloser, error) {
+func (zs *zstdCompressor) Compress(reader io.Reader) (io.ReadCloser, error) {
 
 	pipeReader, pipeWriter := io.Pipe()
 	zw, err := zstd.NewWriter(pipeWriter)
 	if err != nil {
 		return nil, err
 	}
 	go func() {
-		if _, err := system.Copy(zw, reader); err != nil {
+		bytesRead, err := system.Copy(zw, reader)
+		if err != nil {
 			log.Warnf("zstd compress: could not compress layer: %v", err)
 			// #nosec G104
 			_ = pipeWriter.CloseWithError(errors.Wrap(err, "compressing layer"))
 			return
 		}
+		zs.bytesRead = bytesRead
 		if err := zw.Close(); err != nil {
 			log.Warnf("zstd compress: could not close gzip writer: %v", err)
 			// #nosec G104
@@ -114,3 +134,7 @@ func (zs zstdCompressor) Compress(reader io.Reader) (io.ReadCloser, error) {
 func (zs zstdCompressor) MediaTypeSuffix() string {
 	return "zstd"
 }
+
+func (zs zstdCompressor) BytesRead() int64 {
+	return zs.bytesRead
+}

mutate/mutate.go

@@ -24,6 +24,7 @@ package mutate
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"reflect"
 	"time"
@@ -36,6 +37,8 @@ import (
 	"github.com/pkg/errors"
 )
 
+const UmociUncompressedBlobSizeAnnotation = "ci.umo.uncompressed_blob_size"
+
 func configPtr(c ispec.Image) *ispec.Image         { return &c }
 func manifestPtr(m ispec.Manifest) *ispec.Manifest { return &m }
 func timePtr(t time.Time) *time.Time               { return &t }
@@ -295,6 +298,13 @@ func (m *Mutator) Add(ctx context.Context, mediaType string, r io.Reader, histor
 		compressedMediaType = compressedMediaType + "+" + compressor.MediaTypeSuffix()
 	}
 
+	if annotations == nil {
+		annotations = make(map[string]string)
+	}
+	if compressor.BytesRead() >= 0 {
+		annotations[UmociUncompressedBlobSizeAnnotation] = fmt.Sprintf("%d", compressor.BytesRead())
+	}
+
 	// Append to layers.
 	desc = ispec.Descriptor{
 		MediaType:   compressedMediaType,

mutate/mutate_test.go

@@ -214,6 +214,7 @@ func TestMutateAdd(t *testing.T) {
 
 	// This isn't a valid image, but whatever.
 	buffer := bytes.NewBufferString("contents")
+	bufferSize := buffer.Len()
 
 	// Add a new layer.
 	annotations := map[string]string{"hello": "world"}
@@ -253,10 +254,18 @@ func TestMutateAdd(t *testing.T) {
 	if mutator.manifest.Layers[1].Digest == expectedLayerDigest {
 		t.Errorf("manifest.Layers[1].Digest is not the same!")
 	}
-	if len(mutator.manifest.Layers[1].Annotations) != 1 || mutator.manifest.Layers[1].Annotations["hello"] != "world" {
-		t.Errorf("manifest.Layers[1].Annotations was not set correctly!")
+	if len(mutator.manifest.Layers[1].Annotations) != 2 {
+		t.Errorf("manifest.Layers[1].Annotations was not set correctly!: %+v", mutator.manifest.Layers[1].Annotations)
+	}
+	if mutator.manifest.Layers[1].Annotations["hello"] != "world" {
+		t.Errorf("manifest.Layers[1].Annotations['hello'] was not set correctly!: %+v", mutator.manifest.Layers[1].Annotations)
+	}
+	if mutator.manifest.Layers[1].Annotations[UmociUncompressedBlobSizeAnnotation] != fmt.Sprintf("%d", bufferSize) {
+		t.Errorf("manifest.Layers[1].Annotations['%s'] was not set correctly!: %q, expected %d",
+			UmociUncompressedBlobSizeAnnotation,
+			mutator.manifest.Layers[1].Annotations[UmociUncompressedBlobSizeAnnotation],
+			bufferSize)
 	}
-
 	if mutator.manifest.Layers[1].Digest != newLayerDesc.Digest {
 		t.Fatalf("unexpected digest for new layer: %v %v", mutator.manifest.Layers[1].Digest, newLayerDesc.Digest)
 	}