refactor cache.Warm and add tests

pkg/cache/cache.go

@@ -35,40 +35,6 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-type NotFoundErr struct {
-	msg string
-}
-
-func (e NotFoundErr) Error() string {
-	return e.msg
-}
-
-type ExpiredErr struct {
-	msg string
-}
-
-func (e ExpiredErr) Error() string {
-	return e.msg
-}
-
-func IsNotFound(e error) bool {
-	switch e.(type) {
-	case NotFoundErr:
-		return true
-	}
-
-	return false
-}
-
-func IsExpired(e error) bool {
-	switch e.(type) {
-	case ExpiredErr:
-		return true
-	}
-
-	return false
-}
-
 // LayerCache is the layer cache
 type LayerCache interface {
 	RetrieveLayer(string) (v1.Image, error)

pkg/cache/doc_test.go

@@ -0,0 +1,47 @@
+/*
+Copyright 2019 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package cache
+
+import (
+	"bytes"
+	"log"
+
+	"github.com/GoogleContainerTools/kaniko/pkg/config"
+	"github.com/google/go-containerregistry/pkg/v1/remote"
+)
+
+func ExampleWarmer_Warm() {
+	tarBuf := new(bytes.Buffer)
+	manifestBuf := new(bytes.Buffer)
+	w := &Warmer{
+		Remote:         remote.Image,
+		Local:          LocalSource,
+		TarWriter:      tarBuf,
+		ManifestWriter: manifestBuf,
+	}
+
+	options := &config.WarmerOptions{}
+
+	digest, err := w.Warm("ubuntu:latest", options)
+	if err != nil {
+		if !IsAlreadyCached(err) {
+			log.Fatal(err)
+		}
+	}
+
+	log.Printf("digest %v tar len %d\nmanifest:\n%s\n", digest, tarBuf.Len(), manifestBuf.String())
+}

pkg/cache/errors.go

@@ -0,0 +1,79 @@
+/*
+Copyright 2019 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package cache
+
+// IsAlreadyCached returns true if the supplied error is of the type AlreadyCachedErr
+// otherwise it returns false.
+func IsAlreadyCached(err error) bool {
+	switch err.(type) {
+	case AlreadyCachedErr:
+		return true
+	}
+
+	return false
+}
+
+// AlreadyCachedErr is returned when the Docker image requested for caching is already
+// present in the cache.
+type AlreadyCachedErr struct {
+	msg string
+}
+
+func (a AlreadyCachedErr) Error() string {
+	return a.msg
+}
+
+// IsNotFound returns true if the supplied error is of the type NotFoundErr
+// otherwise it returns false.
+func IsNotFound(e error) bool {
+	switch e.(type) {
+	case NotFoundErr:
+		return true
+	}
+
+	return false
+}
+
+// NotFoundErr is returned when the requested Docker image is not present in the cache.
+type NotFoundErr struct {
+	msg string
+}
+
+func (e NotFoundErr) Error() string {
+	return e.msg
+}
+
+// IsExpired returns true if the supplied error is of the type ExpiredErr
+// otherwise it returns false.
+func IsExpired(e error) bool {
+	switch e.(type) {
+	case ExpiredErr:
+		return true
+	}
+
+	return false
+}
+
+// ExpiredErr is returned when the requested Docker image is present in the cache, but is
+// expired according to the supplied TTL.
+type ExpiredErr struct {
+	msg string
+}
+
+func (e ExpiredErr) Error() string {
+	return e.msg
+}

pkg/cache/warm.go

@@ -17,12 +17,15 @@ limitations under the License.
 package cache
 
 import (
-	"fmt"
+	"bytes"
+	"io"
 	"io/ioutil"
+	"os"
 	"path"
 
 	"github.com/GoogleContainerTools/kaniko/pkg/config"
 	"github.com/google/go-containerregistry/pkg/name"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
 	"github.com/google/go-containerregistry/pkg/v1/tarball"
 	"github.com/pkg/errors"
@@ -37,42 +40,111 @@ func WarmCache(opts *config.WarmerOptions) error {
 	logrus.Debugf("%s\n", images)
 
 	for _, image := range images {
-		cacheRef, err := name.NewTag(image, name.WeakValidation)
-		if err != nil {
-			return errors.Wrap(err, fmt.Sprintf("Failed to verify image name: %s", image))
-		}
-		img, err := remote.Image(cacheRef)
-		if err != nil || img == nil {
-			return errors.Wrap(err, fmt.Sprintf("Failed to retrieve image: %s", image))
-		}
+		tarBuf := new(bytes.Buffer)
+		manifestBuf := new(bytes.Buffer)
 
-		digest, err := img.Digest()
-		if err != nil {
-			return errors.Wrap(err, fmt.Sprintf("Failed to retrieve digest: %s", image))
+		cw := &Warmer{
+			Remote:         remote.Image,
+			Local:          LocalSource,
+			TarWriter:      tarBuf,
+			ManifestWriter: manifestBuf,
 		}
-		cachePath := path.Join(cacheDir, digest.String())
 
-		if !opts.Force {
-			_, err := LocalSource(&opts.CacheOptions, digest.String())
-			if err == nil || IsExpired(err) {
-				continue
+		digest, err := cw.Warm(image, opts)
+		if err != nil {
+			if !IsAlreadyCached(err) {
+				return err
 			}
-		}
 
-		err = tarball.WriteToFile(cachePath, cacheRef, img)
-		if err != nil {
-			return errors.Wrap(err, fmt.Sprintf("Failed to write %s to cache", image))
+			continue
 		}
 
-		mfst, err := img.RawManifest()
-		if err != nil {
-			return errors.Wrap(err, fmt.Sprintf("Failed to retrieve manifest for %s", image))
-		}
-		mfstPath := cachePath + ".json"
-		if err := ioutil.WriteFile(mfstPath, mfst, 0666); err != nil {
-			return errors.Wrap(err, fmt.Sprintf("Failed to save manifest for %s", image))
+		cachePath := path.Join(cacheDir, digest.String())
+
+		if err := writeBufsToFile(cachePath, tarBuf, manifestBuf); err != nil {
+			return err
 		}
+
 		logrus.Debugf("Wrote %s to cache", image)
 	}
 	return nil
 }
+
+func writeBufsToFile(cachePath string, tarBuf, manifestBuf *bytes.Buffer) error {
+	f, err := os.Create(cachePath)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	if _, err := f.Write(tarBuf.Bytes()); err != nil {
+		return errors.Wrap(err, "Failed to save tar to file")
+	}
+
+	mfstPath := cachePath + ".json"
+	if err := ioutil.WriteFile(mfstPath, manifestBuf.Bytes(), 0666); err != nil {
+		return errors.Wrap(err, "Failed to save manifest to file")
+	}
+
+	return nil
+}
+
+// FetchRemoteImage retrieves a Docker image manifest from a remote source.
+// github.com/google/go-containerregistry/pkg/v1/remote.Image can be used as
+// this type.
+type FetchRemoteImage func(name.Reference, ...remote.Option) (v1.Image, error)
+
+// FetchLocalSource retrieves a Docker image manifest from a local source.
+// github.com/GoogleContainerTools/kaniko/cache.LocalSource can be used as
+// this type.
+type FetchLocalSource func(*config.CacheOptions, string) (v1.Image, error)
+
+// Warmer is used to prepopulate the cache with a Docker image
+type Warmer struct {
+	Remote         FetchRemoteImage
+	Local          FetchLocalSource
+	TarWriter      io.Writer
+	ManifestWriter io.Writer
+}
+
+// Warm retrieves a Docker image and populates the supplied buffer with the image content and manifest
+// or returns an AlreadyCachedErr if the image is present in the cache.
+func (w *Warmer) Warm(image string, opts *config.WarmerOptions) (v1.Hash, error) {
+	cacheRef, err := name.NewTag(image, name.WeakValidation)
+	if err != nil {
+		return v1.Hash{}, errors.Wrapf(err, "Failed to verify image name: %s", image)
+	}
+
+	img, err := w.Remote(cacheRef)
+	if err != nil || img == nil {
+		return v1.Hash{}, errors.Wrapf(err, "Failed to retrieve image: %s", image)
+	}
+
+	digest, err := img.Digest()
+	if err != nil {
+		return v1.Hash{}, errors.Wrapf(err, "Failed to retrieve digest: %s", image)
+	}
+
+	if !opts.Force {
+		_, err := w.Local(&opts.CacheOptions, digest.String())
+		if err == nil || IsExpired(err) {
+			return v1.Hash{}, AlreadyCachedErr{}
+		}
+	}
+
+	err = tarball.Write(cacheRef, img, w.TarWriter)
+	if err != nil {
+		return v1.Hash{}, errors.Wrapf(err, "Failed to write %s to tar buffer", image)
+	}
+
+	mfst, err := img.RawManifest()
+	if err != nil {
+		return v1.Hash{}, errors.Wrapf(err, "Failed to retrieve manifest for %s", image)
+	}
+
+	if _, err := w.ManifestWriter.Write(mfst); err != nil {
+		return v1.Hash{}, errors.Wrapf(err, "Failed to save manifest to buffer for %s", image)
+	}
+
+	return digest, nil
+}