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.

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)
+}