From 1dffdb819b50de861f196de380bc2316e98655c6 Mon Sep 17 00:00:00 2001 From: zhukunshuai Date: Tue, 15 Sep 2026 19:39:59 +0800 Subject: [PATCH] fix(build): gate the optimize phase's metadata re-upload on finalize's uploads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The optimize phase runs right after finalize, but finalize's PauseAndUpload pushes its uploads (snapshot data and the pre-prefetch metadata) onto the build's UploadErrGroup and returns immediately — the only wait for that group runs after every phase. The optimize phase then collected its mapping and re-uploaded metadata with no ordering against the finalize goroutine, so the two writers raced on the same metadata object for the build: an interleaving where the finalize upload finished last clobbered the just-uploaded prefetch mapping (silent loss, manifests as "prefetch mapping sometimes missing"). finalizeUploadsSettled now waits on the build's UploadErrGroup before the optimize phase re-uploads metadata. A failed finalize upload skips publishing — the remote build is incomplete, and the same error resurfaces at the builder-level wait, so the build still fails loudly. Test: TestFinalizeUploadsSettled covers the settled, failed, blocked (ordering invariant: the gate does not return while the finalize upload is in flight) and nil-group cases. Co-Authored-By: Claude --- .../template/build/phases/optimize/builder.go | 35 ++++++ .../build/phases/optimize/upload_race_test.go | 108 ++++++++++++++++++ 2 files changed, 143 insertions(+) create mode 100644 packages/orchestrator/pkg/template/build/phases/optimize/upload_race_test.go diff --git a/packages/orchestrator/pkg/template/build/phases/optimize/builder.go b/packages/orchestrator/pkg/template/build/phases/optimize/builder.go index 5e256a2a53..b849cbf831 100644 --- a/packages/orchestrator/pkg/template/build/phases/optimize/builder.go +++ b/packages/orchestrator/pkg/template/build/phases/optimize/builder.go @@ -153,6 +153,25 @@ func (pb *OptimizeBuilder) Build( Memory: memoryPrefetchMapping, }) + // Wait for the finalize phase's async uploads (snapshot data and the + // pre-prefetch metadata) before re-uploading metadata with the prefetch + // mapping. Finalize uploads via the build's UploadErrGroup and returns + // immediately, so without this wait the goroutine's metadata upload and + // updateMetadata below race on the same metadata object for this build — + // the later writer wins, and an interleaving that finishes the finalize + // upload last clobbers the prefetch mapping. Waiting is also the earliest + // point where re-publishing makes sense: a failed finalize upload means + // the remote build is incomplete and the mapping would dangle. + if err := pb.finalizeUploadsSettled(); err != nil { + pb.logger.Warn(ctx, "finalize upload failed; skipping prefetch metadata publish", zap.Error(err)) + + return phases.LayerResult{ + Metadata: sourceLayer.Metadata, + Cached: false, + Hash: currentLayer.Hash, + }, nil + } + // Upload the updated metadata err = pb.updateMetadata(ctx, updatedMetadata) if err != nil { @@ -256,6 +275,22 @@ func (pb *OptimizeBuilder) runSandboxAndCollectPrefetch( return prefetchData, nil } +// finalizeUploadsSettled waits for the finalize phase's async uploads and +// reports whether they all succeeded. The build's UploadErrGroup carries the +// finalize goroutine (snapshot data + the pre-prefetch metadata upload); the +// builder's own Wait for the same group runs only after every phase, so +// without this gate the optimize phase's metadata re-upload can interleave +// with — and be clobbered by — the finalize upload writing the same object. +// A non-nil error means the finalize upload failed; the same error +// resurfaces at the builder-level wait, so callers treat it as a soft skip. +func (pb *OptimizeBuilder) finalizeUploadsSettled() error { + if pb.UploadErrGroup == nil { + return nil + } + + return pb.UploadErrGroup.Wait() +} + // updateMetadata updates the template metadata in storages. func (pb *OptimizeBuilder) updateMetadata(ctx context.Context, t metadata.Template) error { err := metadata.UploadMetadata(ctx, pb.templateStorage, t, pb.BuildContext.Config.ObjectMetadata(storage.ObjectOriginTemplateBuild)) diff --git a/packages/orchestrator/pkg/template/build/phases/optimize/upload_race_test.go b/packages/orchestrator/pkg/template/build/phases/optimize/upload_race_test.go new file mode 100644 index 0000000000..320e1db240 --- /dev/null +++ b/packages/orchestrator/pkg/template/build/phases/optimize/upload_race_test.go @@ -0,0 +1,108 @@ +//go:build linux + +package optimize + +import ( + "errors" + "golang.org/x/sync/errgroup" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/buildcontext" +) + +// TestFinalizeUploadsSettled gates the optimize phase's metadata re-upload on +// the finalize phase's async uploads (see Build). The build's UploadErrGroup +// carries the finalize upload goroutine; without the wait, its metadata upload +// and the optimize phase's updateMetadata race on the same metadata object +// for the build, and the later writer wins — clobbering the prefetch mapping. +func TestFinalizeUploadsSettled(t *testing.T) { + t.Parallel() + + t.Run("settled group reports success", func(t *testing.T) { + t.Parallel() + + eg := &errgroup.Group{} + eg.Go(func() error { return nil }) + require.NoError(t, eg.Wait()) // settle it first + + pb := &OptimizeBuilder{ + BuildContext: buildcontext.BuildContext{UploadErrGroup: eg}, + } + + assert.NoError(t, pb.finalizeUploadsSettled()) + }) + + t.Run("failed upload reports error", func(t *testing.T) { + t.Parallel() + + uploadErr := errors.New("upload failed") + eg := &errgroup.Group{} + eg.Go(func() error { return uploadErr }) + require.Error(t, eg.Wait()) // settle it with the failure + + pb := &OptimizeBuilder{ + BuildContext: buildcontext.BuildContext{UploadErrGroup: eg}, + } + + err := pb.finalizeUploadsSettled() + require.Error(t, err) + assert.ErrorIs(t, err, uploadErr) + }) + + t.Run("wait blocks until the finalize upload finishes", func(t *testing.T) { + t.Parallel() + + // The ordering invariant this phase depends on: finalizeUploadsSettled + // must not return while the finalize upload goroutine is still + // running, so updateMetadata (called only after the gate) can never + // interleave with the finalize goroutine's metadata upload. + eg := &errgroup.Group{} + started := make(chan struct{}) + release := make(chan struct{}) + eg.Go(func() error { + close(started) + <-release + return nil + }) + <-started // the goroutine is in flight + + pb := &OptimizeBuilder{ + BuildContext: buildcontext.BuildContext{UploadErrGroup: eg}, + } + + settled := make(chan error, 1) + go func() { settled <- pb.finalizeUploadsSettled() }() + + select { + case err := <-settled: + t.Fatalf("finalizeUploadsSettled returned %v while the finalize upload was still in flight", err) + case <-time.After(50 * time.Millisecond): + // still blocked — the gate is holding, as required. + } + + close(release) + select { + case err := <-settled: + assert.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("finalizeUploadsSettled did not return after the upload finished") + } + }) + + t.Run("nil group is treated as settled", func(t *testing.T) { + t.Parallel() + + // The BuildContext is constructed by the builder; a nil group would + // otherwise panic on Wait. Treat it as settled (nothing to wait for) + // — matches how a build without in-flight uploads behaves. + pb := &OptimizeBuilder{ + BuildContext: buildcontext.BuildContext{}, + } + + assert.NoError(t, pb.finalizeUploadsSettled()) + }) +}