solver: avoid removing references required for future cache

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

cache/manager.go

@@ -105,6 +105,9 @@ func (cm *cacheManager) get(ctx context.Context, id string) (ImmutableRef, error
 		if len(rec.refs) != 0 {
 			return nil, errors.Wrapf(errLocked, "%s is locked", id)
 		}
+		if rec.equalImmutable != nil {
+			return rec.equalImmutable.ref(), nil
+		}
 		return rec.mref().commit(ctx)
 	}
 
@@ -336,15 +339,19 @@ func IsLocked(err error) bool {
 	return errors.Cause(err) == errLocked
 }
 
-type RefOption func(*cacheRecord) error
+type RefOption func(withMetadata) error
 
 type cachePolicy int
 
 const (
 	cachePolicyDefault cachePolicy = iota
-	cachePolicyKeepMutable
+	cachePolicyRetain
 )
 
-func CachePolicyKeepMutable(cr *cacheRecord) error {
-	return setCachePolicy(cr.md, cachePolicyKeepMutable)
+type withMetadata interface {
+	Metadata() *metadata.StorageItem
+}
+
+func CachePolicyRetain(m withMetadata) error {
+	return setCachePolicy(m.Metadata(), cachePolicyRetain)
 }

cache/manager_test.go

@@ -30,7 +30,7 @@ func TestManager(t *testing.T) {
 
 	checkDiskUsage(t, ctx, cm, 0, 0)
 
-	active, err := cm.New(ctx, nil, CachePolicyKeepMutable)
+	active, err := cm.New(ctx, nil, CachePolicyRetain)
 	require.NoError(t, err)
 
 	m, err := active.Mount(ctx, false)
@@ -102,7 +102,7 @@ func TestManager(t *testing.T) {
 	err = snap.Release(ctx)
 	require.NoError(t, err)
 
-	active2, err := cm.New(ctx, snap2, CachePolicyKeepMutable)
+	active2, err := cm.New(ctx, snap2, CachePolicyRetain)
 	require.NoError(t, err)
 
 	checkDiskUsage(t, ctx, cm, 2, 0)
@@ -133,7 +133,7 @@ func TestLazyCommit(t *testing.T) {
 
 	cm := getCacheManager(t, tmpdir)
 
-	active, err := cm.New(ctx, nil, CachePolicyKeepMutable)
+	active, err := cm.New(ctx, nil, CachePolicyRetain)
 	require.NoError(t, err)
 
 	// after commit mutable is locked
@@ -202,7 +202,7 @@ func TestLazyCommit(t *testing.T) {
 	require.NoError(t, err)
 
 	// test restarting after commit
-	active, err = cm.New(ctx, nil, CachePolicyKeepMutable)
+	active, err = cm.New(ctx, nil, CachePolicyRetain)
 	require.NoError(t, err)
 
 	// after commit mutable is locked

cache/refs.go

@@ -18,6 +18,7 @@ type ImmutableRef interface {
 	Size(ctx context.Context) (int64, error)
 	Parent() ImmutableRef
 	Finalize(ctx context.Context) error // Make sure reference is flushed to driver
+	Metadata() *metadata.StorageItem
 	// Prepare() / ChainID() / Meta()
 }
 
@@ -202,6 +203,10 @@ func (sr *immutableRef) Finalize(ctx context.Context) error {
 	return sr.finalize(ctx)
 }
 
+func (sr *cacheRecord) Metadata() *metadata.StorageItem {
+	return sr.md
+}
+
 func (sr *cacheRecord) finalize(ctx context.Context) error {
 	mutable := sr.equalMutable
 	if mutable == nil {
@@ -274,8 +279,11 @@ func (sr *mutableRef) Release(ctx context.Context) error {
 
 func (sr *mutableRef) release(ctx context.Context) error {
 	delete(sr.refs, sr)
-	if getCachePolicy(sr.md) != cachePolicyKeepMutable {
+	if getCachePolicy(sr.md) != cachePolicyRetain {
 		if sr.equalImmutable != nil {
+			if getCachePolicy(sr.equalImmutable.md) == cachePolicyRetain {
+				return nil
+			}
 			if err := sr.equalImmutable.remove(ctx, false); err != nil {
 				return err
 			}

solver/solver.go

@@ -256,6 +256,12 @@ func (s *Solver) getRefs(ctx context.Context, j *job, g *vertex) (retRef []Refer
 					for _, r := range r {
 						refs[i] = append(refs[i], newSharedRef(r))
 					}
+					if ref, ok := toImmutableRef(r[index].(Reference)); ok {
+						// make sure input that is required by next step does not get released in case build is cancelled
+						if err := cache.CachePolicyRetain(ref); err != nil {
+							return err
+						}
+					}
 					return nil
 				})
 			}(i, in.vertex, in.index)

source/git/gitsource.go

@@ -78,7 +78,7 @@ func (gs *gitSource) mountRemote(ctx context.Context, remote string) (target str
 
 	initializeRepo := false
 	if remoteRef == nil {
-		remoteRef, err = gs.cache.New(ctx, nil, cache.CachePolicyKeepMutable)
+		remoteRef, err = gs.cache.New(ctx, nil, cache.CachePolicyRetain)
 		if err != nil {
 			return "", nil, errors.Wrapf(err, "failed to create new mutable for %s", remote)
 		}

source/local/local.go

@@ -104,7 +104,7 @@ func (ls *localSourceHandler) Snapshot(ctx context.Context) (out cache.Immutable
 	}
 
 	if mutable == nil {
-		m, err := ls.cm.New(ctx, nil, cache.CachePolicyKeepMutable)
+		m, err := ls.cm.New(ctx, nil, cache.CachePolicyRetain)
 		if err != nil {
 			return nil, err
 		}