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; that a pinned build-store entry is skipped by
disk-pressure eviction; 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, and that storageTemplate.SchedulingMetadata reads
the durable header non-blockingly (DurableHeaderNow) so memfile_build_ids
exclude the provisional build id and never block the create/resume response —
reporting rootfs-only metadata while the deduped header is still pending.

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/cache_test.go

@@ -271,6 +271,37 @@ func TestDiffStoreDelayEvictionAbort(t *testing.T) { //nolint:paralleltest // ve
 	assert.True(t, found)
 }
 
+// A pinned entry must be skipped by disk-pressure eviction (the next-oldest is
+// chosen instead), and become eligible again once unpinned.
+func TestDiffStorePinnedSkippedByEviction(t *testing.T) {
+	t.Parallel()
+	cachePath := t.TempDir()
+
+	c, err := cfg.Parse()
+	require.NoError(t, err)
+	flags := flagsWithMaxBuildCachePercentage(t, 100)
+	store, err := NewDiffStore(c, flags, cachePath, 60*time.Second, 4*time.Second)
+	require.NoError(t, err)
+
+	oldest := newRootFSDiff(t, cachePath, "pin-oldest")
+	store.Add(oldest)
+	newer := newRootFSDiff(t, cachePath, "pin-newer")
+	store.Add(newer)
+
+	// Pin the oldest → eviction skips it and schedules the next-oldest.
+	store.Pin(oldest.CacheKey())
+	_, err = store.deleteOldestFromCache(t.Context())
+	require.NoError(t, err)
+	assert.False(t, store.isBeingDeleted(oldest.CacheKey()))
+	assert.True(t, store.isBeingDeleted(newer.CacheKey()))
+
+	// Unpin → the oldest is eligible again.
+	store.Unpin(oldest.CacheKey())
+	_, err = store.deleteOldestFromCache(t.Context())
+	require.NoError(t, err)
+	assert.True(t, store.isBeingDeleted(oldest.CacheKey()))
+}
+
 func TestDiffStoreOldestFromCache(t *testing.T) {
 	t.Parallel()
 	cachePath := t.TempDir()

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

@@ -0,0 +1,110 @@
+//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())
+}
+
+// DurableHeaderNow is the non-blocking form: (Header, true) when the durable
+// header is known (no swap pending, or the deduped future already resolved) and
+// (nil, false) while a swap is still pending — used on latency-sensitive paths.
+func TestFile_DurableHeaderNow(t *testing.T) {
+	t.Parallel()
+
+	live := &header.Header{}
+	var f File
+	f.header.Store(live)
+
+	// No durable future: known immediately, returns the live header.
+	h, ok := f.DurableHeaderNow()
+	require.True(t, ok)
+	require.Same(t, live, h)
+
+	// Durable future set but unresolved: not known yet, non-blocking (nil, false).
+	durable := utils.NewSetOnce[*header.Header]()
+	f.SetDurableHeader(durable)
+	h, ok = f.DurableHeaderNow()
+	require.False(t, ok)
+	require.Nil(t, h)
+
+	// Once resolved: returns the durable (deduped) header.
+	deduped := &header.Header{}
+	require.NoError(t, durable.SetValue(deduped))
+	h, ok = f.DurableHeaderNow()
+	require.True(t, ok)
+	require.Same(t, deduped, h)
+}

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()
 

packages/orchestrator/pkg/sandbox/template/storage_template_test.go

@@ -0,0 +1,82 @@
+//go:build linux
+
+package template
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block"
+	blockmocks "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block/mocks"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
+)
+
+// durableMemfileDevice is a ReadonlyDevice that also exposes DurableHeaderNow,
+// like the real *Storage. Header() stands in for the live (provisional) header;
+// DurableHeaderNow returns the deduped header scheduling metadata must use, and
+// its ready flag models whether the deduped header has resolved yet.
+type durableMemfileDevice struct {
+	*blockmocks.MockReadonlyDevice
+
+	durable *header.Header
+	ready   bool
+}
+
+func (d durableMemfileDevice) DurableHeaderNow() (*header.Header, bool) {
+	return d.durable, d.ready
+}
+
+func schedulingTemplate(t *testing.T, mem block.ReadonlyDevice, rootfsBase uuid.UUID) *storageTemplate {
+	t.Helper()
+	rootfsHdr, err := header.NewHeader(&header.Metadata{Version: 3, BlockSize: 4096, Size: 4096, BaseBuildId: rootfsBase}, nil)
+	require.NoError(t, err)
+	rootfsDev := blockmocks.NewMockReadonlyDevice(t)
+	rootfsDev.EXPECT().Header().Return(rootfsHdr)
+
+	tmpl := &storageTemplate{
+		memfile: utils.NewSetOnce[block.ReadonlyDevice](),
+		rootfs:  utils.NewSetOnce[block.ReadonlyDevice](),
+	}
+	require.NoError(t, tmpl.rootfs.SetValue(rootfsDev))
+	require.NoError(t, tmpl.memfile.SetValue(mem))
+
+	return tmpl
+}
+
+// When the deduped header has resolved, SchedulingMetadata reports it (never the
+// live/provisional header) — so memfile build ids reflect the real build id.
+func TestStorageTemplate_SchedulingMetadataUsesDurableHeader(t *testing.T) {
+	t.Parallel()
+
+	dedupedBase := uuid.New()
+	dedupedHdr, err := header.NewHeader(&header.Metadata{Version: 3, BlockSize: 4096, Size: 4096, BaseBuildId: dedupedBase}, nil)
+	require.NoError(t, err)
+	memMock := blockmocks.NewMockReadonlyDevice(t)
+	memMock.EXPECT().Header().Return(nil).Maybe()
+	memDev := durableMemfileDevice{MockReadonlyDevice: memMock, durable: dedupedHdr, ready: true}
+
+	md := schedulingTemplate(t, memDev, uuid.New()).SchedulingMetadata(t.Context())
+	require.NotNil(t, md)
+	assert.Equal(t, dedupedBase.String(), md.GetMemfileBaseBuildId())
+}
+
+// While the deduped header is still pending (provisional window), SchedulingMetadata
+// must NOT block and must NOT emit the provisional build id: it reports rootfs-only
+// metadata (empty memfile base build id).
+func TestStorageTemplate_SchedulingMetadataSkipsPendingMemfile(t *testing.T) {
+	t.Parallel()
+
+	rootfsBase := uuid.New()
+	memMock := blockmocks.NewMockReadonlyDevice(t)
+	memMock.EXPECT().Header().Return(nil).Maybe()
+	memDev := durableMemfileDevice{MockReadonlyDevice: memMock, durable: nil, ready: false}
+
+	md := schedulingTemplate(t, memDev, rootfsBase).SchedulingMetadata(t.Context())
+	require.NotNil(t, md)
+	assert.Empty(t, md.GetMemfileBaseBuildId())
+	assert.Equal(t, rootfsBase.String(), md.GetRootfsBaseBuildId())
+}