fix(orchestrator): clean up newly created cache files on errors

GitOrigin-RevId: d66d85bebb19857ac5e1979c2a286c53991f50ff

Author: snajpa

packages/orchestrator/pkg/sandbox/block/cache.go

@@ -59,13 +59,26 @@ type Cache struct {
 	closed    atomic.Bool
 }
 
-func NewCache(size, blockSize int64, filePath string, dirtyFile bool) (*Cache, error) {
-	f, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE, 0o644)
+func NewCache(size, blockSize int64, filePath string, dirtyFile bool) (_ *Cache, retErr error) {
+	f, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o644)
+	created := err == nil
+	if errors.Is(err, os.ErrExist) {
+		// Existing paths may hold a completed diff. Ownership transfers only
+		// on success, so a failed constructor must leave them to their caller.
+		f, err = os.OpenFile(filePath, os.O_RDWR|os.O_CREATE, 0o644)
+	}
 	if err != nil {
 		return nil, fmt.Errorf("error opening file: %w", err)
 	}
 
-	defer f.Close()
+	defer func() {
+		_ = f.Close()
+		if retErr != nil && created {
+			if err := os.Remove(filePath); err != nil && !errors.Is(err, os.ErrNotExist) {
+				retErr = errors.Join(retErr, fmt.Errorf("remove failed cache file: %w", err))
+			}
+		}
+	}()
 
 	if size == 0 {
 		return &Cache{

packages/orchestrator/pkg/sandbox/block/cache_construction_failure_test.go

@@ -0,0 +1,75 @@
+//go:build linux
+
+package block
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sys/unix"
+)
+
+func TestCacheConstructionFailureRemovesNewOutput(t *testing.T) {
+	t.Parallel()
+	path := filepath.Join(t.TempDir(), "new-output")
+	cache, err := NewCache(-1, int64(os.Getpagesize()), path, false)
+	require.Nil(t, cache)
+	require.ErrorIs(t, err, unix.EINVAL, "the real truncate failure must remain visible")
+	_, statErr := os.Stat(path)
+	require.ErrorIs(t, statErr, os.ErrNotExist, "failed construction must remove its newly created output")
+}
+
+func TestCacheConstructionFailurePreservesExistingOutput(t *testing.T) {
+	t.Parallel()
+	for _, mode := range []string{"regular", "symlink", "dangling-symlink"} {
+		t.Run(mode, func(t *testing.T) {
+			t.Parallel()
+			dir := t.TempDir()
+			target := filepath.Join(dir, "caller-owned")
+			path := target
+			data := []byte("already written diff")
+			if mode == "dangling-symlink" {
+				data = nil
+			} else {
+				require.NoError(t, os.WriteFile(target, data, 0o600))
+			}
+			if mode != "regular" {
+				path = filepath.Join(dir, "existing-link")
+				require.NoError(t, os.Symlink(target, path))
+			}
+			before, err := os.Lstat(path)
+			require.NoError(t, err)
+			cache, err := NewCache(-1, int64(os.Getpagesize()), path, false)
+			require.Nil(t, cache)
+			require.ErrorIs(t, err, unix.EINVAL)
+			after, err := os.Lstat(path)
+			require.NoError(t, err, "failed construction must not remove a pre-existing caller-owned path")
+			require.True(t, os.SameFile(before, after))
+			actual, err := os.ReadFile(target)
+			require.NoError(t, err)
+			assert.Equal(t, string(data), string(actual))
+		})
+	}
+}
+
+func TestCacheConstructionExistingOutputTransfersOnSuccess(t *testing.T) {
+	t.Parallel()
+	path := filepath.Join(t.TempDir(), "written-diff")
+	data := bytes.Repeat([]byte{0x6b}, os.Getpagesize())
+	require.NoError(t, os.WriteFile(path, data, 0o600))
+	cache, err := NewCache(int64(len(data)), int64(os.Getpagesize()), path, true)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = cache.Close() })
+	actual := make([]byte, len(data))
+	n, err := cache.ReadAt(actual, 0)
+	require.NoError(t, err)
+	require.Equal(t, len(data), n)
+	require.Equal(t, data, actual)
+	require.NoError(t, cache.Close())
+	_, err = os.Stat(path)
+	require.ErrorIs(t, err, os.ErrNotExist, "successful constructor transfers ordinary Close ownership")
+}