solver: implement content based cache support

Signed-off-by: Tonis Tiigi <[email protected]>

cache/contenthash/checksum.go

@@ -47,12 +47,12 @@ func Checksum(ctx context.Context, ref cache.ImmutableRef, path string) (digest.
 	return getDefaultManager().Checksum(ctx, ref, path)
 }
 
-func GetCacheContext(ctx context.Context, ref cache.ImmutableRef) (CacheContext, error) {
-	return getDefaultManager().GetCacheContext(ctx, ref)
+func GetCacheContext(ctx context.Context, md *metadata.StorageItem) (CacheContext, error) {
+	return getDefaultManager().GetCacheContext(ctx, md)
 }
 
-func SetCacheContext(ctx context.Context, ref cache.ImmutableRef, cc CacheContext) error {
-	return getDefaultManager().SetCacheContext(ctx, ref, cc)
+func SetCacheContext(ctx context.Context, md *metadata.StorageItem, cc CacheContext) error {
+	return getDefaultManager().SetCacheContext(ctx, md, cc)
 }
 
 type CacheContext interface {
@@ -67,45 +67,57 @@ type Hashed interface {
 type cacheManager struct {
 	locker *locker.Locker
 	lru    *simplelru.LRU
+	lruMu  sync.Mutex
 }
 
 func (cm *cacheManager) Checksum(ctx context.Context, ref cache.ImmutableRef, p string) (digest.Digest, error) {
-	cc, err := cm.GetCacheContext(ctx, ref)
+	cc, err := cm.GetCacheContext(ctx, ensureOriginMetadata(ref.Metadata()))
 	if err != nil {
 		return "", nil
 	}
 	return cc.Checksum(ctx, ref, p)
 }
 
-func (cm *cacheManager) GetCacheContext(ctx context.Context, ref cache.ImmutableRef) (CacheContext, error) {
-	cm.locker.Lock(ref.ID())
-	v, ok := cm.lru.Get(ref.ID())
+func (cm *cacheManager) GetCacheContext(ctx context.Context, md *metadata.StorageItem) (CacheContext, error) {
+	cm.locker.Lock(md.ID())
+	cm.lruMu.Lock()
+	v, ok := cm.lru.Get(md.ID())
+	cm.lruMu.Unlock()
 	if ok {
-		cm.locker.Unlock(ref.ID())
+		cm.locker.Unlock(md.ID())
 		return v.(*cacheContext), nil
 	}
-	cc, err := newCacheContext(ref.Metadata())
+	cc, err := newCacheContext(md)
 	if err != nil {
-		cm.locker.Unlock(ref.ID())
+		cm.locker.Unlock(md.ID())
 		return nil, err
 	}
-	cm.lru.Add(ref.ID(), cc)
-	cm.locker.Unlock(ref.ID())
+	cm.lruMu.Lock()
+	cm.lru.Add(md.ID(), cc)
+	cm.lruMu.Unlock()
+	cm.locker.Unlock(md.ID())
 	return cc, nil
 }
 
-func (cm *cacheManager) SetCacheContext(ctx context.Context, ref cache.ImmutableRef, cci CacheContext) error {
+func (cm *cacheManager) SetCacheContext(ctx context.Context, md *metadata.StorageItem, cci CacheContext) error {
 	cc, ok := cci.(*cacheContext)
 	if !ok {
 		return errors.Errorf("invalid cachecontext: %T", cc)
 	}
-	if ref.ID() != cc.md.ID() {
-		return errors.New("saving cachecontext under different ID not supported")
-	}
-	if err := cc.save(); err != nil {
-		return err
+	if md.ID() != cc.md.ID() {
+		cc = &cacheContext{
+			md:       md,
+			tree:     cci.(*cacheContext).tree,
+			dirtyMap: map[string]struct{}{},
+		}
+	} else {
+		if err := cc.save(); err != nil {
+			return err
+		}
 	}
-	cm.lru.Add(ref.ID(), cc)
+	cm.lruMu.Lock()
+	cm.lru.Add(md.ID(), cc)
+	cm.lruMu.Unlock()
 	return nil
 }
 
@@ -193,7 +205,9 @@ func (cc *cacheContext) save() error {
 	cc.mu.Lock()
 	defer cc.mu.Unlock()
 
-	cc.dirty = true
+	if cc.txn != nil {
+		cc.commitActiveTransaction()
+	}
 
 	var l CacheRecords
 	node := cc.tree.Root()
@@ -231,10 +245,23 @@ func (cc *cacheContext) HandleChange(kind fsutil.ChangeKind, p string, fi os.Fil
 	}
 
 	cc.mu.Lock()
+	defer cc.mu.Unlock()
 	if cc.txn == nil {
 		cc.txn = cc.tree.Txn()
 		cc.node = cc.tree.Root()
+
+		// root is not called by HandleChange. need to fake it
+		if _, ok := cc.node.Get([]byte("/")); !ok {
+			cc.txn.Insert([]byte("/"), &CacheRecord{
+				Type:   CacheRecordTypeDirHeader,
+				Digest: digest.FromBytes(nil),
+			})
+			cc.txn.Insert([]byte(""), &CacheRecord{
+				Type: CacheRecordTypeDir,
+			})
+		}
 	}
+
 	if kind == fsutil.ChangeKindDelete {
 		v, ok := cc.txn.Delete(k)
 		if ok {
@@ -245,7 +272,6 @@ func (cc *cacheContext) HandleChange(kind fsutil.ChangeKind, p string, fi os.Fil
 			d = ""
 		}
 		cc.dirtyMap[d] = struct{}{}
-		cc.mu.Unlock()
 		return
 	}
 
@@ -256,7 +282,6 @@ func (cc *cacheContext) HandleChange(kind fsutil.ChangeKind, p string, fi os.Fil
 
 	h, ok := fi.(Hashed)
 	if !ok {
-		cc.mu.Unlock()
 		return errors.Errorf("invalid fileinfo: %s", p)
 	}
 
@@ -287,7 +312,6 @@ func (cc *cacheContext) HandleChange(kind fsutil.ChangeKind, p string, fi os.Fil
 		d = ""
 	}
 	cc.dirtyMap[d] = struct{}{}
-	cc.mu.Unlock()
 
 	return nil
 }
@@ -405,12 +429,12 @@ func (cc *cacheContext) checksum(ctx context.Context, root *iradix.Node, txn *ir
 	switch cr.Type {
 	case CacheRecordTypeDir:
 		h := sha256.New()
-		iter := root.Iterator()
 		next := append(k, []byte("/")...)
-		iter.SeekPrefix(next)
+		iter := root.Seek(next)
+		subk := next
+		ok := true
 		for {
-			subk, _, ok := iter.Next()
-			if !ok || bytes.Compare(next, subk) > 0 {
+			if !ok || !bytes.HasPrefix(subk, next) {
 				break
 			}
 			h.Write(bytes.TrimPrefix(subk, k))
@@ -422,9 +446,10 @@ func (cc *cacheContext) checksum(ctx context.Context, root *iradix.Node, txn *ir
 
 			h.Write([]byte(subcr.Digest))
 			if subcr.Type == CacheRecordTypeDir { // skip subfiles
-				next = append(k, []byte("/\xff")...)
-				iter.SeekPrefix(next)
+				next := append(subk, []byte("/\xff")...)
+				iter = root.Seek(next)
 			}
+			subk, _, ok = iter.Next()
 		}
 		dgst = digest.NewDigest(digest.SHA256, h)
 	default:
@@ -565,3 +590,19 @@ func addParentToMap(d string, m map[string]struct{}) {
 	m[d] = struct{}{}
 	addParentToMap(d, m)
 }
+
+func ensureOriginMetadata(md *metadata.StorageItem) *metadata.StorageItem {
+	v := md.Get("cache.equalMutable") // TODO: const
+	if v == nil {
+		return md
+	}
+	var mutable string
+	if err := v.Unmarshal(&mutable); err != nil {
+		return md
+	}
+	si, ok := md.Storage().Get(mutable)
+	if ok {
+		return &si
+	}
+	return md
+}

cache/contenthash/checksum_test.go

@@ -92,7 +92,7 @@ func TestChecksumBasicFile(t *testing.T) {
 	dgst, err = cc.Checksum(context.TODO(), ref, "/")
 	assert.NoError(t, err)
 
-	assert.Equal(t, digest.Digest("sha256:0d87c8c2a606f961483cd4c5dc0350a4136a299b4066eea4a969d6ed756614cd"), dgst)
+	assert.Equal(t, digest.Digest("sha256:7378af5287e8b417b6cbc63154d300e130983bfc645e35e86fdadf6f5060468a"), dgst)
 
 	dgst, err = cc.Checksum(context.TODO(), ref, "d0")
 	assert.NoError(t, err)

cache/instructioncache/cache.go

@@ -1,22 +1,26 @@
 package instructioncache
 
 import (
+	"strings"
+
 	"github.com/boltdb/bolt"
 	"github.com/moby/buildkit/cache"
 	"github.com/moby/buildkit/cache/metadata"
+	digest "github.com/opencontainers/go-digest"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/net/context"
 )
 
 const cacheKey = "buildkit.instructioncache"
+const contentCacheKey = "buildkit.instructioncache.content"
 
 type LocalStore struct {
 	MetadataStore *metadata.Store
 	Cache         cache.Accessor
 }
 
-func (ls *LocalStore) Set(key string, value interface{}) error {
+func (ls *LocalStore) Set(key digest.Digest, value interface{}) error {
 	ref, ok := value.(cache.ImmutableRef)
 	if !ok {
 		return errors.Errorf("invalid ref")
@@ -25,21 +29,21 @@ func (ls *LocalStore) Set(key string, value interface{}) error {
 	if err != nil {
 		return err
 	}
-	v.Index = index(key)
+	v.Index = index(key.String())
 	si, _ := ls.MetadataStore.Get(ref.ID())
 	return si.Update(func(b *bolt.Bucket) error {
-		return si.SetValue(b, index(key), v)
+		return si.SetValue(b, v.Index, v)
 	})
 }
 
-func (ls *LocalStore) Lookup(ctx context.Context, key string) (interface{}, error) {
-	snaps, err := ls.MetadataStore.Search(index(key))
+func (ls *LocalStore) Lookup(ctx context.Context, key digest.Digest) (interface{}, error) {
+	snaps, err := ls.MetadataStore.Search(index(key.String()))
 	if err != nil {
 		return nil, err
 	}
 
 	for _, s := range snaps {
-		v := s.Get(index(key))
+		v := s.Get(index(key.String()))
 		if v != nil {
 			var id string
 			if err = v.Unmarshal(&id); err != nil {
@@ -56,6 +60,42 @@ func (ls *LocalStore) Lookup(ctx context.Context, key string) (interface{}, erro
 	return nil, nil
 }
 
+func (ls *LocalStore) SetContentMapping(key digest.Digest, value interface{}) error {
+	ref, ok := value.(cache.ImmutableRef)
+	if !ok {
+		return errors.Errorf("invalid ref")
+	}
+	v, err := metadata.NewValue(ref.ID())
+	if err != nil {
+		return err
+	}
+	v.Index = contentIndex(key.String())
+	si, _ := ls.MetadataStore.Get(ref.ID())
+	return si.Update(func(b *bolt.Bucket) error {
+		return si.SetValue(b, v.Index, v)
+	})
+}
+
+func (ls *LocalStore) GetContentMapping(key digest.Digest) ([]digest.Digest, error) {
+	snaps, err := ls.MetadataStore.Search(contentIndex(key.String()))
+	if err != nil {
+		return nil, err
+	}
+	var out []digest.Digest
+	for _, s := range snaps {
+		for _, k := range s.Keys() {
+			if strings.HasPrefix(k, index("")) {
+				out = append(out, digest.Digest(strings.TrimPrefix(k, index("")))) // TODO: type
+			}
+		}
+	}
+	return out, nil
+}
+
 func index(k string) string {
 	return cacheKey + "::" + k
 }
+
+func contentIndex(k string) string {
+	return contentCacheKey + "::" + k
+}

cache/metadata.go

@@ -18,7 +18,6 @@ import (
 const sizeUnknown int64 = -1
 const keySize = "snapshot.size"
 const keyEqualMutable = "cache.equalMutable"
-const keyEqualImmutable = "cache.equalImmutable"
 const keyCachePolicy = "cache.cachePolicy"
 const keyDescription = "cache.description"
 const keyCreatedAt = "cache.createdAt"

cache/metadata/metadata.go

@@ -205,6 +205,10 @@ func newStorageItem(id string, b *bolt.Bucket, s *Store) (StorageItem, error) {
 	return si, nil
 }
 
+func (s *StorageItem) Storage() *Store { // TODO: used in local source. how to remove this?
+	return s.storage
+}
+
 func (s *StorageItem) ID() string {
 	return s.id
 }

cache/refs.go

@@ -28,6 +28,7 @@ type MutableRef interface {
 	Commit(context.Context) (ImmutableRef, error)
 	Release(context.Context) error
 	Size(ctx context.Context) (int64, error)
+	Metadata() *metadata.StorageItem
 }
 
 type Mountable interface {

session/filesync/diffcopy.go

@@ -21,10 +21,11 @@ func recvDiffCopy(ds grpc.Stream, dest string, cu CacheUpdater, progress progres
 		logrus.Debugf("diffcopy took: %v", time.Since(st))
 	}()
 	var cf fsutil.ChangeFunc
+	var ch fsutil.ContentHasher
 	if cu != nil {
 		cu.MarkSupported(true)
 		cf = cu.HandleChange
+		ch = cu.ContentHasher()
 	}
-	_ = cf
-	return fsutil.Receive(ds.Context(), ds, dest, nil, nil, progress)
+	return fsutil.Receive(ds.Context(), ds, dest, cf, ch, progress)
 }

session/filesync/filesync.go

@@ -147,6 +147,7 @@ type FSSendRequestOpt struct {
 type CacheUpdater interface {
 	MarkSupported(bool)
 	HandleChange(fsutil.ChangeKind, string, os.FileInfo, error) error
+	ContentHasher() fsutil.ContentHasher
 }
 
 // FSSync initializes a transfer of files