inline remote cache support

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

cache/remotecache/import.go

@@ -4,13 +4,19 @@ import (
 	"context"
 	"encoding/json"
 	"io"
+	"sync"
 
 	"github.com/containerd/containerd/content"
+	"github.com/containerd/containerd/images"
 	v1 "github.com/moby/buildkit/cache/remotecache/v1"
 	"github.com/moby/buildkit/solver"
+	"github.com/moby/buildkit/util/imageutil"
 	"github.com/moby/buildkit/worker"
+	digest "github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+	"golang.org/x/sync/errgroup"
 )
 
 // ResolveCacheImporterFunc returns importer and descriptor.
@@ -55,7 +61,7 @@ func (ci *contentCacheImporter) Resolve(ctx context.Context, desc ocispec.Descri
 	}
 
 	if configDesc.Digest == "" {
-		return nil, errors.Errorf("invalid build cache from %+v, lacks manifest with MediaType=%s", desc, v1.CacheConfigMediaTypeV0)
+		return ci.importInlineCache(ctx, dt, id, w)
 	}
 
 	dt, err = readBlob(ctx, ci.provider, configDesc)
@@ -95,3 +101,126 @@ func readBlob(ctx context.Context, provider content.Provider, desc ocispec.Descr
 	}
 	return dt, err
 }
+
+func (ci *contentCacheImporter) importInlineCache(ctx context.Context, dt []byte, id string, w worker.Worker) (solver.CacheManager, error) {
+	m := map[digest.Digest][]byte{}
+
+	if err := ci.allDistributionManifests(ctx, dt, m); err != nil {
+		return nil, err
+	}
+
+	var mu sync.Mutex
+	cc := v1.NewCacheChains()
+
+	eg, ctx := errgroup.WithContext(ctx)
+	for dgst, dt := range m {
+		func(dgst digest.Digest, dt []byte) {
+			eg.Go(func() error {
+				var m ocispec.Manifest
+
+				if err := json.Unmarshal(dt, &m); err != nil {
+					return err
+				}
+
+				if m.Config.Digest == "" || len(m.Layers) == 0 {
+					return nil
+				}
+
+				p, err := content.ReadBlob(ctx, ci.provider, m.Config)
+				if err != nil {
+					return err
+				}
+
+				var img struct {
+					Rootfs struct {
+						DiffIDs []digest.Digest `json:"diff_ids"`
+					} `json:"rootfs"`
+					Cache json.RawMessage `json:"moby.buildkit.cache.v0"`
+				}
+
+				if err := json.Unmarshal(p, &img); err != nil {
+					return err
+				}
+
+				if len(img.Rootfs.DiffIDs) != len(m.Layers) {
+					logrus.Warnf("invalid image with mismatching manifest and config")
+					return nil
+				}
+
+				if img.Cache == nil {
+					return nil
+				}
+
+				var config v1.CacheConfig
+				if err := json.Unmarshal(img.Cache, &config.Records); err != nil {
+					return err
+				}
+
+				layers := v1.DescriptorProvider{}
+				for i, m := range m.Layers {
+					if m.Annotations == nil {
+						m.Annotations = map[string]string{}
+					}
+					m.Annotations["containerd.io/uncompressed"] = img.Rootfs.DiffIDs[i].String()
+					layers[m.Digest] = v1.DescriptorProviderPair{
+						Descriptor: m,
+						Provider:   ci.provider,
+					}
+					config.Layers = append(config.Layers, v1.CacheLayer{
+						Blob:        m.Digest,
+						ParentIndex: i - 1,
+					})
+				}
+
+				mu.Lock()
+				if err := v1.ParseConfig(config, layers, cc); err != nil {
+					return err
+				}
+				mu.Unlock()
+				return nil
+			})
+		}(dgst, dt)
+	}
+
+	if err := eg.Wait(); err != nil {
+		return nil, err
+	}
+
+	keysStorage, resultStorage, err := v1.NewCacheKeyStorage(cc, w)
+	if err != nil {
+		return nil, err
+	}
+	return solver.NewCacheManager(id, keysStorage, resultStorage), nil
+}
+
+func (ci *contentCacheImporter) allDistributionManifests(ctx context.Context, dt []byte, m map[digest.Digest][]byte) error {
+	mt, err := imageutil.DetectManifestBlobMediaType(dt)
+	if err != nil {
+		return err
+	}
+
+	switch mt {
+	case images.MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest:
+		m[digest.FromBytes(dt)] = dt
+	case images.MediaTypeDockerSchema2ManifestList, ocispec.MediaTypeImageIndex:
+		var index ocispec.Index
+		if err := json.Unmarshal(dt, &index); err != nil {
+			return err
+		}
+
+		for _, d := range index.Manifests {
+			if _, ok := m[d.Digest]; ok {
+				continue
+			}
+			p, err := content.ReadBlob(ctx, ci.provider, d)
+			if err != nil {
+				return err
+			}
+			if err := ci.allDistributionManifests(ctx, p, m); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}

cache/remotecache/inline/inline.go

@@ -0,0 +1,70 @@
+package registry
+
+import (
+	"context"
+	"encoding/json"
+
+	"github.com/moby/buildkit/cache/remotecache"
+	v1 "github.com/moby/buildkit/cache/remotecache/v1"
+	"github.com/moby/buildkit/solver"
+	digest "github.com/opencontainers/go-digest"
+	"github.com/sirupsen/logrus"
+)
+
+func ResolveCacheExporterFunc() remotecache.ResolveCacheExporterFunc {
+	return func(ctx context.Context, _ map[string]string) (remotecache.Exporter, error) {
+		return NewExporter(), nil
+	}
+}
+
+func NewExporter() remotecache.Exporter {
+	cc := v1.NewCacheChains()
+	return &exporter{CacheExporterTarget: cc, chains: cc}
+}
+
+type exporter struct {
+	solver.CacheExporterTarget
+	chains *v1.CacheChains
+}
+
+func (ce *exporter) Finalize(ctx context.Context) (map[string]string, error) {
+	return nil, nil
+}
+
+func (ce *exporter) ExportForLayers(layers []digest.Digest) ([]byte, error) {
+	config, descs, err := ce.chains.Marshal()
+	if err != nil {
+		return nil, err
+	}
+
+	descs2 := map[digest.Digest]v1.DescriptorProviderPair{}
+	for _, k := range layers {
+		if v, ok := descs[k]; ok {
+			descs2[k] = v
+		}
+	}
+
+	cc := v1.NewCacheChains()
+	if err := v1.ParseConfig(*config, descs, cc); err != nil {
+		return nil, err
+	}
+
+	cfg, _, err := cc.Marshal()
+	if err != nil {
+		return nil, err
+	}
+
+	if len(cfg.Layers) == 0 {
+		logrus.Warn("failed to match any cache with layers")
+		return nil, nil
+	}
+
+	// TODO: are the layers always ordered already or should we check?
+
+	dt, err := json.Marshal(cfg.Records)
+	if err != nil {
+		return nil, err
+	}
+
+	return dt, nil
+}

cache/remotecache/v1/doc.go

@@ -6,7 +6,7 @@ package cacheimport
 // https://github.com/opencontainers/image-spec/blob/master/image-index.md .
 // Manifests array contains descriptors to the cache layers and one instance of
 // build cache config with media type application/vnd.buildkit.cacheconfig.v0 .
-// The cache layer descripts need to have an annotation with uncompressed digest
+// The cache layer descriptors need to have an annotation with uncompressed digest
 // to allow deduplication on extraction and optionally "buildkit/createdat"
 // annotation to support maintaining original timestamps.
 //

cache/remotecache/v1/parse.go

@@ -15,6 +15,10 @@ func Parse(configJSON []byte, provider DescriptorProvider, t solver.CacheExporte
 		return err
 	}
 
+	return ParseConfig(config, provider, t)
+}
+
+func ParseConfig(config CacheConfig, provider DescriptorProvider, t solver.CacheExporterTarget) error {
 	cache := map[int]solver.CacheExporterRecord{}
 
 	for i := range config.Records {
@@ -22,7 +26,6 @@ func Parse(configJSON []byte, provider DescriptorProvider, t solver.CacheExporte
 			return err
 		}
 	}
-
 	return nil
 }
 
@@ -57,7 +60,9 @@ func parseRecord(cc CacheConfig, idx int, provider DescriptorProvider, t solver.
 		if err != nil {
 			return nil, err
 		}
-		r.AddResult(res.CreatedAt, remote)
+		if remote != nil {
+			r.AddResult(res.CreatedAt, remote)
+		}
 	}
 
 	cache[idx] = r
@@ -78,7 +83,7 @@ func getRemoteChain(layers []CacheLayer, idx int, provider DescriptorProvider, v
 
 	descPair, ok := provider[l.Blob]
 	if !ok {
-		return nil, errors.Errorf("missing blob for %s", l.Blob)
+		return nil, nil
 	}
 
 	var r *solver.Remote
@@ -88,6 +93,9 @@ func getRemoteChain(layers []CacheLayer, idx int, provider DescriptorProvider, v
 		if err != nil {
 			return nil, err
 		}
+		if r == nil {
+			return nil, nil
+		}
 		r.Descriptors = append(r.Descriptors, descPair.Descriptor)
 		mp := contentutil.NewMultiProvider(r.Provider)
 		mp.Add(descPair.Descriptor.Digest, descPair.Provider)

cmd/buildkitd/main.go

@@ -22,6 +22,7 @@ import (
 	"github.com/docker/go-connections/sockets"
 	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
 	"github.com/moby/buildkit/cache/remotecache"
+	inlineremotecache "github.com/moby/buildkit/cache/remotecache/inline"
 	localremotecache "github.com/moby/buildkit/cache/remotecache/local"
 	registryremotecache "github.com/moby/buildkit/cache/remotecache/registry"
 	"github.com/moby/buildkit/client"
@@ -513,6 +514,7 @@ func newController(c *cli.Context, cfg *config.Config) (*control.Controller, err
 	remoteCacheExporterFuncs := map[string]remotecache.ResolveCacheExporterFunc{
 		"registry": registryremotecache.ResolveCacheExporterFunc(sessionManager, resolverFn),
 		"local":    localremotecache.ResolveCacheExporterFunc(sessionManager),
+		"inline":   inlineremotecache.ResolveCacheExporterFunc(),
 	}
 	remoteCacheImporterFuncs := map[string]remotecache.ResolveCacheImporterFunc{
 		"registry": registryremotecache.ResolveCacheImporterFunc(sessionManager, resolverFn),

exporter/containerimage/exptypes/types.go

@@ -3,6 +3,7 @@ package exptypes
 import specs "github.com/opencontainers/image-spec/specs-go/v1"
 
 const ExporterImageConfigKey = "containerimage.config"
+const ExporterInlineCache = "containerimage.inlinecache"
 const ExporterPlatformsKey = "refs.platforms"
 
 type Platforms struct {

exporter/containerimage/writer.go

@@ -56,7 +56,7 @@ func (ic *ImageWriter) Commit(ctx context.Context, inp exporter.Source, oci bool
 		if err != nil {
 			return nil, err
 		}
-		return ic.commitDistributionManifest(ctx, inp.Ref, inp.Metadata[exptypes.ExporterImageConfigKey], layers[0], oci)
+		return ic.commitDistributionManifest(ctx, inp.Ref, inp.Metadata[exptypes.ExporterImageConfigKey], layers[0], oci, inp.Metadata[exptypes.ExporterInlineCache])
 	}
 
 	var p exptypes.Platforms
@@ -108,7 +108,7 @@ func (ic *ImageWriter) Commit(ctx context.Context, inp exporter.Source, oci bool
 		}
 		config := inp.Metadata[fmt.Sprintf("%s/%s", exptypes.ExporterImageConfigKey, p.ID)]
 
-		desc, err := ic.commitDistributionManifest(ctx, r, config, layers[layersMap[p.ID]], oci)
+		desc, err := ic.commitDistributionManifest(ctx, r, config, layers[layersMap[p.ID]], oci, inp.Metadata[fmt.Sprintf("%s/%s", exptypes.ExporterInlineCache, p.ID)])
 		if err != nil {
 			return nil, err
 		}
@@ -173,7 +173,7 @@ func (ic *ImageWriter) exportLayers(ctx context.Context, refs ...cache.Immutable
 	return out, nil
 }
 
-func (ic *ImageWriter) commitDistributionManifest(ctx context.Context, ref cache.ImmutableRef, config []byte, layers []blobs.DiffPair, oci bool) (*ocispec.Descriptor, error) {
+func (ic *ImageWriter) commitDistributionManifest(ctx context.Context, ref cache.ImmutableRef, config []byte, layers []blobs.DiffPair, oci bool, cache []byte) (*ocispec.Descriptor, error) {
 	if len(config) == 0 {
 		var err error
 		config, err = emptyImageConfig()
@@ -189,7 +189,7 @@ func (ic *ImageWriter) commitDistributionManifest(ctx context.Context, ref cache
 
 	diffPairs, history := normalizeLayersAndHistory(layers, history, ref)
 
-	config, err = patchImageConfig(config, diffPairs, history)
+	config, err = patchImageConfig(config, diffPairs, history, cache)
 	if err != nil {
 		return nil, err
 	}
@@ -312,7 +312,7 @@ func parseHistoryFromConfig(dt []byte) ([]ocispec.History, error) {
 	return config.History, nil
 }
 
-func patchImageConfig(dt []byte, dps []blobs.DiffPair, history []ocispec.History) ([]byte, error) {
+func patchImageConfig(dt []byte, dps []blobs.DiffPair, history []ocispec.History, cache []byte) ([]byte, error) {
 	m := map[string]json.RawMessage{}
 	if err := json.Unmarshal(dt, &m); err != nil {
 		return nil, errors.Wrap(err, "failed to parse image config for patch")
@@ -349,6 +349,10 @@ func patchImageConfig(dt []byte, dps []blobs.DiffPair, history []ocispec.History
 		m["created"] = dt
 	}
 
+	if cache != nil {
+		m["moby.buildkit.cache.v0"] = cache
+	}
+
 	dt, err = json.Marshal(m)
 	return dt, errors.Wrap(err, "failed to marshal config after patch")
 }

solver/llbsolver/solver.go

@@ -2,6 +2,7 @@ package llbsolver
 
 import (
 	"context"
+	"fmt"
 	"strings"
 	"time"
 
@@ -10,6 +11,7 @@ import (
 	"github.com/moby/buildkit/client"
 	controlgateway "github.com/moby/buildkit/control/gateway"
 	"github.com/moby/buildkit/exporter"
+	"github.com/moby/buildkit/exporter/containerimage/exptypes"
 	"github.com/moby/buildkit/frontend"
 	"github.com/moby/buildkit/frontend/gateway"
 	"github.com/moby/buildkit/identity"
@@ -138,7 +140,7 @@ func (s *Solver) Solve(ctx context.Context, id string, req frontend.SolveRequest
 	}()
 
 	var exporterResponse map[string]string
-	if exp := exp.Exporter; exp != nil {
+	if e := exp.Exporter; e != nil {
 		inp := exporter.Source{
 			Metadata: res.Metadata,
 		}
@@ -151,6 +153,14 @@ func (s *Solver) Solve(ctx context.Context, id string, req frontend.SolveRequest
 				return nil, errors.Errorf("invalid reference: %T", res.Sys())
 			}
 			inp.Ref = workerRef.ImmutableRef
+
+			dt, err := inlineCache(ctx, exp.CacheExporter, res)
+			if err != nil {
+				return nil, err
+			}
+			if dt != nil {
+				inp.Metadata[exptypes.ExporterInlineCache] = dt
+			}
 		}
 		if res.Refs != nil {
 			m := make(map[string]cache.ImmutableRef, len(res.Refs))
@@ -163,13 +173,21 @@ func (s *Solver) Solve(ctx context.Context, id string, req frontend.SolveRequest
 						return nil, errors.Errorf("invalid reference: %T", res.Sys())
 					}
 					m[k] = workerRef.ImmutableRef
+
+					dt, err := inlineCache(ctx, exp.CacheExporter, res)
+					if err != nil {
+						return nil, err
+					}
+					if dt != nil {
+						inp.Metadata[fmt.Sprintf("%s/%s", exptypes.ExporterInlineCache, k)] = dt
+					}
 				}
 			}
 			inp.Refs = m
 		}
 
-		if err := inVertexContext(j.Context(ctx), exp.Name(), "", func(ctx context.Context) error {
-			exporterResponse, err = exp.Export(ctx, inp)
+		if err := inVertexContext(j.Context(ctx), e.Name(), "", func(ctx context.Context) error {
+			exporterResponse, err = e.Export(ctx, inp)
 			return err
 		}); err != nil {
 			return nil, err
@@ -218,6 +236,37 @@ func (s *Solver) Solve(ctx context.Context, id string, req frontend.SolveRequest
 	}, nil
 }
 
+func inlineCache(ctx context.Context, e remotecache.Exporter, res solver.CachedResult) ([]byte, error) {
+	if efl, ok := e.(interface {
+		ExportForLayers([]digest.Digest) ([]byte, error)
+	}); ok {
+		workerRef, ok := res.Sys().(*worker.WorkerRef)
+		if !ok {
+			return nil, errors.Errorf("invalid reference: %T", res.Sys())
+		}
+
+		remote, err := workerRef.Worker.GetRemote(ctx, workerRef.ImmutableRef, true)
+		if err != nil || remote == nil {
+			return nil, nil
+		}
+
+		digests := make([]digest.Digest, 0, len(remote.Descriptors))
+		for _, desc := range remote.Descriptors {
+			digests = append(digests, desc.Digest)
+		}
+
+		if _, err := res.CacheKeys()[0].Exporter.ExportTo(ctx, e, solver.CacheExportOpt{
+			Convert: workerRefConverter,
+			Mode:    solver.CacheExportModeMin,
+		}); err != nil {
+			return nil, err
+		}
+
+		return efl.ExportForLayers(digests)
+	}
+	return nil, nil
+}
+
 func (s *Solver) Status(ctx context.Context, id string, statusChan chan *client.SolveStatus) error {
 	j, err := s.solver.Get(id)
 	if err != nil {

util/imageutil/config.go

@@ -135,17 +135,21 @@ func childrenConfigHandler(provider content.Provider, platform platforms.MatchCo
 func DetectManifestMediaType(ra content.ReaderAt) (string, error) {
 	// TODO: schema1
 
-	p := make([]byte, ra.Size())
-	if _, err := ra.ReadAt(p, 0); err != nil {
+	dt := make([]byte, ra.Size())
+	if _, err := ra.ReadAt(dt, 0); err != nil {
 		return "", err
 	}
 
+	return DetectManifestBlobMediaType(dt)
+}
+
+func DetectManifestBlobMediaType(dt []byte) (string, error) {
 	var mfst struct {
 		MediaType string          `json:"mediaType"`
 		Config    json.RawMessage `json:"config"`
 	}
 
-	if err := json.Unmarshal(p, &mfst); err != nil {
+	if err := json.Unmarshal(dt, &mfst); err != nil {
 		return "", err
 	}