solver: add selector support to cache

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

solver-next/edge.go

@@ -122,12 +122,16 @@ func (e *edge) commitOptions() (CacheKey, []Result) {
 		return NewCacheKey(e.cacheMap.Digest, e.edge.Index, nil), nil
 	}
 
-	inputs := make([]CacheKey, len(e.deps))
+	inputs := make([]CacheKeyWithSelector, len(e.deps))
 	results := make([]Result, len(e.deps))
 	for i, dep := range e.deps {
-		inputs[i] = dep.result.CacheKey()
+		inputs[i] = CacheKeyWithSelector{CacheKey: dep.result.CacheKey(), Selector: e.cacheMap.Deps[i].Selector}
 		if dep.slowCacheKey != nil {
-			inputs[i] = NewCacheKey("", 0, []CacheKey{inputs[i], dep.slowCacheKey})
+			inputs[i] = CacheKeyWithSelector{CacheKey: NewCacheKey("", 0, []CacheKeyWithSelector{
+				inputs[i],
+				{CacheKey: dep.slowCacheKey, Selector: NoSelector},
+			},
+			)}
 		}
 		results[i] = dep.result
 	}
@@ -165,7 +169,7 @@ func (e *edge) probeCache(d *dep, keys []CacheKey) {
 	if e.op.IgnoreCache() {
 		return
 	}
-	records, err := e.op.Cache().Query(keys, d.index, e.cacheMap.Digest, e.edge.Index)
+	records, err := e.op.Cache().Query(keys, d.index, e.cacheMap.Digest, e.edge.Index, e.cacheMap.Deps[d.index].Selector)
 	if err != nil {
 		e.err = errors.Wrap(err, "error on cache query")
 	}
@@ -291,7 +295,7 @@ func (e *edge) processUpdate(upt pipe.Receiver) (depChanged bool) {
 			if len(e.deps) == 0 {
 				if !e.op.IgnoreCache() {
 					k := NewCacheKey(e.cacheMap.Digest, e.edge.Index, nil)
-					records, err := e.op.Cache().Query(nil, 0, e.cacheMap.Digest, e.edge.Index)
+					records, err := e.op.Cache().Query(nil, 0, e.cacheMap.Digest, e.edge.Index, "")
 					if err != nil {
 						logrus.Error(errors.Wrap(err, "invalid query response")) // make the build fail for this error
 					} else {

solver-next/index.go

@@ -186,7 +186,8 @@ func getUniqueID(k CacheKey) digest.Digest {
 
 	dgstr := digest.SHA256.Digester()
 	for _, inp := range k.Deps() {
-		dgstr.Hash().Write([]byte(getUniqueID(inp)))
+		dgstr.Hash().Write([]byte(getUniqueID(inp.CacheKey)))
+		dgstr.Hash().Write([]byte(inp.Selector))
 	}
 
 	dgstr.Hash().Write([]byte(k.Digest()))

solver-next/memorycache.go

@@ -17,6 +17,8 @@ type internalMemoryKeyT string
 
 var internalMemoryKey = internalMemoryKeyT("buildkit/memory-cache-id")
 
+var NoSelector = digest.FromBytes(nil)
+
 func NewInMemoryCacheManager() CacheManager {
 	return &inMemoryCacheManager{
 		byID: map[string]*inMemoryCacheKey{},
@@ -29,13 +31,13 @@ type inMemoryCacheKey struct {
 	id     string
 	dgst   digest.Digest
 	output Index
-	deps   []CacheKey // only []*inMemoryCacheManager
+	deps   []CacheKeyWithSelector // only []*inMemoryCacheKey
 
 	results map[Index]map[string]savedResult
 	links   map[link]map[string]struct{}
 }
 
-func (ck *inMemoryCacheKey) Deps() []CacheKey {
+func (ck *inMemoryCacheKey) Deps() []CacheKeyWithSelector {
 	return ck.deps
 }
 func (ck *inMemoryCacheKey) Digest() digest.Digest {
@@ -53,6 +55,7 @@ type savedResult struct {
 type link struct {
 	input, output Index
 	digest        digest.Digest
+	selector      digest.Digest
 }
 
 type inMemoryCacheManager struct {
@@ -65,7 +68,7 @@ func (c *inMemoryCacheManager) ID() string {
 	return c.id
 }
 
-func (c *inMemoryCacheManager) Query(deps []CacheKey, input Index, dgst digest.Digest, output Index) ([]*CacheRecord, error) {
+func (c *inMemoryCacheManager) Query(deps []CacheKey, input Index, dgst digest.Digest, output Index, selector digest.Digest) ([]*CacheRecord, error) {
 	c.mu.RLock()
 	defer c.mu.RUnlock()
 
@@ -75,10 +78,13 @@ func (c *inMemoryCacheManager) Query(deps []CacheKey, input Index, dgst digest.D
 	for _, dep := range deps {
 		ck, err := c.getInternalKey(dep, false)
 		if err == nil {
-			for key := range ck.links[link{input, output, dgst}] {
+			for key := range ck.links[link{input, output, dgst, selector}] {
 				refs[key] = struct{}{}
 			}
-			for key := range ck.links[link{Index(-1), Index(0), ""}] {
+			for key := range ck.links[link{Index(-1), Index(0), "", selector}] {
+				sublinks[key] = struct{}{}
+			}
+			for key := range ck.links[link{Index(-1), Index(0), "", NoSelector}] {
 				sublinks[key] = struct{}{}
 			}
 		}
@@ -86,7 +92,7 @@ func (c *inMemoryCacheManager) Query(deps []CacheKey, input Index, dgst digest.D
 
 	for id := range sublinks {
 		if ck, ok := c.byID[id]; ok {
-			for key := range ck.links[link{input, output, dgst}] {
+			for key := range ck.links[link{input, output, dgst, ""}] {
 				refs[key] = struct{}{}
 			}
 		}
@@ -165,15 +171,15 @@ func (c *inMemoryCacheManager) getInternalKey(k CacheKey, createIfNotExist bool)
 		}
 		return ck, nil
 	}
-	inputs := make([]CacheKey, len(k.Deps()))
+	inputs := make([]CacheKeyWithSelector, len(k.Deps()))
 	dgstr := digest.SHA256.Digester()
 	for i, inp := range k.Deps() {
-		ck, err := c.getInternalKey(inp, createIfNotExist)
+		ck, err := c.getInternalKey(inp.CacheKey, createIfNotExist)
 		if err != nil {
 			return nil, err
 		}
-		inputs[i] = ck
+		inputs[i] = CacheKeyWithSelector{CacheKey: ck, Selector: inp.Selector}
-		if _, err := dgstr.Hash().Write([]byte(ck.id)); err != nil {
+		if _, err := dgstr.Hash().Write([]byte(fmt.Sprintf("%s:%s,", ck.id, inp.Selector))); err != nil {
 			return nil, err
 		}
 	}
@@ -209,7 +215,7 @@ func (c *inMemoryCacheManager) getInternalKey(k CacheKey, createIfNotExist bool)
 		if ck.dgst == "" {
 			i = -1
 		}
-		if err := c.addLink(link{Index(i), ck.output, ck.dgst}, inp.(*inMemoryCacheKey), ck); err != nil {
+		if err := c.addLink(link{Index(i), ck.output, ck.dgst, inp.Selector}, inp.CacheKey.(*inMemoryCacheKey), ck); err != nil {
 			return nil, err
 		}
 	}
@@ -259,14 +265,14 @@ func (cm *combinedCacheManager) ID() string {
 	return cm.id
 }
 
-func (cm *combinedCacheManager) Query(inp []CacheKey, inputIndex Index, dgst digest.Digest, outputIndex Index) ([]*CacheRecord, error) {
+func (cm *combinedCacheManager) Query(inp []CacheKey, inputIndex Index, dgst digest.Digest, outputIndex Index, selector digest.Digest) ([]*CacheRecord, error) {
 	eg, _ := errgroup.WithContext(context.TODO())
 	res := make(map[string]*CacheRecord, len(cm.cms))
 	var mu sync.Mutex
 	for i, c := range cm.cms {
 		func(i int, c CacheManager) {
 			eg.Go(func() error {
-				recs, err := c.Query(inp, inputIndex, dgst, outputIndex)
+				recs, err := c.Query(inp, inputIndex, dgst, outputIndex, selector)
 				if err != nil {
 					return err
 				}

solver-next/memorycache_test.go

@@ -17,7 +17,7 @@ func TestInMemoryCache(t *testing.T) {
 	cacheFoo, err := m.Save(NewCacheKey(dgst("foo"), 0, nil), testResult("result0"))
 	require.NoError(t, err)
 
-	matches, err := m.Query(nil, 0, dgst("foo"), 0)
+	matches, err := m.Query(nil, 0, dgst("foo"), 0, "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -29,7 +29,7 @@ func TestInMemoryCache(t *testing.T) {
 	cacheBar, err := m.Save(NewCacheKey(dgst("bar"), 0, nil), testResult("result1"))
 	require.NoError(t, err)
 
-	matches, err = m.Query(nil, 0, dgst("bar"), 0)
+	matches, err = m.Query(nil, 0, dgst("bar"), 0, "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -38,30 +38,30 @@ func TestInMemoryCache(t *testing.T) {
 	require.Equal(t, "result1", unwrap(res))
 
 	// invalid request
-	matches, err = m.Query(nil, 0, dgst("baz"), 0)
+	matches, err = m.Query(nil, 0, dgst("baz"), 0, "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 0)
 
 	// second level
-	k := NewCacheKey(dgst("baz"), Index(1), []CacheKey{
+	k := NewCacheKey(dgst("baz"), Index(1), []CacheKeyWithSelector{
-		cacheFoo, cacheBar,
+		{CacheKey: cacheFoo}, {CacheKey: cacheBar},
 	})
 	cacheBaz, err := m.Save(k, testResult("result2"))
 	require.NoError(t, err)
 
-	matches, err = m.Query(nil, 0, dgst("baz"), 0)
+	matches, err = m.Query(nil, 0, dgst("baz"), 0, "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 0)
 
-	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("baz"), 0)
+	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("baz"), 0, "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 0)
 
-	matches, err = m.Query([]CacheKey{cacheFoo}, 1, dgst("baz"), Index(1))
+	matches, err = m.Query([]CacheKey{cacheFoo}, 1, dgst("baz"), Index(1), "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 0)
 
-	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("baz"), Index(1))
+	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("baz"), Index(1), "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
@@ -69,50 +69,129 @@ func TestInMemoryCache(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, "result2", unwrap(res))
 
-	matches2, err := m.Query([]CacheKey{cacheBar}, 1, dgst("baz"), Index(1))
+	matches2, err := m.Query([]CacheKey{cacheBar}, 1, dgst("baz"), Index(1), "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches2), 1)
 
 	require.Equal(t, matches[0].ID, matches2[0].ID)
 
-	k = NewCacheKey(dgst("baz"), Index(1), []CacheKey{
+	k = NewCacheKey(dgst("baz"), Index(1), []CacheKeyWithSelector{
-		cacheFoo,
+		{CacheKey: cacheFoo},
 	})
 	_, err = m.Save(k, testResult("result3"))
 	require.NoError(t, err)
 
-	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("baz"), Index(1))
+	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("baz"), Index(1), "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 2)
 
 	// combination save
-	k2 := NewCacheKey("", 0, []CacheKey{
+	k2 := NewCacheKey("", 0, []CacheKeyWithSelector{
-		cacheFoo, cacheBaz,
+		{CacheKey: cacheFoo}, {CacheKey: cacheBaz},
 	})
 
-	k = NewCacheKey(dgst("bax"), 0, []CacheKey{
+	k = NewCacheKey(dgst("bax"), 0, []CacheKeyWithSelector{
-		k2, cacheBar,
+		{CacheKey: k2}, {CacheKey: cacheBar},
 	})
 	_, err = m.Save(k, testResult("result4"))
 	require.NoError(t, err)
 
 	// foo, bar, baz should all point to result4
-	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("bax"), 0)
+	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("bax"), 0, "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 
 	id := matches[0].ID
 
-	matches, err = m.Query([]CacheKey{cacheBar}, 1, dgst("bax"), 0)
+	matches, err = m.Query([]CacheKey{cacheBar}, 1, dgst("bax"), 0, "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 	require.Equal(t, matches[0].ID, id)
 
-	matches, err = m.Query([]CacheKey{cacheBaz}, 0, dgst("bax"), 0)
+	matches, err = m.Query([]CacheKey{cacheBaz}, 0, dgst("bax"), 0, "")
 	require.NoError(t, err)
 	require.Equal(t, len(matches), 1)
 	require.Equal(t, matches[0].ID, id)
+}
 
+func TestInMemoryCacheSelector(t *testing.T) {
+	ctx := context.TODO()
+
+	m := NewInMemoryCacheManager()
+
+	cacheFoo, err := m.Save(NewCacheKey(dgst("foo"), 0, nil), testResult("result0"))
+	require.NoError(t, err)
+
+	_, err = m.Save(NewCacheKey(dgst("bar"), 0, []CacheKeyWithSelector{
+		{CacheKey: cacheFoo, Selector: dgst("sel0")},
+	}), testResult("result1"))
+	require.NoError(t, err)
+
+	matches, err := m.Query([]CacheKey{cacheFoo}, 0, dgst("bar"), 0, "")
+	require.NoError(t, err)
+	require.Equal(t, len(matches), 0)
+
+	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("bar"), 0, "sel-invalid")
+	require.NoError(t, err)
+	require.Equal(t, len(matches), 0)
+
+	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("bar"), 0, dgst("sel0"))
+	require.NoError(t, err)
+	require.Equal(t, len(matches), 1)
+
+	res, err := m.Load(ctx, matches[0])
+	require.NoError(t, err)
+	require.Equal(t, "result1", unwrap(res))
+}
+
+func TestInMemoryCacheSelectorNested(t *testing.T) {
+	ctx := context.TODO()
+
+	m := NewInMemoryCacheManager()
+
+	cacheFoo, err := m.Save(NewCacheKey(dgst("foo"), 0, nil), testResult("result0"))
+	require.NoError(t, err)
+
+	k2 := NewCacheKey("", 0, []CacheKeyWithSelector{
+		{CacheKey: cacheFoo, Selector: dgst("sel0")},
+		{CacheKey: NewCacheKey(dgst("second"), 0, nil), Selector: NoSelector},
+	})
+
+	_, err = m.Save(NewCacheKey(dgst("bar"), 0, []CacheKeyWithSelector{
+		{CacheKey: k2},
+	}), testResult("result1"))
+	require.NoError(t, err)
+
+	matches, err := m.Query([]CacheKey{cacheFoo}, 0, dgst("bar"), 0, dgst("sel0"))
+	require.NoError(t, err)
+	require.Equal(t, len(matches), 1)
+	res, err := m.Load(ctx, matches[0])
+	require.NoError(t, err)
+	require.Equal(t, "result1", unwrap(res))
+
+	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("bar"), 0, "")
+	require.NoError(t, err)
+	require.Equal(t, len(matches), 0)
+
+	matches, err = m.Query([]CacheKey{cacheFoo}, 0, dgst("bar"), 0, dgst("bar"))
+	require.NoError(t, err)
+	require.Equal(t, len(matches), 0)
+
+	// keys with NoSelector always match
+	matches, err = m.Query([]CacheKey{NewCacheKey(dgst("second"), 0, nil)}, 0, dgst("bar"), 0, dgst("sel0"))
+	require.NoError(t, err)
+	require.Equal(t, len(matches), 1)
+	res, err = m.Load(ctx, matches[0])
+	require.NoError(t, err)
+	require.Equal(t, "result1", unwrap(res))
+
+	matches, err = m.Query([]CacheKey{NewCacheKey(dgst("second"), 0, nil)}, 0, dgst("bar"), 0, "")
+	require.NoError(t, err)
+	require.Equal(t, len(matches), 1)
+
+	matches, err = m.Query([]CacheKey{NewCacheKey(dgst("second"), 0, nil)}, 0, dgst("bar"), 0, dgst("bar"))
+	require.NoError(t, err)
+	require.Equal(t, len(matches), 1)
 }
 
 func dgst(s string) digest.Digest {

solver-next/types.go

@@ -68,7 +68,6 @@ type CacheMap struct {
 	Digest digest.Digest
 	Deps   []struct {
 		// Optional digest that is merged with the cache key of the input
-		// TODO(tonistiigi): not implemented
 		Selector digest.Digest
 		// Optional function that returns a digest for the input based on its
 		// return value
@@ -79,7 +78,7 @@ type CacheMap struct {
 // CacheKey is an identifier for storing/loading build cache
 type CacheKey interface {
 	// Deps are dependant cache keys
-	Deps() []CacheKey
+	Deps() []CacheKeyWithSelector
 	// Base digest for operation. Usually CacheMap.Digest
 	Digest() digest.Digest
 	// Index for the output that is cached
@@ -107,7 +106,7 @@ type CacheManager interface {
 	ID() string
 	// Query searches for cache paths from one cache key to the output of a
 	// possible match.
-	Query(inp []CacheKey, inputIndex Index, dgst digest.Digest, outputIndex Index) ([]*CacheRecord, error)
+	Query(inp []CacheKey, inputIndex Index, dgst digest.Digest, outputIndex Index, selector digest.Digest) ([]*CacheRecord, error)
 	// Load pulls and returns the cached result
 	Load(ctx context.Context, rec *CacheRecord) (Result, error)
 	// Save saves a result based on a cache key
@@ -115,7 +114,7 @@ type CacheManager interface {
 }
 
 // NewCacheKey creates a new cache key for a specific output index
-func NewCacheKey(dgst digest.Digest, index Index, deps []CacheKey) CacheKey {
+func NewCacheKey(dgst digest.Digest, index Index, deps []CacheKeyWithSelector) CacheKey {
 	return &cacheKey{
 		dgst:   dgst,
 		deps:   deps,
@@ -124,10 +123,17 @@ func NewCacheKey(dgst digest.Digest, index Index, deps []CacheKey) CacheKey {
 	}
 }
 
+// CacheKeyWithSelector combines a cache key with an optional selector digest.
+// Used to limit the matches for dependency cache key.
+type CacheKeyWithSelector struct {
+	Selector digest.Digest
+	CacheKey CacheKey
+}
+
 type cacheKey struct {
 	dgst   digest.Digest
 	index  Index
-	deps   []CacheKey
+	deps   []CacheKeyWithSelector
 	values *sync.Map
 }
 
@@ -140,7 +146,7 @@ func (ck *cacheKey) GetValue(key interface{}) interface{} {
 	return v
 }
 
-func (ck *cacheKey) Deps() []CacheKey {
+func (ck *cacheKey) Deps() []CacheKeyWithSelector {
 	return ck.deps
 }