Rename, refactor, & basic caching.

solver/jobs.go

@@ -2,7 +2,10 @@ package solver
 
 import (
 	"context"
+	"crypto/sha256"
 	"fmt"
+	"hash"
+	"io"
 	"sync"
 	"time"
 
@@ -529,36 +532,48 @@ func (jl *Solver) deleteIfUnreferenced(k digest.Digest, st *state) {
 	}
 }
 
-type dumbBuilder struct {
+type simpleSolver struct {
 	resolveOpFunc ResolveOpFunc
 	solver        *Solver
 	job           *Job
 }
 
-var cache = map[string]CachedResult{}
+var cache = map[string]Result{}
+var cacheMaps = map[string]*simpleCacheMap{}
+var execInProgress = map[string]struct{}{}
 var mu = sync.Mutex{}
 
-var execInProgress = map[string]struct{}{}
+type cacheMapDep struct {
+	selector string
+	computed string
+}
 
-func (b *dumbBuilder) build(ctx context.Context, e Edge) (CachedResult, error) {
+type simpleCacheMap struct {
+	digest string
+	inputs []string
+	deps   []cacheMapDep
+}
+
+func (s *simpleSolver) build(ctx context.Context, e Edge) (CachedResult, error) {
 
 	// Ordered list of vertices to build.
-	digests, vertices := b.exploreVertices(e)
+	digests, vertices := s.exploreVertices(e)
 
-	var ret CachedResult
+	var ret Result
 
 	for _, d := range digests {
+		fmt.Println()
+
 		vertex, ok := vertices[d]
 		if !ok {
 			return nil, errors.Errorf("digest %s not found", d)
 		}
 
-		dgst := d.String()
 		mu.Lock()
 		// TODO: replace busy-wait loop with a wait-for-channel-to-close approach
 		for {
-			if _, shouldWait := execInProgress[dgst]; !shouldWait {
-				execInProgress[dgst] = struct{}{}
+			if _, shouldWait := execInProgress[d.String()]; !shouldWait {
+				execInProgress[d.String()] = struct{}{}
 				mu.Unlock()
 				break
 			}
@@ -567,97 +582,213 @@ func (b *dumbBuilder) build(ctx context.Context, e Edge) (CachedResult, error) {
 			mu.Lock()
 		}
 
-		defaultCache := NewInMemoryCacheManager()
+		st := s.createState(vertex)
 
-		st := &state{
-			opts:         SolverOpt{DefaultCache: defaultCache, ResolveOpFunc: b.resolveOpFunc},
-			parents:      map[digest.Digest]struct{}{},
-			childVtx:     map[digest.Digest]struct{}{},
-			allPw:        map[progress.Writer]struct{}{},
-			mpw:          progress.NewMultiWriter(progress.WithMetadata("vertex", d)),
-			mspan:        tracing.NewMultiSpan(),
-			vtx:          vertex,
-			clientVertex: initClientVertex(vertex),
-			edges:        map[Index]*edge{},
-			index:        b.solver.index,
-			mainCache:    defaultCache,
-			cache:        map[string]CacheManager{},
-			solver:       b.solver,
-			origDigest:   vertex.Digest(),
-		}
+		notifyCompleted := notifyStarted(ctx, &st.clientVertex, true)
 
-		st.jobs = map[*Job]struct{}{
-			b.job: {},
-		}
+		op := newSharedOp(st.opts.ResolveOpFunc, st.opts.DefaultCache, st)
 
-		st.mpw.Add(b.job.pw)
-
-		//fmt.Println("Processing vertex", vertex.Name(), d.String())
+		// CacheMap populates required fields in SourceOp.
+		cm, err := op.CacheMap(ctx, int(e.Index))
+		if err != nil {
+			return nil, err
+		}
 
-		edge := st.getEdge(e.Index)
+		fmt.Println(vertex.Name())
+		fmt.Println("LLB digest:", d.String())
+		fmt.Println("CacheMap digest:", cm.Digest)
 
-		edge.deps = make([]*dep, 0, len(vertex.Inputs()))
-		inputs := vertex.Inputs()
-		for i := range inputs {
-			dep := newDep(Index(i))
-			if v, ok := cache[inputs[i].Vertex.Digest().String()]; ok {
-				dep.result = NewSharedCachedResult(v)
-			}
-			edge.deps = append(edge.deps, dep)
+		inputs, err := s.preprocessInputs(ctx, st, vertex, cm.CacheMap)
+		if err != nil {
+			return nil, err
 		}
 
-		cm, err := edge.op.CacheMap(ctx, int(e.Index))
+		cacheKey, err := s.cacheKey(ctx, d.String())
 		if err != nil {
 			return nil, err
 		}
 
-		edge.cacheMap = cm.CacheMap
+		fmt.Println("Computed cache key:", cacheKey)
 
-		if r, ok := cache[dgst]; ok {
-			st.clientVertex.Cached = true
-			st.mpw.Write(identity.NewID(), st.clientVertex)
-			ret = r
-			mu.Lock()
-			delete(execInProgress, dgst)
+		mu.Lock()
+		if v, ok := cache[cacheKey]; ok && v != nil {
+			fmt.Println("Cache hit!")
+			delete(execInProgress, d.String())
 			mu.Unlock()
+			notifyCompleted(nil, true)
+			ret = v
 			continue
 		}
+		mu.Unlock()
 
-		res, err := edge.execOp(ctx)
+		results, _, err := op.Exec(ctx, inputs)
 		if err != nil {
 			mu.Lock()
-			delete(execInProgress, dgst)
+			delete(execInProgress, d.String())
 			mu.Unlock()
+			notifyCompleted(err, false)
 			return nil, err
 		}
 
-		cachedResult := res.(CachedResult)
+		notifyCompleted(nil, false)
+
+		res := results[int(e.Index)]
+		ret = res
+
+		mu.Lock()
+		delete(execInProgress, d.String())
+		cache[cacheKey] = res
+		mu.Unlock()
+	}
+
+	return NewCachedResult(ret, []ExportableCacheKey{}), nil
+}
+
+func (s *simpleSolver) createState(vertex Vertex) *state {
+	defaultCache := NewInMemoryCacheManager()
+
+	st := &state{
+		opts:         SolverOpt{DefaultCache: defaultCache, ResolveOpFunc: s.resolveOpFunc},
+		parents:      map[digest.Digest]struct{}{},
+		childVtx:     map[digest.Digest]struct{}{},
+		allPw:        map[progress.Writer]struct{}{},
+		mpw:          progress.NewMultiWriter(progress.WithMetadata("vertex", vertex.Digest())),
+		mspan:        tracing.NewMultiSpan(),
+		vtx:          vertex,
+		clientVertex: initClientVertex(vertex),
+		edges:        map[Index]*edge{},
+		index:        s.solver.index,
+		mainCache:    defaultCache,
+		cache:        map[string]CacheManager{},
+		solver:       s.solver,
+		origDigest:   vertex.Digest(),
+	}
+
+	st.jobs = map[*Job]struct{}{
+		s.job: {},
+	}
+
+	st.mpw.Add(s.job.pw)
 
-		st.mpw.Write(identity.NewID(), st.clientVertex)
+	return st
+}
+
+func (s *simpleSolver) preprocessInputs(ctx context.Context, st *state, vertex Vertex, cm *CacheMap) ([]Result, error) {
+	// This struct is used to reconstruct a cache key from an LLB digest & all
+	// parents using consistent digests that depend on the full dependency chain.
+	// TODO: handle cm.Opts (CacheOpts)?
+	scm := simpleCacheMap{
+		digest: cm.Digest.String(),
+		deps:   make([]cacheMapDep, len(cm.Deps)),
+		inputs: make([]string, len(cm.Deps)),
+	}
 
-		edge.result = NewSharedCachedResult(cachedResult)
-		edge.state = edgeStatusComplete
+	var inputs []Result
 
-		ret = cachedResult
+	for i, in := range vertex.Inputs() {
+		// Compute a cache key given the LLB digest value.
+		cacheKey, err := s.cacheKey(ctx, in.Vertex.Digest().String())
+		if err != nil {
+			return nil, err
+		}
 
+		// Lookup the result for that cache key.
 		mu.Lock()
-		cache[dgst] = cachedResult
-		delete(execInProgress, dgst)
+		res, ok := cache[cacheKey]
 		mu.Unlock()
+		if !ok {
+			return nil, errors.Errorf("cache key not found: %s", cacheKey)
+		}
+
+		dep := cm.Deps[i]
+
+		// Unlazy the result.
+		if dep.PreprocessFunc != nil {
+			err = dep.PreprocessFunc(ctx, res, st)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		// Add selectors (usually file references) to the struct.
+		scm.deps[i] = cacheMapDep{
+			selector: dep.Selector.String(),
+		}
+
+		// ComputeDigestFunc will usually checksum files. This is then used as
+		// part of the cache key to ensure it's consistent & distinct for this
+		// operation.
+		if dep.ComputeDigestFunc != nil {
+			compDigest, err := dep.ComputeDigestFunc(ctx, res, st)
+			if err != nil {
+				return nil, err
+			}
+			scm.deps[i].computed = compDigest.String()
+		}
+
+		// Add input references to the struct as to link dependencies.
+		scm.inputs[i] = in.Vertex.Digest().String()
+
+		// Add the cached result to the input set. These inputs are used to
+		// reconstruct dependencies (mounts, etc.) for a new container run.
+		inputs = append(inputs, res)
 	}
 
-	return ret, nil
+	mu.Lock()
+	cacheMaps[vertex.Digest().String()] = &scm
+	mu.Unlock()
+
+	return inputs, nil
+}
+
+func (s *simpleSolver) cacheKey(ctx context.Context, d string) (string, error) {
+	h := sha256.New()
+
+	err := s.calcCacheKey(ctx, d, h)
+	if err != nil {
+		return "", err
+	}
+
+	return fmt.Sprintf("%x", h.Sum(nil)), nil
+}
+
+func (s *simpleSolver) calcCacheKey(ctx context.Context, d string, h hash.Hash) error {
+	mu.Lock()
+	c, ok := cacheMaps[d]
+	mu.Unlock()
+	if !ok {
+		return errors.New("missing cache map key")
+	}
+
+	for _, in := range c.inputs {
+		err := s.calcCacheKey(ctx, in, h)
+		if err != nil {
+			return err
+		}
+	}
+
+	io.WriteString(h, c.digest)
+	for _, dep := range c.deps {
+		if dep.selector != "" {
+			io.WriteString(h, dep.selector)
+		}
+		if dep.computed != "" {
+			io.WriteString(h, dep.computed)
+		}
+	}
+
+	return nil
 }
 
-func (b *dumbBuilder) exploreVertices(e Edge) ([]digest.Digest, map[digest.Digest]Vertex) {
+func (s *simpleSolver) exploreVertices(e Edge) ([]digest.Digest, map[digest.Digest]Vertex) {
 
 	digests := []digest.Digest{e.Vertex.Digest()}
 	vertices := map[digest.Digest]Vertex{
 		e.Vertex.Digest(): e.Vertex,
 	}
 
 	for _, edge := range e.Vertex.Inputs() {
-		d, v := b.exploreVertices(edge)
+		d, v := s.exploreVertices(edge)
 		digests = append(d, digests...)
 		for key, value := range v {
 			vertices[key] = value
@@ -681,7 +812,7 @@ func (j *Job) Build(ctx context.Context, e Edge) (CachedResultWithProvenance, er
 		j.span = span
 	}
 
-	b := &dumbBuilder{
+	b := &simpleSolver{
 		resolveOpFunc: j.list.opts.ResolveOpFunc,
 		solver:        j.list,
 		job:           j,