perf(orchestrator): avoid V3-ancestor header refresh

packages/orchestrator/pkg/sandbox/build_upload_test.go

@@ -5,11 +5,17 @@ package sandbox
 import (
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 
+	blockmocks "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block/mocks"
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/build"
+	templatemocks "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/template/mocks"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
+	headers "github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
 func newV4HeaderFF(t *testing.T, on bool) *featureflags.Client {
@@ -55,3 +61,92 @@ func TestResolveCompressConfig_V4_FlagOn(t *testing.T) {
 	ff := newV4HeaderFF(t, true)
 	require.True(t, resolveV4(t, ff))
 }
+
+// putV3Header registers a V3 ancestor in the fake cache. V3 headers carry no
+// Builds map at all — appendAncestorBuilds must synthesize a placeholder from
+// Metadata.Size so descendants don't pay a refresh roundtrip per cold read.
+func putV3Header(t *testing.T, cache *fakeCache, buildID uuid.UUID, fileType build.DiffType, size uint64) {
+	t.Helper()
+	tpl := templatemocks.NewMockTemplate(t)
+	dev := blockmocks.NewMockReadonlyDevice(t)
+	dev.EXPECT().Header().Return(&headers.Header{
+		Metadata: &headers.Metadata{Version: 3, Size: size},
+	}).Maybe()
+
+	switch fileType {
+	case build.Memfile:
+		tpl.EXPECT().Memfile(mock.Anything).Return(dev, nil).Maybe()
+	case build.Rootfs:
+		tpl.EXPECT().Rootfs().Return(dev, nil).Maybe()
+	}
+
+	cache.put(buildID.String(), tpl)
+}
+
+func mappingTo(t *testing.T, blockSize uint64, ancestorID uuid.UUID, length uint64) headers.Mapping {
+	t.Helper()
+	m, err := headers.NewMapping(blockSize, []headers.BuildMap{{
+		Offset: 0, Length: length, BuildId: ancestorID, BuildStorageOffset: 0,
+	}})
+	require.NoError(t, err)
+
+	return m
+}
+
+// V4 descendant of a V3 ancestor: appendAncestorBuilds writes a sentinel
+// empty BuildData{} so descendants take createDiff's hasEntry branch and
+// resolve() short-circuits on the UncompressedFrameTable hint from
+// GetBuildFrameData. Size is left zero on purpose — Metadata.Size is the
+// virtual size, not the diff size, and asking storage at upload time across
+// a long chain would multiply roundtrips. createDiff already falls back to
+// upstream.Size when bd.Size == 0.
+func TestAppendAncestorBuilds_V3AncestorSynthesizesEntry(t *testing.T) {
+	t.Parallel()
+
+	uploads, cache := newUploads(t)
+	ancestorID := uuid.New()
+	putV3Header(t, cache, ancestorID, build.Memfile, 512*1024*1024)
+
+	u := &Upload{buildID: uuid.New(), uploads: uploads}
+	dst := map[uuid.UUID]headers.BuildData{}
+
+	err := u.appendAncestorBuilds(t.Context(), dst, mappingTo(t, 4096, ancestorID, 4096), build.Memfile)
+	require.NoError(t, err)
+
+	bd, ok := dst[ancestorID]
+	require.True(t, ok, "V3 ancestor must produce a Builds entry")
+	require.Equal(t, int64(0), bd.Size, "Size must stay 0 — diff size is queried at read time via upstream.Size")
+	require.Nil(t, bd.FrameData, "FrameData must be nil so GetBuildFrameData returns UncompressedFrameTable")
+}
+
+// V4 ancestor: existing behavior — copy the ancestor's own self-entry verbatim.
+func TestAppendAncestorBuilds_V4AncestorCopiesEntry(t *testing.T) {
+	t.Parallel()
+
+	uploads, cache := newUploads(t)
+	ancestorID := uuid.New()
+	putHeader(t, cache, ancestorID, build.Memfile, false)
+
+	u := &Upload{buildID: uuid.New(), uploads: uploads}
+	dst := map[uuid.UUID]headers.BuildData{}
+
+	err := u.appendAncestorBuilds(t.Context(), dst, mappingTo(t, 4096, ancestorID, 4096), build.Memfile)
+	require.NoError(t, err)
+
+	_, ok := dst[ancestorID]
+	require.True(t, ok, "V4 ancestor's self-entry must be copied into dst")
+}
+
+// V3 caller (dst=nil): the barrier still runs, but no entry is written —
+// preserves the existing contract.
+func TestAppendAncestorBuilds_NilDstSkipsSynthesis(t *testing.T) {
+	t.Parallel()
+
+	uploads, cache := newUploads(t)
+	ancestorID := uuid.New()
+	putV3Header(t, cache, ancestorID, build.Memfile, 1024)
+
+	u := &Upload{buildID: uuid.New(), uploads: uploads}
+	err := u.appendAncestorBuilds(t.Context(), nil, mappingTo(t, 4096, ancestorID, 4096), build.Memfile)
+	require.NoError(t, err)
+}

packages/orchestrator/pkg/sandbox/build_upload_v4.go

@@ -128,8 +128,16 @@ func (u *Upload) uploadFramed(
 // (excluding self) — gating publish on parents' header finalization — and,
 // when dst is non-nil, writes the freshest BuildData into it. Existing dst
 // entries are overwritten (Wait is more authoritative than CloneForUpload).
-// Skips silently when Wait returns nil or the ancestor carries no Builds
-// entry (V3 ancestor); pre-existing dst entries are preserved.
+// Skips silently when Wait returns nil.
+//
+// V3 ancestors carry no Builds map, so a sentinel empty BuildData{} is
+// written — the entry's presence alone is what matters: GetBuildFrameData
+// returns UncompressedFrameTable (nil FrameData → sentinel), and createDiff's
+// hasEntry branch handles size=0 by falling back to upstream.Size. We avoid
+// computing the diff size here on purpose: it's not in Metadata.Size (that's
+// the virtual size), and asking storage at upload time across a long
+// ancestor chain would multiply roundtrips. The fallback amortizes into the
+// read that's about to happen anyway.
 //
 // V3 callers pass dst=nil — they need the barrier but have no Builds map.
 //
@@ -162,6 +170,12 @@ func (u *Upload) appendAncestorBuilds(
 
 		if bd, ok := h.Builds[buildID]; ok {
 			dst[buildID] = bd
+
+			continue
+		}
+
+		if h.Metadata.Version < headers.MetadataVersionV4 {
+			dst[buildID] = headers.BuildData{}
 		}
 	}