solver: add support for multiple cache keys

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

solver-next/boltdbcachestorage/storage.go

@@ -156,6 +156,32 @@ func (s *Store) AddResult(id string, res solver.CacheResult) error {
 	})
 }
 
+func (s *Store) WalkIDsByResult(resultID string, fn func(string) error) error {
+	ids := map[string]struct{}{}
+	if err := s.db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket([]byte(byResultBucket))
+		if b == nil {
+			return nil
+		}
+		b = b.Bucket([]byte(resultID))
+		if b == nil {
+			return nil
+		}
+		return b.ForEach(func(k, v []byte) error {
+			ids[string(k)] = struct{}{}
+			return nil
+		})
+	}); err != nil {
+		return err
+	}
+	for id := range ids {
+		if err := fn(id); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (s *Store) Release(resultID string) error {
 	return s.db.Update(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(byResultBucket))

solver-next/cachestorage.go

@@ -19,6 +19,7 @@ type CacheKeyStorage interface {
 	Load(id string, resultID string) (CacheResult, error)
 	AddResult(id string, res CacheResult) error
 	Release(resultID string) error
+	WalkIDsByResult(resultID string, fn func(string) error) error
 
 	AddLink(id string, link CacheInfoLink, target string) error
 	WalkLinks(id string, link CacheInfoLink, fn func(id string) error) error

solver-next/cachestorage_testsuite.go

@@ -19,6 +19,7 @@ func RunCacheStorageTests(t *testing.T, st func() (CacheKeyStorage, func())) {
 		testResultReleaseSingleLevel,
 		testResultReleaseMultiLevel,
 		testBacklinks,
+		testWalkIDsByResult,
 	} {
 		runStorageTest(t, tc, st)
 	}
@@ -326,6 +327,44 @@ func testResultReleaseMultiLevel(t *testing.T, st CacheKeyStorage) {
 	require.False(t, st.Exists("foo"))
 }
 
+func testWalkIDsByResult(t *testing.T, st CacheKeyStorage) {
+	t.Parallel()
+
+	err := st.AddResult("foo", CacheResult{
+		ID:        "foo-result",
+		CreatedAt: time.Now(),
+	})
+	require.NoError(t, err)
+
+	err = st.AddResult("foo2", CacheResult{
+		ID:        "foo-result",
+		CreatedAt: time.Now(),
+	})
+	require.NoError(t, err)
+
+	err = st.AddResult("bar", CacheResult{
+		ID:        "bar-result",
+		CreatedAt: time.Now(),
+	})
+	require.NoError(t, err)
+
+	m := map[string]struct{}{}
+	err = st.WalkIDsByResult("foo-result", func(id string) error {
+		m[id] = struct{}{}
+		return nil
+	})
+	require.NoError(t, err)
+
+	_, ok := m["foo"]
+	require.True(t, ok)
+
+	_, ok = m["foo2"]
+	require.True(t, ok)
+
+	_, ok = m["bar"]
+	require.False(t, ok)
+}
+
 func getFunctionName(i interface{}) string {
 	fullname := runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
 	dot := strings.LastIndex(fullname, ".") + 1

solver-next/edge.go

@@ -43,11 +43,14 @@ type edge struct {
 	depRequests map[pipe.Receiver]*dep
 	deps        []*dep
 
-	cacheMapReq  pipe.Receiver
-	execReq      pipe.Receiver
-	err          error
-	cacheRecords map[string]*CacheRecord
-	keyMap       map[string]*CacheKey
+	cacheMapReq     pipe.Receiver
+	cacheMapDone    bool
+	cacheMapIndex   int
+	cacheMapDigests []digest.Digest
+	execReq         pipe.Receiver
+	err             error
+	cacheRecords    map[string]*CacheRecord
+	keyMap          map[string]*CacheKey
 
 	noCacheMatchPossible      bool
 	allDepsCompletedCacheFast bool
@@ -132,10 +135,14 @@ func (e *edge) release() {
 }
 
 // commitOptions returns parameters for the op execution
-func (e *edge) commitOptions() (*CacheKey, []CachedResult) {
+func (e *edge) commitOptions() ([]*CacheKey, []CachedResult) {
 	k := NewCacheKey(e.cacheMap.Digest, e.edge.Index)
-	if e.deps == nil {
-		return k, nil
+	if len(e.deps) == 0 {
+		keys := make([]*CacheKey, 0, len(e.cacheMapDigests))
+		for _, dgst := range e.cacheMapDigests {
+			keys = append(keys, NewCacheKey(dgst, e.edge.Index))
+		}
+		return keys, nil
 	}
 
 	inputs := make([][]CacheKeyWithSelector, len(e.deps))
@@ -149,7 +156,7 @@ func (e *edge) commitOptions() (*CacheKey, []CachedResult) {
 	}
 
 	k.deps = inputs
-	return k, results
+	return []*CacheKey{k}, results
 }
 
 // isComplete returns true if edge state is final and will never change
@@ -315,9 +322,10 @@ func (e *edge) unpark(incoming []pipe.Sender, updates, allPipes []pipe.Receiver,
 	}
 
 	// set up new outgoing requests if needed
-	if e.cacheMapReq == nil {
+	if e.cacheMapReq == nil && (e.cacheMap == nil || len(e.cacheRecords) == 0) {
+		index := e.cacheMapIndex
 		e.cacheMapReq = f.NewFuncRequest(func(ctx context.Context) (interface{}, error) {
-			return e.op.CacheMap(ctx)
+			return e.op.CacheMap(ctx, index)
 		})
 	}
 
@@ -359,8 +367,12 @@ func (e *edge) processUpdate(upt pipe.Receiver) (depChanged bool) {
 				e.err = err
 			}
 		} else {
-			e.cacheMap = upt.Status().Value.(*CacheMap)
+			resp := upt.Status().Value.(*cacheMapResp)
+			e.cacheMap = resp.CacheMap
+			e.cacheMapDone = resp.complete
+			e.cacheMapIndex++
 			if len(e.deps) == 0 {
+				e.cacheMapDigests = append(e.cacheMapDigests, e.cacheMap.Digest)
 				if !e.op.IgnoreCache() {
 					keys, err := e.op.Cache().Query(nil, 0, e.cacheMap.Digest, e.edge.Index)
 					if err != nil {
@@ -391,6 +403,9 @@ func (e *edge) processUpdate(upt pipe.Receiver) (depChanged bool) {
 				e.probeCache(dep, withSelector(dep.keys, e.cacheMap.Deps[i].Selector))
 				e.checkDepMatchPossible(dep)
 			}
+			if !e.cacheMapDone {
+				e.cacheMapReq = nil
+			}
 		}
 		return true
 	}
@@ -573,11 +588,14 @@ func (e *edge) recalcCurrentState() {
 		if stHigh > e.state {
 			e.state = stHigh
 		}
+		if !e.cacheMapDone && len(e.keys) == 0 {
+			e.state = edgeStatusInitial
+		}
 
-		e.allDepsCompletedCacheFast = allDepsCompletedCacheFast
-		e.allDepsCompletedCacheSlow = allDepsCompletedCacheSlow
-		e.allDepsStateCacheSlow = allDepsStateCacheSlow
-		e.allDepsCompleted = allDepsCompleted
+		e.allDepsCompletedCacheFast = e.cacheMapDone && allDepsCompletedCacheFast
+		e.allDepsCompletedCacheSlow = e.cacheMapDone && allDepsCompletedCacheSlow
+		e.allDepsStateCacheSlow = e.cacheMapDone && allDepsStateCacheSlow
+		e.allDepsCompleted = e.cacheMapDone && allDepsCompleted
 	}
 }
 
@@ -784,7 +802,7 @@ func (e *edge) loadCache(ctx context.Context) (interface{}, error) {
 
 // execOp creates a request to execute the vertex operation
 func (e *edge) execOp(ctx context.Context) (interface{}, error) {
-	cacheKey, inputs := e.commitOptions()
+	cacheKeys, inputs := e.commitOptions()
 	results, subExporters, err := e.op.Exec(ctx, toResultSlice(inputs))
 	if err != nil {
 		return nil, err
@@ -802,21 +820,27 @@ func (e *edge) execOp(ctx context.Context) (interface{}, error) {
 			go results[i].Release(context.TODO())
 		}
 	}
-	ck, err := e.op.Cache().Save(cacheKey, res)
-	if err != nil {
-		return nil, err
-	}
 
-	exps := make([]Exporter, 0, len(subExporters))
-	for _, exp := range subExporters {
-		exps = append(exps, exp.Exporter)
-	}
+	var exporters []Exporter
 
-	if len(subExporters) > 0 {
-		ck = &ExportableCacheKey{
-			CacheKey: ck.CacheKey,
-			Exporter: &mergedExporter{exporters: append([]Exporter{ck.Exporter}, exps...)},
+	for _, cacheKey := range cacheKeys {
+		ck, err := e.op.Cache().Save(cacheKey, res)
+		if err != nil {
+			return nil, err
 		}
+
+		exps := make([]Exporter, 0, len(subExporters))
+		for _, exp := range subExporters {
+			exps = append(exps, exp.Exporter)
+		}
+
+		exporters = append(exporters, ck.Exporter)
+		exporters = append(exporters, exps...)
+	}
+
+	ck := &ExportableCacheKey{
+		CacheKey: cacheKeys[0],
+		Exporter: &mergedExporter{exporters: exporters},
 	}
 
 	return NewCachedResult(res, *ck), nil

solver-next/exporter.go

@@ -100,6 +100,7 @@ func (e *exporter) ExportTo(ctx context.Context, t ExporterTarget, converter fun
 	}
 
 	rec := t.Add(rootKey(e.k.Digest(), e.k.Output()))
+	allRec := []ExporterRecord{rec}
 
 	for i, srcs := range srcs {
 		for _, src := range srcs {
@@ -119,9 +120,24 @@ func (e *exporter) ExportTo(ctx context.Context, t ExporterTarget, converter fun
 
 	var remote *Remote
 
+	if v := e.record; v != nil && len(deps) == 0 {
+		cm := v.cacheManager
+		key := cm.getID(v.key)
+		if err := cm.backend.WalkIDsByResult(v.ID, func(id string) error {
+			if id == key {
+				return nil
+			}
+			allRec = append(allRec, t.Add(digest.Digest(id)))
+			return nil
+		}); err != nil {
+			return nil, err
+		}
+	}
+
 	if v := e.record; v != nil && len(e.k.Deps()) > 0 {
 		cm := v.cacheManager
-		res, err := cm.backend.Load(cm.getID(v.key), v.ID)
+		key := cm.getID(v.key)
+		res, err := cm.backend.Load(key, v.ID)
 		if err != nil {
 			return nil, err
 		}
@@ -144,10 +160,12 @@ func (e *exporter) ExportTo(ctx context.Context, t ExporterTarget, converter fun
 		}
 
 		if remote != nil {
-			rec.AddResult(v.CreatedAt, remote)
+			for _, rec := range allRec {
+				rec.AddResult(v.CreatedAt, remote)
+			}
 		}
 	}
-	e.res = []ExporterRecord{rec}
+	e.res = allRec
 	t.Visit(e)
 
 	return e.res, nil

solver-next/jobs.go

@@ -446,8 +446,13 @@ func (j *Job) Call(ctx context.Context, name string, fn func(ctx context.Context
 	return inVertexContext(ctx, name, fn)
 }
 
+type cacheMapResp struct {
+	*CacheMap
+	complete bool
+}
+
 type activeOp interface {
-	CacheMap(context.Context) (*CacheMap, error)
+	CacheMap(context.Context, int) (*cacheMapResp, error)
 	LoadCache(ctx context.Context, rec *CacheRecord) (Result, error)
 	Exec(ctx context.Context, inputs []Result) (outputs []Result, exporters []ExportableCacheKey, err error)
 	IgnoreCache() bool
@@ -483,8 +488,9 @@ type sharedOp struct {
 	execRes *execRes
 	execErr error
 
-	cacheRes *CacheMap
-	cacheErr error
+	cacheRes  []*CacheMap
+	cacheDone bool
+	cacheErr  error
 
 	slowMu       sync.Mutex
 	slowCacheRes map[Index]digest.Digest
@@ -553,13 +559,13 @@ func (s *sharedOp) CalcSlowCache(ctx context.Context, index Index, f ResultBased
 	return key.(digest.Digest), nil
 }
 
-func (s *sharedOp) CacheMap(ctx context.Context) (*CacheMap, error) {
+func (s *sharedOp) CacheMap(ctx context.Context, index int) (*cacheMapResp, error) {
 	op, err := s.getOp()
 	if err != nil {
 		return nil, err
 	}
 	res, err := s.g.Do(ctx, "cachemap", func(ctx context.Context) (ret interface{}, retErr error) {
-		if s.cacheRes != nil {
+		if s.cacheRes != nil && s.cacheDone || index < len(s.cacheRes) {
 			return s.cacheRes, nil
 		}
 		if s.cacheErr != nil {
@@ -576,7 +582,7 @@ func (s *sharedOp) CacheMap(ctx context.Context) (*CacheMap, error) {
 				notifyCompleted(ctx, &s.st.clientVertex, retErr, false)
 			}()
 		}
-		res, err := op.CacheMap(ctx)
+		res, done, err := op.CacheMap(ctx, len(s.cacheRes))
 		complete := true
 		if err != nil {
 			canceled := false
@@ -591,16 +597,22 @@ func (s *sharedOp) CacheMap(ctx context.Context) (*CacheMap, error) {
 		}
 		if complete {
 			if err == nil {
-				s.cacheRes = res
+				s.cacheRes = append(s.cacheRes, res)
+				s.cacheDone = done
 			}
 			s.cacheErr = err
 		}
-		return res, err
+		return s.cacheRes, err
 	})
 	if err != nil {
 		return nil, err
 	}
-	return res.(*CacheMap), nil
+
+	if len(res.([]*CacheMap)) <= index {
+		return s.CacheMap(ctx, index)
+	}
+
+	return &cacheMapResp{CacheMap: res.([]*CacheMap)[index], complete: s.cacheDone}, nil
 }
 
 func (s *sharedOp) Exec(ctx context.Context, inputs []Result) (outputs []Result, exporters []ExportableCacheKey, err error) {

solver-next/memorycachestorage.go

@@ -116,6 +116,24 @@ func (s *inMemoryStore) AddResult(id string, res CacheResult) error {
 	return nil
 }
 
+func (s *inMemoryStore) WalkIDsByResult(resultID string, fn func(string) error) error {
+	s.mu.Lock()
+
+	ids := map[string]struct{}{}
+	for id := range s.byResult[resultID] {
+		ids[id] = struct{}{}
+	}
+	s.mu.Unlock()
+
+	for id := range ids {
+		if err := fn(id); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func (s *inMemoryStore) Release(resultID string) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()

solver-next/scheduler_test.go

@@ -2264,6 +2264,125 @@ func TestSlowCacheAvoidAccess(t *testing.T) {
 	require.Equal(t, int64(1), cacheManager.loadCounter)
 }
 
+func TestCacheMultipleMaps(t *testing.T) {
+	t.Parallel()
+	ctx := context.TODO()
+
+	cacheManager := newTrackingCacheManager(NewInMemoryCacheManager())
+
+	l := NewJobList(SolverOpt{
+		ResolveOpFunc: testOpResolver,
+		DefaultCache:  cacheManager,
+	})
+	defer l.Close()
+
+	j0, err := l.NewJob("j0")
+	require.NoError(t, err)
+
+	defer func() {
+		if j0 != nil {
+			j0.Discard()
+		}
+	}()
+
+	g0 := Edge{
+		Vertex: vtx(vtxOpt{
+			name:         "v0",
+			cacheKeySeed: "seed0",
+			cacheKeySeeds: []func() string{
+				func() string { return "seed1" },
+				func() string { return "seed2" },
+			},
+			value: "result0",
+		}),
+	}
+	res, err := j0.Build(ctx, g0)
+	require.NoError(t, err)
+	require.Equal(t, unwrap(res), "result0")
+
+	require.NoError(t, j0.Discard())
+	j0 = nil
+
+	expTarget := newTestExporterTarget()
+
+	_, err = res.CacheKey().Exporter.ExportTo(ctx, expTarget, testConvertToRemote)
+	require.NoError(t, err)
+
+	expTarget.normalize()
+	require.Equal(t, len(expTarget.records), 3)
+
+	j1, err := l.NewJob("j1")
+	require.NoError(t, err)
+
+	defer func() {
+		if j1 != nil {
+			j1.Discard()
+		}
+	}()
+
+	called := false
+	g1 := Edge{
+		Vertex: vtx(vtxOpt{
+			name:         "v1",
+			cacheKeySeed: "seed1",
+			cacheKeySeeds: []func() string{
+				func() string { called = true; return "seed3" },
+			},
+			value: "result0-not-cached",
+		}),
+	}
+
+	res, err = j1.Build(ctx, g1)
+	require.NoError(t, err)
+	require.Equal(t, unwrap(res), "result0")
+
+	require.NoError(t, j1.Discard())
+	j1 = nil
+
+	expTarget = newTestExporterTarget()
+
+	_, err = res.CacheKey().Exporter.ExportTo(ctx, expTarget, testConvertToRemote)
+	require.NoError(t, err)
+
+	require.Equal(t, len(expTarget.records), 3)
+	require.Equal(t, called, false)
+
+	j2, err := l.NewJob("j2")
+	require.NoError(t, err)
+
+	defer func() {
+		if j2 != nil {
+			j2.Discard()
+		}
+	}()
+
+	g2 := Edge{
+		Vertex: vtx(vtxOpt{
+			name:         "v2",
+			cacheKeySeed: "seed3",
+			cacheKeySeeds: []func() string{
+				func() string { called = true; return "seed2" },
+			},
+			value: "result0-not-cached",
+		}),
+	}
+
+	res, err = j2.Build(ctx, g2)
+	require.NoError(t, err)
+	require.Equal(t, unwrap(res), "result0")
+
+	require.NoError(t, j2.Discard())
+	j2 = nil
+
+	expTarget = newTestExporterTarget()
+
+	_, err = res.CacheKey().Exporter.ExportTo(ctx, expTarget, testConvertToRemote)
+	require.NoError(t, err)
+
+	require.Equal(t, len(expTarget.records), 3)
+	require.Equal(t, called, true)
+}
+
 func TestCacheExportingPartialSelector(t *testing.T) {
 	t.Parallel()
 	ctx := context.TODO()
@@ -2589,6 +2708,7 @@ func generateSubGraph(nodes int) (Edge, int) {
 type vtxOpt struct {
 	name             string
 	cacheKeySeed     string
+	cacheKeySeeds    []func() string
 	execDelay        time.Duration
 	cacheDelay       time.Duration
 	cachePreFunc     func(context.Context) error
@@ -2686,11 +2806,16 @@ func (v *vertex) cacheMap(ctx context.Context) error {
 	return nil
 }
 
-func (v *vertex) CacheMap(ctx context.Context) (*CacheMap, error) {
-	if err := v.cacheMap(ctx); err != nil {
-		return nil, err
+func (v *vertex) CacheMap(ctx context.Context, index int) (*CacheMap, bool, error) {
+	if index == 0 {
+		if err := v.cacheMap(ctx); err != nil {
+			return nil, false, err
+		}
+		return v.makeCacheMap(), len(v.opt.cacheKeySeeds) == index, nil
 	}
-	return v.makeCacheMap(), nil
+	return &CacheMap{
+		Digest: digest.FromBytes([]byte(fmt.Sprintf("seed:%s", v.opt.cacheKeySeeds[index-1]()))),
+	}, len(v.opt.cacheKeySeeds) == index, nil
 }
 
 func (v *vertex) exec(ctx context.Context, inputs []Result) error {

solver-next/types.go

@@ -89,8 +89,9 @@ type CacheLink struct {
 
 // Op is an implementation for running a vertex
 type Op interface {
-	// CacheMap returns structure describing how the operation is cached
-	CacheMap(context.Context) (*CacheMap, error)
+	// CacheMap returns structure describing how the operation is cached.
+	// Currently only roots are allowed to return multiple cache maps per op.
+	CacheMap(context.Context, int) (*CacheMap, bool, error)
 	// Exec runs an operation given results from previous operations.
 	Exec(ctx context.Context, inputs []Result) (outputs []Result, err error)
 }