Add compression option to tarball.Layer and stream.Layer

Signed-off-by: Javier Romero <jromero@pivotal.io>
google · Oct 17, 2019 · f2b0120 · f2b0120
1 parent ad4d343
commit f2b0120
Show file tree

Hide file tree

Showing 6 changed files with 108 additions and 28 deletions.
diff --git a/go.mod b/go.mod
diff --git a/go.sum b/go.sum
diff --git a/pkg/v1/stream/layer.go b/pkg/v1/stream/layer.go
@@ -39,8 +39,9 @@ var (
 
 // Layer is a streaming implementation of v1.Layer.
 type Layer struct {
-	blob     io.ReadCloser
-	consumed bool
+	blob        io.ReadCloser
+	consumed    bool
+	compression int
 
 	mu             sync.Mutex
 	digest, diffID *v1.Hash
@@ -49,8 +50,29 @@ type Layer struct {
 
 var _ v1.Layer = (*Layer)(nil)
 
+// LayerOption applies options to layer
+type LayerOption func(*Layer)
+
+// WithCompressionLevel sets the gzip compression. See `gzip.NewWriterLevel` for possible values.
+func WithCompressionLevel(level int) LayerOption {
+	return func(l *Layer) {
+		l.compression = level
+	}
+}
+
 // NewLayer creates a Layer from an io.ReadCloser.
-func NewLayer(rc io.ReadCloser) *Layer { return &Layer{blob: rc} }
+func NewLayer(rc io.ReadCloser, opts ...LayerOption) *Layer {
+	layer := &Layer{
+		blob:        rc,
+		compression: gzip.BestSpeed,
+	}
+
+	for _, opt := range opts {
+		opt(layer)
+	}
+
+	return layer
+}
 
 // Digest implements v1.Layer.
 func (l *Layer) Digest() (v1.Hash, error) {
@@ -121,7 +143,7 @@ func newCompressedReader(l *Layer) (*compressedReader, error) {
 	// capture compressed digest, and a countWriter to capture compressed
 	// size.
 	pr, pw := io.Pipe()
-	zw, err := gzip.NewWriterLevel(io.MultiWriter(pw, zh, count), gzip.BestSpeed)
+	zw, err := gzip.NewWriterLevel(io.MultiWriter(pw, zh, count), l.compression)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/v1/stream/layer_test.go b/pkg/v1/stream/layer_test.go
@@ -84,6 +84,25 @@ func TestStreamVsBuffer(t *testing.T) {
 	} else if s != wantSize {
 		t.Errorf("stream Size got %d, want %d", s, wantSize)
 	}
+
+	// Test with different compression
+	l2 := NewLayer(newBlob(), WithCompressionLevel(2))
+	l2WantDigest := "sha256:c9afe7b0da6783232e463e12328cb306142548384accf3995806229c9a6a707f"
+	if c, err := l2.Compressed(); err != nil {
+		t.Errorf("Compressed: %v", err)
+	} else {
+		if _, err := io.Copy(ioutil.Discard, c); err != nil {
+			t.Errorf("error reading Compressed: %v", err)
+		}
+		if err := c.Close(); err != nil {
+			t.Errorf("Close: %v", err)
+		}
+	}
+	if d, err := l2.Digest(); err != nil {
+		t.Errorf("Digest: %v", err)
+	} else if d.String() != l2WantDigest {
+		t.Errorf("stream Digest got %q, want %q", d.String(), l2WantDigest)
+	}
 }
 
 func TestLargeStream(t *testing.T) {

diff --git a/pkg/v1/tarball/layer.go b/pkg/v1/tarball/layer.go
@@ -27,11 +27,12 @@ import (
 )
 
 type layer struct {
-	digest     v1.Hash
-	diffID     v1.Hash
-	size       int64
-	opener     Opener
-	compressed bool
+	digest      v1.Hash
+	diffID      v1.Hash
+	size        int64
+	opener      Opener
+	compressed  bool
+	compression int
 }
 
 func (l *layer) Digest() (v1.Hash, error) {
@@ -45,7 +46,7 @@ func (l *layer) DiffID() (v1.Hash, error) {
 func (l *layer) Compressed() (io.ReadCloser, error) {
 	rc, err := l.opener()
 	if err == nil && !l.compressed {
-		return v1util.GzipReadCloser(rc), nil
+		return v1util.GzipReadCloserLevel(rc, l.compression), nil
 	}
 
 	return rc, err
@@ -68,16 +69,26 @@ func (l *layer) MediaType() (types.MediaType, error) {
 	return types.DockerLayer, nil
 }
 
+// LayerOption applies options to layer
+type LayerOption func(*layer)
+
+// WithCompressionLevel sets the gzip compression. See `gzip.NewWriterLevel` for possible values.
+func WithCompressionLevel(level int) LayerOption {
+	return func(l *layer) {
+		l.compression = level
+	}
+}
+
 // LayerFromFile returns a v1.Layer given a tarball
-func LayerFromFile(path string) (v1.Layer, error) {
+func LayerFromFile(path string, opts ...LayerOption) (v1.Layer, error) {
 	opener := func() (io.ReadCloser, error) {
 		return os.Open(path)
 	}
-	return LayerFromOpener(opener)
+	return LayerFromOpener(opener, opts...)
 }
 
 // LayerFromOpener returns a v1.Layer given an Opener function
-func LayerFromOpener(opener Opener) (v1.Layer, error) {
+func LayerFromOpener(opener Opener, opts ...LayerOption) (v1.Layer, error) {
 	rc, err := opener()
 	if err != nil {
 		return nil, err
@@ -89,24 +100,25 @@ func LayerFromOpener(opener Opener) (v1.Layer, error) {
 		return nil, err
 	}
 
-	var digest v1.Hash
-	var size int64
-	if digest, size, err = computeDigest(opener, compressed); err != nil {
+	layer := &layer{
+		compressed:  compressed,
+		compression: gzip.BestSpeed,
+		opener:      opener,
+	}
+
+	for _, opt := range opts {
+		opt(layer)
+	}
+
+	if layer.digest, layer.size, err = computeDigest(opener, compressed, layer.compression); err != nil {
 		return nil, err
 	}
 
-	diffID, err := computeDiffID(opener, compressed)
-	if err != nil {
+	if layer.diffID, err = computeDiffID(opener, compressed); err != nil {
 		return nil, err
 	}
 
-	return &layer{
-		digest:     digest,
-		diffID:     diffID,
-		size:       size,
-		compressed: compressed,
-		opener:     opener,
-	}, nil
+	return layer, nil
 }
 
 // LayerFromReader returns a v1.Layer given a io.Reader.
@@ -121,7 +133,7 @@ func LayerFromReader(reader io.Reader) (v1.Layer, error) {
 	})
 }
 
-func computeDigest(opener Opener, compressed bool) (v1.Hash, int64, error) {
+func computeDigest(opener Opener, compressed bool, compression int) (v1.Hash, int64, error) {
 	rc, err := opener()
 	if err != nil {
 		return v1.Hash{}, 0, err
@@ -132,7 +144,7 @@ func computeDigest(opener Opener, compressed bool) (v1.Hash, int64, error) {
 		return v1.SHA256(rc)
 	}
 
-	return v1.SHA256(v1util.GzipReadCloser(ioutil.NopCloser(rc)))
+	return v1.SHA256(v1util.GzipReadCloserLevel(ioutil.NopCloser(rc), compression))
 }
 
 func computeDiffID(opener Opener, compressed bool) (v1.Hash, error) {

diff --git a/pkg/v1/tarball/layer_test.go b/pkg/v1/tarball/layer_test.go
@@ -47,9 +47,34 @@ func TestLayerFromFile(t *testing.T) {
 	if err := validate.Layer(tarLayer); err != nil {
 		t.Errorf("validate.Layer(tarLayer): %v", err)
 	}
+
 	if err := validate.Layer(tarGzLayer); err != nil {
 		t.Errorf("validate.Layer(tarGzLayer): %v", err)
 	}
+
+	tarLayerDefaultCompression, err := LayerFromFile("testdata/content.tar", WithCompressionLevel(gzip.DefaultCompression))
+	if err != nil {
+		t.Fatalf("Unable to create layer with 'Default' compression from tar file: %v", err)
+	}
+
+	defaultDigest, err := tarLayerDefaultCompression.Digest()
+	if err != nil {
+		t.Fatal("Unable to generate digest with 'Default' compression", err)
+	}
+
+	tarLayerSpeedCompression, err := LayerFromFile("testdata/content.tar", WithCompressionLevel(gzip.BestSpeed))
+	if err != nil {
+		t.Fatalf("Unable to create layer with 'BestSpeed' compression from tar file: %v", err)
+	}
+
+	speedDigest, err := tarLayerSpeedCompression.Digest()
+	if err != nil {
+		t.Fatal("Unable to generate digest with 'BestSpeed' compression", err)
+	}
+
+	if defaultDigest.String() == speedDigest.String() {
+		t.Errorf("expected digests to differ: %s", defaultDigest.String())
+	}
 }
 
 func TestLayerFromOpenerReader(t *testing.T) {