Merge pull request #7332 from tonistiigi/fix-commit-cancel

cache: avoid losing finalize commit on cancel

cache/manager_test.go

@@ -3015,6 +3015,108 @@ func TestLoadHalfFinalizedRef(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestFinalizeAlreadyCommitted(t *testing.T) {
+	// This test simulates a snapshotter commit that was applied but reported
+	// as failed (e.g. a canceled containerd RPC). Finalize must notice the
+	// committed snapshot instead of failing on the missing active snapshot.
+	t.Parallel()
+
+	ctx := namespaces.WithNamespace(t.Context(), "buildkit-test")
+
+	tmpdir := t.TempDir()
+
+	snapshotter, err := native.NewSnapshotter(filepath.Join(tmpdir, "snapshots"))
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, snapshotter.Close())
+	})
+
+	co, cleanup, err := newCacheManager(ctx, t, cmOpt{
+		tmpdir:          tmpdir,
+		snapshotter:     snapshotter,
+		snapshotterName: "native",
+	})
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+	cm := co.manager.(*cacheManager)
+
+	mref, err := cm.New(ctx, nil, nil, CachePolicyRetain)
+	require.NoError(t, err)
+	mutRef := mref.(*mutableRef)
+
+	iref, err := mutRef.Commit(ctx)
+	require.NoError(t, err)
+	immutRef := iref.(*immutableRef)
+
+	err = cm.Snapshotter.Commit(ctx, immutRef.getSnapshotID(), mutRef.getSnapshotID())
+	require.NoError(t, err)
+
+	require.NoError(t, iref.Finalize(ctx))
+	require.Nil(t, immutRef.equalMutable)
+
+	// the committed snapshot must stay protected by the ref's lease
+	resources, err := co.lm.ListResources(ctx, leases.Lease{ID: immutRef.ID()})
+	require.NoError(t, err)
+	require.Contains(t, resources, leases.Resource{
+		ID:   immutRef.getSnapshotID(),
+		Type: "snapshots/" + cm.Snapshotter.Name(),
+	})
+
+	m, err := iref.Mount(ctx, true, nil)
+	require.NoError(t, err)
+	mnts, release, err := m.Mount()
+	require.NoError(t, err)
+	require.NotEmpty(t, mnts)
+	require.NoError(t, release())
+
+	require.NoError(t, mref.Release(ctx))
+	require.NoError(t, iref.Release(ctx))
+}
+
+func TestFinalizeCanceledContext(t *testing.T) {
+	// Finalize ignores cancellation, as a remote snapshotter may apply a commit
+	// that the caller observes as canceled.
+	t.Parallel()
+
+	ctx := namespaces.WithNamespace(t.Context(), "buildkit-test")
+
+	tmpdir := t.TempDir()
+
+	snapshotter, err := native.NewSnapshotter(filepath.Join(tmpdir, "snapshots"))
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, snapshotter.Close())
+	})
+
+	co, cleanup, err := newCacheManager(ctx, t, cmOpt{
+		tmpdir:          tmpdir,
+		snapshotter:     snapshotter,
+		snapshotterName: "native",
+	})
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+	cm := co.manager.(*cacheManager)
+
+	mref, err := cm.New(ctx, nil, nil, CachePolicyRetain)
+	require.NoError(t, err)
+
+	iref, err := mref.Commit(ctx)
+	require.NoError(t, err)
+	immutRef := iref.(*immutableRef)
+
+	cctx, cancel := context.WithCancelCause(ctx)
+	cancel(errors.WithStack(context.Canceled))
+	require.NoError(t, iref.Finalize(cctx))
+	require.Nil(t, immutRef.equalMutable)
+
+	info, err := cm.Snapshotter.Stat(ctx, immutRef.getSnapshotID())
+	require.NoError(t, err)
+	require.Equal(t, snapshots.KindCommitted, info.Kind)
+
+	require.NoError(t, mref.Release(ctx))
+	require.NoError(t, iref.Release(ctx))
+}
+
 func TestMountReadOnly(t *testing.T) {
 	t.Parallel()
 	if runtime.GOOS != "linux" {

cache/refs.go

@@ -1538,6 +1538,11 @@ func (cr *cacheRecord) finalize(ctx context.Context) error {
 		return nil
 	}
 
+	// The snapshotter may apply the commit even when the caller sees a
+	// cancellation (e.g. a containerd RPC canceled in flight), leaving the
+	// mutable snapshot gone while the record still points to it.
+	ctx = context.WithoutCancel(ctx)
+
 	_, err := cr.cm.LeaseManager.Create(ctx, func(l *leases.Lease) error {
 		l.ID = cr.ID()
 		l.Labels = map[string]string{
@@ -1555,13 +1560,23 @@ func (cr *cacheRecord) finalize(ctx context.Context) error {
 		ID:   cr.getSnapshotID(),
 		Type: "snapshots/" + cr.cm.Snapshotter.Name(),
 	}); err != nil {
-		_ = cr.cm.LeaseManager.Delete(context.WithoutCancel(ctx), leases.Lease{ID: cr.ID()})
+		_ = cr.cm.LeaseManager.Delete(ctx, leases.Lease{ID: cr.ID()})
 		return errors.Wrapf(err, "failed to add snapshot %s to lease", cr.getSnapshotID())
 	}
 
 	if err := cr.cm.Snapshotter.Commit(ctx, cr.getSnapshotID(), mutable.getSnapshotID()); err != nil {
-		_ = cr.cm.LeaseManager.Delete(context.WithoutCancel(ctx), leases.Lease{ID: cr.ID()})
-		return errors.Wrapf(err, "failed to commit %s to %s during finalize", mutable.getSnapshotID(), cr.getSnapshotID())
+		// the commit may still have been applied, e.g. if the response was lost
+		info, statErr := cr.cm.Snapshotter.Stat(ctx, cr.getSnapshotID())
+		if statErr != nil || info.Kind != snapshots.KindCommitted {
+			// keep the lease if it is unknown whether the commit was applied
+			if statErr == nil || errors.Is(statErr, cerrdefs.ErrNotFound) {
+				_ = cr.cm.LeaseManager.Delete(ctx, leases.Lease{ID: cr.ID()})
+			} else {
+				bklog.G(ctx).WithError(statErr).Warnf("failed to check snapshot %s after commit error, keeping lease", cr.getSnapshotID())
+			}
+			return errors.Wrapf(err, "failed to commit %s to %s during finalize", mutable.getSnapshotID(), cr.getSnapshotID())
+		}
+		bklog.G(ctx).WithError(err).Warnf("snapshot %s was committed to %s despite commit error", mutable.getSnapshotID(), cr.getSnapshotID())
 	}
 	cr.mountCache = nil
 
@@ -1569,9 +1584,8 @@ func (cr *cacheRecord) finalize(ctx context.Context) error {
 	go func() {
 		cr.cm.mu.Lock()
 		defer cr.cm.mu.Unlock()
-		cleanupCtx := context.WithoutCancel(ctx)
-		if err := mutable.remove(cleanupCtx, true); err != nil {
-			bklog.G(cleanupCtx).Error(err)
+		if err := mutable.remove(ctx, true); err != nil {
+			bklog.G(ctx).Error(err)
 		}
 	}()