cache: improve diff ref release logic.

Before this change, the lower and upper parents provided to the cache
manager Diff method were not cloned, which resulted in some code paths
incorrectly providing them directly as the parents of the returned ref.
This meant that if they were released after the call to Diff, the diff
ref could become incorrectly invalidated.

Now, the lower and upper are cloned and unit test coverage has been
added to test that ref release is handled correctly.

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

cache/manager.go

@@ -694,6 +694,7 @@ func (cm *cacheManager) Merge(ctx context.Context, inputParents []ImmutableRef,
 			parent = p.(*immutableRef)
 			defer parent.Release(context.TODO())
 		}
+		// On success, cloned parents will be not be released and will be owned by the returned ref
 		switch parent.kind() {
 		case Merge:
 			// if parent is itself a merge, flatten it out by just setting our parents directly to its parents
@@ -708,9 +709,9 @@ func (cm *cacheManager) Merge(ctx context.Context, inputParents []ImmutableRef,
 		}
 	}
 
+	// On success, createMergeRef takes ownership of parents
 	mergeRef, err := cm.createMergeRef(ctx, parents, dhs, opts...)
 	if err != nil {
-		parents.release(context.TODO())
 		return nil, err
 	}
 	return mergeRef, nil
@@ -821,10 +822,11 @@ func (cm *cacheManager) Diff(ctx context.Context, lower, upper ImmutableRef, opt
 			parent = p.(*immutableRef)
 			defer parent.Release(context.TODO())
 		}
+		// On success, cloned parents will not be released and will be owned by the returned ref
 		if i == 0 {
-			dps.lower = parent
+			dps.lower = parent.clone()
 		} else {
-			dps.upper = parent
+			dps.upper = parent.clone()
 		}
 		for dgst, handler := range parent.descHandlers {
 			dhs[dgst] = handler
@@ -861,6 +863,7 @@ func (cm *cacheManager) Diff(ctx context.Context, lower, upper ImmutableRef, opt
 			for i := len(lowerLayers); i < len(upperLayers); i++ {
 				subUpper := upperLayers[i]
 				subLower := subUpper.layerParent
+				// On success, cloned refs will not be released and will be owned by the returned ref
 				if subLower == nil {
 					mergeParents.mergeParents[i-len(lowerLayers)] = subUpper.clone()
 				} else {
@@ -874,6 +877,7 @@ func (cm *cacheManager) Diff(ctx context.Context, lower, upper ImmutableRef, opt
 					mergeParents.mergeParents[i-len(lowerLayers)] = diffRef
 				}
 			}
+			// On success, createMergeRef takes ownership of mergeParents
 			mergeRef, err := cm.createMergeRef(ctx, mergeParents, dhs)
 			if err != nil {
 				return nil, err
@@ -883,9 +887,9 @@ func (cm *cacheManager) Diff(ctx context.Context, lower, upper ImmutableRef, opt
 		}
 	}
 
+	// On success, createDiffRef takes ownership of parents
 	diffRef, err := cm.createDiffRef(ctx, parents, dhs, opts...)
 	if err != nil {
-		parents.release(context.TODO())
 		return nil, err
 	}
 	return diffRef, nil

cache/manager_test.go

@@ -1684,6 +1684,98 @@ func TestMergeOp(t *testing.T) {
 	checkDiskUsage(ctx, t, cm, 0, 0)
 }
 
+func TestDiffOp(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("Depends on unimplemented diff-op support on Windows")
+	}
+
+	// This just tests the basic Diff method and some of the logic with releasing diff refs.
+	// Tests for the fs diff logic are in client_test and snapshotter_test.
+	t.Parallel()
+
+	ctx := namespaces.WithNamespace(context.Background(), "buildkit-test")
+
+	tmpdir, err := ioutil.TempDir("", "cachemanager")
+	require.NoError(t, err)
+	defer os.RemoveAll(tmpdir)
+
+	snapshotter, err := native.NewSnapshotter(filepath.Join(tmpdir, "snapshots"))
+	require.NoError(t, err)
+
+	co, cleanup, err := newCacheManager(ctx, cmOpt{
+		snapshotter:     snapshotter,
+		snapshotterName: "native",
+	})
+	require.NoError(t, err)
+	defer cleanup()
+	cm := co.manager
+
+	newLower, err := cm.New(ctx, nil, nil)
+	require.NoError(t, err)
+	lowerA, err := newLower.Commit(ctx)
+	require.NoError(t, err)
+	newUpper, err := cm.New(ctx, nil, nil)
+	require.NoError(t, err)
+	upperA, err := newUpper.Commit(ctx)
+	require.NoError(t, err)
+
+	// verify that releasing parents does not invalidate a diff ref until it is released
+	diff, err := cm.Diff(ctx, lowerA, upperA)
+	require.NoError(t, err)
+	checkDiskUsage(ctx, t, cm, 3, 0)
+	require.NoError(t, lowerA.Release(ctx))
+	require.NoError(t, upperA.Release(ctx))
+	checkDiskUsage(ctx, t, cm, 3, 0)
+	require.NoError(t, cm.Prune(ctx, nil, client.PruneInfo{All: true}))
+	checkDiskUsage(ctx, t, cm, 3, 0)
+	_, err = diff.Mount(ctx, true, nil)
+	require.NoError(t, err)
+	require.NoError(t, diff.Release(ctx))
+	checkDiskUsage(ctx, t, cm, 0, 3)
+	require.NoError(t, cm.Prune(ctx, nil, client.PruneInfo{All: true}))
+	checkDiskUsage(ctx, t, cm, 0, 0)
+
+	// test "unmerge" diffs that are defined as a merge of single-layer diffs
+	newRef, err := cm.New(ctx, nil, nil)
+	require.NoError(t, err)
+	a, err := newRef.Commit(ctx)
+	require.NoError(t, err)
+	newRef, err = cm.New(ctx, a, nil)
+	require.NoError(t, err)
+	b, err := newRef.Commit(ctx)
+	require.NoError(t, err)
+	newRef, err = cm.New(ctx, b, nil)
+	require.NoError(t, err)
+	c, err := newRef.Commit(ctx)
+	require.NoError(t, err)
+	newRef, err = cm.New(ctx, c, nil)
+	require.NoError(t, err)
+	d, err := newRef.Commit(ctx)
+	require.NoError(t, err)
+	newRef, err = cm.New(ctx, d, nil)
+	require.NoError(t, err)
+	e, err := newRef.Commit(ctx)
+	require.NoError(t, err)
+
+	diff, err = cm.Diff(ctx, c, e)
+	require.NoError(t, err)
+	checkDiskUsage(ctx, t, cm, 8, 0) // 5 base refs + 2 diffs + 1 merge
+	require.NoError(t, a.Release(ctx))
+	require.NoError(t, b.Release(ctx))
+	require.NoError(t, c.Release(ctx))
+	require.NoError(t, d.Release(ctx))
+	require.NoError(t, e.Release(ctx))
+	checkDiskUsage(ctx, t, cm, 8, 0)
+	require.NoError(t, cm.Prune(ctx, nil, client.PruneInfo{All: true}))
+	checkDiskUsage(ctx, t, cm, 8, 0)
+	_, err = diff.Mount(ctx, true, nil)
+	require.NoError(t, err)
+	require.NoError(t, diff.Release(ctx))
+	checkDiskUsage(ctx, t, cm, 0, 8)
+	require.NoError(t, cm.Prune(ctx, nil, client.PruneInfo{All: true}))
+	checkDiskUsage(ctx, t, cm, 0, 0)
+}
+
 func TestLoadHalfFinalizedRef(t *testing.T) {
 	// This test simulates the situation where a ref w/ an equalMutable has its
 	// snapshot committed but there is a crash before the metadata is updated to