cache: lazily pull contents to populate cache (google#951)

Author: imjasonh

pkg/v1/cache/cache.go

@@ -3,9 +3,11 @@ package cache
 
 import (
 	"errors"
+	"io"
 
 	"github.com/google/go-containerregistry/pkg/logs"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/types"
 )
 
 // Cache encapsulates methods to interact with cached layers.
@@ -55,46 +57,66 @@ func (i *image) Layers() ([]v1.Layer, error) {
 
 	var out []v1.Layer
 	for _, l := range ls {
-		// Check if this layer is present in the cache in compressed
-		// form.
-		digest, err := l.Digest()
-		if err != nil {
-			return nil, err
-		}
-		if cl, err := i.c.Get(digest); err == nil {
-			// Layer found in the cache.
-			logs.Progress.Printf("Layer %s found (compressed) in cache", digest)
-			out = append(out, cl)
-			continue
-		} else if err != nil && err != ErrNotFound {
-			return nil, err
-		}
+		out = append(out, &lazyLayer{inner: l, c: i.c})
+	}
+	return out, nil
+}
 
-		// Check if this layer is present in the cache in
-		// uncompressed form.
-		diffID, err := l.DiffID()
-		if err != nil {
-			return nil, err
-		}
-		if cl, err := i.c.Get(diffID); err == nil {
-			// Layer found in the cache.
-			logs.Progress.Printf("Layer %s found (uncompressed) in cache", diffID)
-			out = append(out, cl)
-		} else if err != nil && err != ErrNotFound {
-			return nil, err
-		}
+type lazyLayer struct {
+	inner v1.Layer
+	c     Cache
+}
 
-		// Not cached, fall through to real layer.
-		l, err = i.c.Put(l)
-		if err != nil {
-			return nil, err
-		}
-		out = append(out, l)
+func (l *lazyLayer) Compressed() (io.ReadCloser, error) {
+	digest, err := l.inner.Digest()
+	if err != nil {
+		return nil, err
+	}
 
+	if cl, err := l.c.Get(digest); err == nil {
+		// Layer found in the cache.
+		logs.Progress.Printf("Layer %s found (compressed) in cache", digest)
+		return cl.Compressed()
+	} else if err != nil && err != ErrNotFound {
+		return nil, err
 	}
-	return out, nil
+
+	// Not cached, pull and return the real layer.
+	logs.Progress.Printf("Layer %s not found (compressed) in cache, getting", digest)
+	rl, err := l.c.Put(l.inner)
+	if err != nil {
+		return nil, err
+	}
+	return rl.Compressed()
 }
 
+func (l *lazyLayer) Uncompressed() (io.ReadCloser, error) {
+	diffID, err := l.inner.DiffID()
+	if err != nil {
+		return nil, err
+	}
+	if cl, err := l.c.Get(diffID); err == nil {
+		// Layer found in the cache.
+		logs.Progress.Printf("Layer %s found (uncompressed) in cache", diffID)
+		return cl.Uncompressed()
+	} else if err != nil && err != ErrNotFound {
+		return nil, err
+	}
+
+	// Not cached, pull and return the real layer.
+	logs.Progress.Printf("Layer %s not found (uncompressed) in cache, getting", diffID)
+	rl, err := l.c.Put(l.inner)
+	if err != nil {
+		return nil, err
+	}
+	return rl.Uncompressed()
+}
+
+func (l *lazyLayer) Size() (int64, error)                { return l.inner.Size() }
+func (l *lazyLayer) DiffID() (v1.Hash, error)            { return l.inner.DiffID() }
+func (l *lazyLayer) Digest() (v1.Hash, error)            { return l.inner.Digest() }
+func (l *lazyLayer) MediaType() (types.MediaType, error) { return l.inner.MediaType() }
+
 func (i *image) LayerByDigest(h v1.Hash) (v1.Layer, error) {
 	l, err := i.c.Get(h)
 	if err == ErrNotFound {

pkg/v1/cache/cache_test.go

@@ -2,51 +2,58 @@ package cache
 
 import (
 	"errors"
+	"io"
+	"io/ioutil"
 	"testing"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/random"
 	"github.com/google/go-containerregistry/pkg/v1/validate"
 )
 
-// TestCache tests that the cache is populated when LayerByDigest is called.
-func TestCache(t *testing.T) {
-	numLayers := 5
-	img, err := random.Image(10, int64(numLayers))
+func TestImage(t *testing.T) {
+	img, err := random.Image(1024, 5)
 	if err != nil {
 		t.Fatalf("random.Image: %v", err)
 	}
 	m := &memcache{map[v1.Hash]v1.Layer{}}
 	img = Image(img, m)
 
-	// Cache is empty.
-	if len(m.m) != 0 {
-		t.Errorf("Before consuming, cache is non-empty: %+v", m.m)
-	}
-
-	// Consume each layer, cache gets populated.
-	if _, err := img.Layers(); err != nil {
-		t.Fatalf("Layers: %v", err)
+	// Validate twice to hit the cache.
+	if err := validate.Image(img); err != nil {
+		t.Errorf("Validate: %v", err)
 	}
-	if got, want := len(m.m), numLayers; got != want {
-		t.Errorf("Cache has %d entries, want %d", got, want)
+	if err := validate.Image(img); err != nil {
+		t.Errorf("Validate: %v", err)
 	}
 }
 
-func TestImage(t *testing.T) {
+func TestLayersLazy(t *testing.T) {
 	img, err := random.Image(1024, 5)
 	if err != nil {
 		t.Fatalf("random.Image: %v", err)
 	}
 	m := &memcache{map[v1.Hash]v1.Layer{}}
 	img = Image(img, m)
 
-	// Validate twice to hit the cache.
-	if err := validate.Image(img); err != nil {
-		t.Errorf("Validate: %v", err)
+	layers, err := img.Layers()
+	if err != nil {
+		t.Fatalf("img.Layers: %v", err)
 	}
-	if err := validate.Image(img); err != nil {
-		t.Errorf("Validate: %v", err)
+
+	// After calling Layers, nothing is cached.
+	if got, want := len(m.m), 0; got != want {
+		t.Errorf("Cache has %d entries, want %d", got, want)
+	}
+
+	rc, err := layers[0].Uncompressed()
+	if err != nil {
+		t.Fatalf("layer.Uncompressed: %v", err)
+	}
+	io.Copy(ioutil.Discard, rc)
+
+	if got, expected := len(m.m), 1; got != expected {
+		t.Errorf("expected %v layers in cache after reading, got %v", expected, got)
 	}
 }