diff --git a/packages/orchestrator/pkg/sandbox/template/storage_template.go b/packages/orchestrator/pkg/sandbox/template/storage_template.go index dc400fe392..66d14c626b 100644 --- a/packages/orchestrator/pkg/sandbox/template/storage_template.go +++ b/packages/orchestrator/pkg/sandbox/template/storage_template.go @@ -340,5 +340,5 @@ func (t *storageTemplate) UpdateMetadata(meta metadata.Template) error { return fmt.Errorf("failed to get metafile: %w", err) } - return meta.ToFile(metafile.Path()) + return meta.ReplaceFile(metafile.Path()) } diff --git a/packages/orchestrator/pkg/server/prefetch_harvest.go b/packages/orchestrator/pkg/server/prefetch_harvest.go index a739e49da6..55b5a8d802 100644 --- a/packages/orchestrator/pkg/server/prefetch_harvest.go +++ b/packages/orchestrator/pkg/server/prefetch_harvest.go @@ -2,6 +2,7 @@ package server import ( "context" + "errors" "fmt" "maps" "slices" @@ -263,21 +264,25 @@ func (h *prefetchHarvester) run( // update is enough for a same-node resume, so do it regardless of whether the // remote upload succeeded. meta = meta.WithPrefetch(&metadata.Prefetch{Memory: mapping}) + var localUpdateErr error if err := h.templates.UpdateMetadata(buildID, meta); err != nil { - return pages, harvestSuccess, fmt.Errorf("update local metadata: %w", err) + localUpdateErr = fmt.Errorf("update local metadata: %w", err) + if !errors.Is(err, metadata.ErrReplaceCommitted) { + return pages, harvestSuccess, localUpdateErr + } } // Only enrich the remote metadata if the snapshot actually landed; on upload // failure the remote build is incomplete, so there is nothing to enrich (the // local update above still lets a same-node resume prefetch). if uploadErr != nil { - return pages, harvestSuccess, nil //nolint:nilerr // remote snapshot did not land; the local update is the most we can do + return pages, harvestSuccess, localUpdateErr } if err := h.uploadMetadata(ctx, meta, objectMetadata); err != nil { - return pages, harvestSuccess, fmt.Errorf("re-upload metadata: %w", err) + return pages, harvestSuccess, errors.Join(localUpdateErr, fmt.Errorf("re-upload metadata: %w", err)) } - return pages, harvestSuccess, nil + return pages, harvestSuccess, localUpdateErr } // resumeMapping resumes a throwaway warm copy of the just-paused snapshot, diff --git a/packages/orchestrator/pkg/server/prefetch_harvest_test.go b/packages/orchestrator/pkg/server/prefetch_harvest_test.go index 8b7ad79de1..5ba5c58822 100644 --- a/packages/orchestrator/pkg/server/prefetch_harvest_test.go +++ b/packages/orchestrator/pkg/server/prefetch_harvest_test.go @@ -5,6 +5,7 @@ package server import ( "context" "errors" + "fmt" "testing" "time" @@ -341,6 +342,21 @@ func TestHarvestRun_UploadFailedKeepsLocalSkipsRemote(t *testing.T) { require.False(t, p.uploadCalled, "remote re-upload is skipped when the snapshot did not land") } +func TestHarvestRun_CommittedLocalUpdateErrorStillUpdatesRemote(t *testing.T) { + t.Parallel() + + p := newHarvestProbe() + p.tmpls.updateErr = fmt.Errorf("directory sync failed: %w", metadata.ErrReplaceCommitted) + + pages, outcome, err := p.run(t.Context(), true) + + require.ErrorIs(t, err, metadata.ErrReplaceCommitted) + require.Equal(t, harvestSuccess, outcome) + require.Equal(t, 2, pages) + require.True(t, p.tmpls.updated) + require.True(t, p.uploadCalled, "remote metadata should be updated after the local replacement commits") +} + // TestHarvestRun_DeadlineDuringWaitLeavesMetadataUntouched: if the harvest // deadline fires while waiting for the upload, neither local nor remote metadata // is touched (the upload may still be reading/writing them). diff --git a/packages/orchestrator/pkg/template/metadata/prefetch_carry_test.go b/packages/orchestrator/pkg/template/metadata/prefetch_carry_test.go index 0d7354530a..56c009d944 100644 --- a/packages/orchestrator/pkg/template/metadata/prefetch_carry_test.go +++ b/packages/orchestrator/pkg/template/metadata/prefetch_carry_test.go @@ -57,35 +57,36 @@ func TestWithPrefetchSetsMappingAndPreservesFields(t *testing.T) { } // TestPrefetchSurvivesMetadataFileRoundTrip is the persist→resume link the -// consume path depends on: the harvest writes the mapping via ToFile/UploadMetadata +// consume path depends on: the harvest writes the mapping via ReplaceFile/UploadMetadata // and the resume reads it via FromFile. If Prefetch did not survive JSON // serialization (e.g. a missing/renamed tag), the whole feature would silently // no-op and every in-memory test above would still pass — so assert the round trip. func TestPrefetchSurvivesMetadataFileRoundTrip(t *testing.T) { t.Parallel() - orig := Template{ + base := Template{ Version: CurrentVersion, Template: TemplateMetadata{BuildID: "build-rt", KernelVersion: "6.1", FirecrackerVersion: "1.14"}, Context: Context{User: "root"}, - Prefetch: &Prefetch{Memory: &MemoryPrefetchMapping{ - Indices: []uint64{3, 1, 2}, - AccessTypes: []AccessType{"r", "w", "p"}, - BlockSize: 2 << 20, - }}, } + updated := base.WithPrefetch(&Prefetch{Memory: &MemoryPrefetchMapping{ + Indices: []uint64{3, 1, 2}, + AccessTypes: []AccessType{"r", "w", "p"}, + BlockSize: 2 << 20, + }}) path := filepath.Join(t.TempDir(), "metadata.json") - require.NoError(t, orig.ToFile(path)) + require.NoError(t, base.ToFile(path)) + require.NoError(t, updated.ReplaceFile(path)) got, err := FromFile(path) require.NoError(t, err) require.NotNil(t, got.Prefetch, "Prefetch must survive ToFile/FromFile") require.NotNil(t, got.Prefetch.Memory) - assert.Equal(t, orig.Prefetch.Memory.Indices, got.Prefetch.Memory.Indices, "ordered indices must be preserved") - assert.Equal(t, orig.Prefetch.Memory.AccessTypes, got.Prefetch.Memory.AccessTypes) - assert.Equal(t, orig.Prefetch.Memory.BlockSize, got.Prefetch.Memory.BlockSize) + assert.Equal(t, updated.Prefetch.Memory.Indices, got.Prefetch.Memory.Indices, "ordered indices must be preserved") + assert.Equal(t, updated.Prefetch.Memory.AccessTypes, got.Prefetch.Memory.AccessTypes) + assert.Equal(t, updated.Prefetch.Memory.BlockSize, got.Prefetch.Memory.BlockSize) } func TestMemoryPrefetchMappingCount(t *testing.T) { diff --git a/packages/orchestrator/pkg/template/metadata/template_metadata.go b/packages/orchestrator/pkg/template/metadata/template_metadata.go index 2a5f344d4e..4160e52201 100644 --- a/packages/orchestrator/pkg/template/metadata/template_metadata.go +++ b/packages/orchestrator/pkg/template/metadata/template_metadata.go @@ -30,6 +30,10 @@ const ( FilesystemOnlyVersion = DeprecatedVersion + 1 ) +// ErrReplaceCommitted means the replacement is visible, but parent-directory +// durability could not be confirmed. +var ErrReplaceCommitted = ioutils.ErrAtomicWriteCommitted + var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/orchestrator/pkg/template/metadata") // AccessType is a compact representation of block access type for JSON serialization. @@ -221,6 +225,20 @@ func (t Template) ToFile(path string) error { return nil } +func (t Template) ReplaceFile(path string) error { + mr, err := serialize(t) + if err != nil { + return err + } + + err = ioutils.WriteToFileFromReaderAtomically(path, mr) + if err != nil { + return fmt.Errorf("failed to replace metadata file: %w", err) + } + + return nil +} + func FromFile(path string) (Template, error) { f, err := os.Open(path) if err != nil { diff --git a/packages/shared/pkg/ioutils/write_to_file.go b/packages/shared/pkg/ioutils/write_to_file.go index 2413ddb6bc..7a9597eaf6 100644 --- a/packages/shared/pkg/ioutils/write_to_file.go +++ b/packages/shared/pkg/ioutils/write_to_file.go @@ -1,26 +1,98 @@ package ioutils import ( + "errors" + "fmt" "io" "os" + "path/filepath" ) +var ErrAtomicWriteCommitted = errors.New("atomic write committed") + func WriteToFileFromReader(path string, r io.Reader) (err error) { - // Create (truncate if exists) with 0644 perms f, err := os.Create(path) if err != nil { return err } - // Make sure we return a close error if that's the only error. defer func() { - if cerr := f.Close(); err == nil && cerr != nil { - err = cerr + if closeErr := f.Close(); err == nil && closeErr != nil { + err = closeErr + } + }() + + if _, err = io.Copy(f, r); err != nil { + return err + } + + return f.Sync() +} + +// WriteToFileFromReaderAtomically replaces an existing regular file without +// exposing partial contents. Errors wrapping ErrAtomicWriteCommitted mean the +// replacement is visible, but parent-directory durability could not be confirmed. +func WriteToFileFromReaderAtomically(path string, r io.Reader) (err error) { + info, err := os.Lstat(path) + if err != nil { + return err + } + if !info.Mode().IsRegular() { + return fmt.Errorf("atomic write target %q is not a regular file", path) + } + + dir := filepath.Dir(path) + f, err := os.CreateTemp(dir, ".atomic-write-*") + if err != nil { + return err + } + tempPath := f.Name() + closed := false + committed := false + defer func() { + if !closed { + if closeErr := f.Close(); closeErr != nil { + err = errors.Join(err, fmt.Errorf("close temporary file: %w", closeErr)) + } + } + if !committed { + if removeErr := os.Remove(tempPath); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) { + err = errors.Join(err, fmt.Errorf("remove temporary file: %w", removeErr)) + } } }() if _, err = io.Copy(f, r); err != nil { return err } + if err = f.Chmod(info.Mode().Perm()); err != nil { + return err + } + if err = f.Sync(); err != nil { + return err + } + err = f.Close() + closed = true + if err != nil { + return err + } + + if err = os.Rename(tempPath, path); err != nil { + return err + } + committed = true + + if err = syncDirectory(dir); err != nil { + return fmt.Errorf("%w: sync parent directory: %w", ErrAtomicWriteCommitted, err) + } + + return nil +} + +func syncDirectory(path string) error { + dir, err := os.Open(path) + if err != nil { + return err + } - return f.Sync() // ensure contents hit disk + return errors.Join(dir.Sync(), dir.Close()) } diff --git a/packages/shared/pkg/ioutils/write_to_file_test.go b/packages/shared/pkg/ioutils/write_to_file_test.go new file mode 100644 index 0000000000..c7200e6087 --- /dev/null +++ b/packages/shared/pkg/ioutils/write_to_file_test.go @@ -0,0 +1,146 @@ +package ioutils + +import ( + "errors" + "io" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "testing/iotest" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestWriteToFileFromReaderAtomicallyReplacesExistingFile(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "metadata.json") + require.NoError(t, os.WriteFile(path, []byte("old metadata"), 0o600)) + + reader := &blockingReader{ + data: []byte("new metadata"), + started: make(chan struct{}), + release: make(chan struct{}), + } + errCh := make(chan error, 1) + go func() { + errCh <- WriteToFileFromReaderAtomically(path, reader) + }() + + released := false + finished := false + defer func() { + if !released { + close(reader.release) + } + if !finished { + waitForWrite(t, errCh) + } + }() + + select { + case <-reader.started: + case err := <-errCh: + finished = true + t.Fatalf("writer completed before reading input: %v", err) + case <-time.After(5 * time.Second): + t.Fatal("writer did not read input") + } + + contents, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, []byte("old metadata"), contents) + + close(reader.release) + released = true + err = waitForWrite(t, errCh) + finished = true + require.NoError(t, err) + + contents, err = os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, []byte("new metadata"), contents) + + info, err := os.Stat(path) + require.NoError(t, err) + assert.Equal(t, os.FileMode(0o600), info.Mode().Perm()) +} + +func TestWriteToFileFromReaderFailureKeepsExistingFile(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + path := filepath.Join(dir, "metadata.json") + require.NoError(t, os.WriteFile(path, []byte("old metadata"), 0o644)) + + readErr := errors.New("read failed") + err := WriteToFileFromReaderAtomically(path, io.MultiReader(strings.NewReader("partial"), iotest.ErrReader(readErr))) + require.ErrorIs(t, err, readErr) + + contents, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, []byte("old metadata"), contents) + + entries, err := os.ReadDir(dir) + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "metadata.json", entries[0].Name()) +} + +func TestWriteToFileFromReaderAtomicallyRequiresExistingRegularFile(t *testing.T) { + t.Parallel() + + t.Run("missing file", func(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "metadata.json") + err := WriteToFileFromReaderAtomically(path, strings.NewReader("metadata")) + require.ErrorIs(t, err, os.ErrNotExist) + }) + + t.Run("directory", func(t *testing.T) { + t.Parallel() + + path := t.TempDir() + err := WriteToFileFromReaderAtomically(path, strings.NewReader("metadata")) + require.Error(t, err) + assert.Contains(t, err.Error(), "is not a regular file") + }) +} + +func waitForWrite(t *testing.T, errCh <-chan error) error { + t.Helper() + + select { + case err := <-errCh: + return err + case <-time.After(5 * time.Second): + t.Fatal("writer did not finish") + + return nil + } +} + +type blockingReader struct { + data []byte + started chan struct{} + release chan struct{} + once sync.Once +} + +func (r *blockingReader) Read(p []byte) (int, error) { + r.once.Do(func() { close(r.started) }) + <-r.release + if len(r.data) == 0 { + return 0, io.EOF + } + + n := copy(p, r.data) + r.data = r.data[n:] + + return n, nil +}