test(orch): cover MemfdIdentitySource.Slice and provisional fallback

Extend the MemfdIdentitySource test to exercise Slice (identity read while
mapped, and BytesNotAvailableError after release), and add a test that
buildProvisionalMemfile declines (returns nil, nil, nil) for a non-memfd-dedup source
or when disabled, so the caller falls back to the deduped header. Also cover
build.File.DurableHeader: it returns the live header when no durable future is
set and otherwise waits for the deduped future, which is what keeps a pause from
parenting off a provisional header, and SwapHeaderIfCurrent, which swaps only
when the current header still matches so a late provisional swap can't clobber a
header the upload already finalized. Cover the peer-to-peer header source too:
it serves the durable (deduped) header, so a peer never receives provisional
synthetic-build-id mappings.

Signed-off-by: Nikita Kalyazin <nikita.kalyazin@e2b.dev>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: kalyazin

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

@@ -464,12 +464,18 @@ func TestMemfdIdentitySource_ServesThenReleases(t *testing.T) {
 		require.Equal(t, int(ps), n)
 		require.Equal(t, data[page*ps:(page+1)*ps], b)
 	}
+	// Slice takes the same identity path and returns a copy of the memfd bytes.
+	sl, err := src.Slice(2*ps, ps)
+	require.NoError(t, err)
+	require.Equal(t, data[2*ps:3*ps], sl)
 	require.True(t, src.IsCached(t.Context(), 0, ps*4))
 
 	require.NoError(t, d.releaseMemfd())
 
-	_, err := src.ReadAt(make([]byte, ps), 0)
 	var bna BytesNotAvailableError
+	_, err = src.ReadAt(make([]byte, ps), 0)
+	require.ErrorAs(t, err, &bna)
+	_, err = src.Slice(0, ps)
 	require.ErrorAs(t, err, &bna)
 	require.False(t, src.IsCached(t.Context(), 0, ps))
 }

packages/orchestrator/pkg/sandbox/build/durable_header_test.go

@@ -0,0 +1,80 @@
+//go:build linux
+
+package build
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
+)
+
+// DurableHeader must return the live header when no durable future is set, and
+// otherwise wait for and return the durable (deduped) future's value — this is
+// what keeps a Pause from parenting a snapshot off a provisional header.
+func TestFile_DurableHeaderPrefersDurableFuture(t *testing.T) {
+	t.Parallel()
+
+	live := &header.Header{}
+	var f File
+	f.header.Store(live)
+
+	// No durable future: returns the live header immediately.
+	got, err := f.DurableHeader(t.Context())
+	require.NoError(t, err)
+	require.Same(t, live, got)
+
+	// With a durable future: blocks until it resolves, then returns it (not live).
+	durable := utils.NewSetOnce[*header.Header]()
+	f.SetDurableHeader(durable)
+
+	type res struct {
+		h   *header.Header
+		err error
+	}
+	ch := make(chan res, 1)
+	go func() {
+		h, e := f.DurableHeader(t.Context())
+		ch <- res{h, e}
+	}()
+
+	select {
+	case <-ch:
+		t.Fatal("DurableHeader returned before the durable future resolved")
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	deduped := &header.Header{}
+	require.NoError(t, durable.SetValue(deduped))
+
+	r := <-ch
+	require.NoError(t, r.err)
+	require.Same(t, deduped, r.h)
+	require.NotSame(t, live, r.h)
+}
+
+// SwapHeaderIfCurrent must replace the header only when it still matches, so a
+// late provisional→deduped swap cannot clobber a header another writer (e.g. an
+// upload finalizing the build) has already installed.
+func TestFile_SwapHeaderIfCurrent(t *testing.T) {
+	t.Parallel()
+
+	provisional := &header.Header{}
+	var f File
+	f.header.Store(provisional)
+
+	deduped := &header.Header{}
+	finalized := &header.Header{}
+
+	// Current still matches → swaps.
+	require.True(t, f.SwapHeaderIfCurrent(provisional, deduped))
+	require.Same(t, deduped, f.Header())
+
+	// Current already advanced (upload finalized it) → no-op, newer header kept.
+	f.header.Store(finalized)
+	require.False(t, f.SwapHeaderIfCurrent(provisional, deduped))
+	require.Same(t, finalized, f.Header())
+}

packages/orchestrator/pkg/sandbox/provisional_test.go

@@ -0,0 +1,41 @@
+//go:build linux
+
+package sandbox
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// fakeDiffSource is a block.DiffSource that is not a *block.DedupedMemfdCache,
+// so buildProvisionalMemfile must decline to produce a provisional header.
+type fakeDiffSource struct{}
+
+func (fakeDiffSource) Close() error                            { return nil }
+func (fakeDiffSource) ReadAt([]byte, int64) (int, error)       { return 0, nil }
+func (fakeDiffSource) Slice(int64, int64) ([]byte, error)      { return nil, nil }
+func (fakeDiffSource) Size() (int64, error)                    { return 0, nil }
+func (fakeDiffSource) FileSize(context.Context) (int64, error) { return 0, nil }
+func (fakeDiffSource) BlockSize() int64                        { return 0 }
+func (fakeDiffSource) Path(context.Context) (string, error)    { return "", nil }
+
+// buildProvisionalMemfile must fall back (return nil, nil, nil) when it can't
+// apply, so the caller keeps using the deduped header: a non-memfd-dedup source,
+// and a disabled flag.
+func TestBuildProvisionalMemfile_FallsBack(t *testing.T) {
+	t.Parallel()
+
+	// Source is not a *block.DedupedMemfdCache → no provisional serving.
+	h, d, swapDone := buildProvisionalMemfile(t.Context(), fakeDiffSource{}, true, nil, nil, nil)
+	require.Nil(t, h)
+	require.Nil(t, d)
+	require.Nil(t, swapDone)
+
+	// Disabled flag → no provisional serving.
+	h, d, swapDone = buildProvisionalMemfile(t.Context(), fakeDiffSource{}, false, nil, nil, nil)
+	require.Nil(t, h)
+	require.Nil(t, d)
+	require.Nil(t, swapDone)
+}

packages/orchestrator/pkg/sandbox/template/peerserver/header_test.go

@@ -3,6 +3,7 @@
 package peerserver
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -16,6 +17,20 @@ import (
 	storageheader "github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
+// durableDevice is a ReadonlyDevice that also exposes DurableHeader, like the
+// real *template.Storage. Header() stands in for the live (provisional) header,
+// DurableHeader for the deduped one a pause/peer must use instead.
+type durableDevice struct {
+	*blockmocks.MockReadonlyDevice
+
+	durable    *storageheader.Header
+	durableErr error
+}
+
+func (d durableDevice) DurableHeader(context.Context) (*storageheader.Header, error) {
+	return d.durable, d.durableErr
+}
+
 func TestHeaderSource_Stream(t *testing.T) {
 	t.Parallel()
 
@@ -59,6 +74,34 @@ func TestHeaderSource_Stream_NilHeader(t *testing.T) {
 	assert.ErrorIs(t, err, ErrNotAvailable)
 }
 
+// When the device exposes DurableHeader, Stream must serve that (the deduped
+// header) rather than the live one, so a peer never receives provisional
+// synthetic-build-id mappings it can't resolve. The live Header() returns nil
+// here, so the stream succeeds only if DurableHeader is used.
+func TestHeaderSource_Stream_ServesDurableHeader(t *testing.T) {
+	t.Parallel()
+
+	deduped, err := storageheader.NewHeader(&storageheader.Metadata{Version: 3, BlockSize: 4096, Size: 4096}, nil)
+	require.NoError(t, err)
+
+	mockDev := blockmocks.NewMockReadonlyDevice(t)
+	mockDev.EXPECT().Header().Return(nil).Maybe()
+	dev := durableDevice{MockReadonlyDevice: mockDev, durable: deduped}
+
+	tmplMock := templatemocks.NewMockTemplate(t)
+	tmplMock.EXPECT().Memfile(mock.Anything).Return(dev, nil)
+
+	cache := peerservermocks.NewMockCache(t)
+	cache.EXPECT().GetCachedTemplate("build-1").Return(tmplMock, true)
+
+	src, err := ResolveBlob(cache, "build-1", storage.MemfileName+storage.HeaderSuffix)
+	require.NoError(t, err)
+
+	sender := &collectSender{}
+	require.NoError(t, src.Stream(t.Context(), sender))
+	assert.NotEmpty(t, sender.data)
+}
+
 func TestHeaderSource_Stream_Rootfs(t *testing.T) {
 	t.Parallel()