e2b-dev · kalyazin · Jul 10, 2026 · Jul 10, 2026 · Jul 10, 2026 · Jul 10, 2026
@@ -0,0 +1,144 @@
+//go:build linux
+
+package sandbox
+
+import (
+	"slices"
+
+	"github.com/google/uuid"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/metadata"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+	"github.com/e2b-dev/infra/packages/shared/pkg/units"
+)
+
+// resume-prefetch-source flag values (see featureflags.ResumePrefetchSourceFlag).
+const (
+	resumePrefetchOff       = "off"
+	resumePrefetchInit      = "init"
+	resumePrefetchLastCycle = "last-cycle"
+	resumePrefetchBoth      = "both"
+)
+
+// selectResumePrefetch maps the resume-prefetch-source flag value to which
+// traces the resume prefetcher replays. "init" (and any unknown value) keeps
+// today's behavior — replay only the build-time/harvested init trace — so the
+// default is a no-op-equivalent.
+func selectResumePrefetch(source string) (useInit, useLastCycle bool) {
+	switch source {
+	case resumePrefetchOff:
+		return false, false
+	case resumePrefetchLastCycle:
+		return false, true
+	case resumePrefetchBoth:
+		return true, true
+	default: // resumePrefetchInit and any unrecognized value
+		return true, false
+	}
+}
+
+// capResumePrefetch bounds a prefetch mapping to at most maxMiB of blocks,
+// keeping the earliest (offset-order) blocks and leaving the rest to
+// demand-fault. maxMiB < 0 means uncapped; m is returned unchanged when it
+// already fits, the cap is disabled, or the block size is unknown. A nil
+// mapping returns nil.
+func capResumePrefetch(m *metadata.MemoryPrefetchMapping, maxMiB int) *metadata.MemoryPrefetchMapping {
+	if m == nil || maxMiB < 0 || m.BlockSize <= 0 {
+		return m
+	}
+
+	maxBlocks := units.MBToBytes(int64(maxMiB)) / m.BlockSize
+	if int64(len(m.Indices)) <= maxBlocks {
+		return m
+	}
+
+	return &metadata.MemoryPrefetchMapping{
+		Indices:   m.Indices[:maxBlocks],
+		BlockSize: m.BlockSize,
+	}
+}
+
+// buildDiffMemoryPrefetchMapping builds a prefetch mapping over only the
+// blocks owned by this snapshot's own pause diff (BuildMap.BuildId ==
+// Metadata.BuildId) — i.e. the pages the last resume→pause cycle actually
+// wrote (the sandbox's "last-cycle" working set).
+//
+// Pause/resume carries no build-time prefetch mapping (SameVersionTemplate
+// drops it on the resumed snapshot's metadata), so today it demand-faults
+// everything cold. But the pages a resume→pause cycle wrote are already
+// recorded: they ARE the pause diff, present in the merged memfile header as
+// the BuildMap entries whose BuildId equals this header's own build ID. No
+// separate dirty-bitmap capture at pause time is needed — deriving the
+// mapping from the header at resume is sufficient and reuses data that
+// already exists.
+//
+// Blocks are enumerated at the memfile's block size (the prefetcher's fetch
+// unit) and returned in offset order — readahead- and coalescing-friendly,
+// and there is no access-order predictor for pause/resume anyway.
+//
+// Returns nil if the diff is empty, the header/metadata isn't usable, or the
+// header has no base layer. The last guard matters: a base / build template
+// (as opposed to a paused snapshot) maps its entire memfile to its own
+// BuildId, so without it a create/build resume with resume-prefetch-source
+// including last-cycle would prefetch the whole template image, not a diff.
+func buildDiffMemoryPrefetchMapping(h *header.Header) *metadata.MemoryPrefetchMapping {
+	if h == nil || h.Metadata == nil || h.Metadata.BlockSize == 0 {
+		return nil
+	}
+
+	blockSize := int64(h.Metadata.BlockSize)
+	own := h.Metadata.BuildId
+
+	// A merged/dedup memfile header has overlapping BuildMap entries across
+	// layers (precedence is resolved at read time), so the same block offset
+	// can appear in multiple entries. Dedup by block index: fetch each block
+	// once (source.Slice resolves the top layer itself). Offsets/lengths are
+	// uint64; walk in uint64 to avoid an int64 cast that could overflow.
+	uBlockSize := uint64(blockSize)
+	seen := make(map[uint64]struct{})
+	hasBase := false
+	for _, bm := range h.Mapping.All() {
+		if bm.BuildId == uuid.Nil {
+			continue
+		}
+		if bm.BuildId != own {
+			hasBase = true // a distinct base/parent layer exists
+
+			continue
+		}
+		if bm.Length == 0 {
+			continue
+		}
+
+		// Enumerate by block index, not by byte stride. A merged/dedup entry
+		// can start unaligned to the block size (dedup emits PageSize-granular
+		// entries, and NormalizeMappings merges adjacent same-build runs), so
+		// stepping by uBlockSize from an unaligned Offset would skip the final
+		// block whenever the entry straddles a block boundary.
+		firstBlock := bm.Offset / uBlockSize
+		lastBlock := (bm.Offset + bm.Length - 1) / uBlockSize
+		for idx := firstBlock; idx <= lastBlock; idx++ {
+			seen[idx] = struct{}{}
+		}
+	}
+
+	// A header that maps every block to its own BuildId has no base layer — it
+	// is a full image (a base / build template), not a pause diff. Returning
+	// its "own" blocks would prefetch the entire template on a normal
+	// create/build resume, so only treat own-BuildId blocks as a last-cycle
+	// diff when the snapshot is layered on a distinct base.
+	if !hasBase || len(seen) == 0 {
+		return nil
+	}
+
+	indices := make([]uint64, 0, len(seen))
+	for idx := range seen {
+		indices = append(indices, idx)
+	}
+	slices.Sort(indices)
+
+	return &metadata.MemoryPrefetchMapping{
+		Indices:   indices,
+		BlockSize: blockSize,
+	}
+}
@@ -0,0 +1,224 @@
+//go:build linux
+
+package sandbox
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/metadata"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+)
+
+const testBlockSize = 4096
+
+// buildTestHeader assembles a header.Header at testBlockSize for the given
+// own build ID and raw BuildMap entries (deliberately allowed to overlap,
+// since real merged/dedup headers can have overlapping entries across
+// layers).
+func buildTestHeader(t *testing.T, own uuid.UUID, entries []header.BuildMap) *header.Header {
+	t.Helper()
+
+	mapping, err := header.NewMapping(testBlockSize, entries)
+	require.NoError(t, err)
+
+	return &header.Header{
+		Metadata: &header.Metadata{
+			BlockSize: testBlockSize,
+			BuildId:   own,
+		},
+		Mapping: mapping,
+	}
+}
+
+func TestBuildDiffMemoryPrefetchMapping(t *testing.T) {
+	t.Parallel()
+
+	own := uuid.New()
+	other := uuid.New()
+
+	t.Run("selects only own-build blocks and excludes other builds", func(t *testing.T) {
+		t.Parallel()
+
+		h := buildTestHeader(t, own, []header.BuildMap{
+			{Offset: 0, Length: testBlockSize, BuildId: own},
+			{Offset: 2 * testBlockSize, Length: testBlockSize, BuildId: other},
+			{Offset: 3 * testBlockSize, Length: testBlockSize, BuildId: own},
+		})
+
+		got := buildDiffMemoryPrefetchMapping(h)
+
+		require.NotNil(t, got)
+		assert.Equal(t, []uint64{0, 3}, got.Indices, "only own-BuildId blocks, offset order")
+		assert.Equal(t, int64(testBlockSize), got.BlockSize)
+	})
+
+	t.Run("dedups a block index covered by multiple own-build entries", func(t *testing.T) {
+		t.Parallel()
+
+		// Simulates a merged header where the same block offset is covered by
+		// more than one BuildMap entry attributed to this build (e.g.
+		// overlapping ranges folded in across layers): block index 1 is
+		// covered by both own entries below and must be counted once. The
+		// trailing base entry makes this a diff (layered on a base), not a full
+		// image.
+		h := buildTestHeader(t, own, []header.BuildMap{
+			{Offset: 0, Length: 2 * testBlockSize, BuildId: own},             // blocks 0, 1
+			{Offset: testBlockSize, Length: 2 * testBlockSize, BuildId: own}, // blocks 1, 2
+			{Offset: 5 * testBlockSize, Length: testBlockSize, BuildId: other},
+		})
+
+		got := buildDiffMemoryPrefetchMapping(h)
+
+		require.NotNil(t, got)
+		assert.Equal(t, []uint64{0, 1, 2}, got.Indices, "each block index counted exactly once")
+	})
+
+	t.Run("full image with no base layer (base/build template) returns nil", func(t *testing.T) {
+		t.Parallel()
+
+		// Every block owned by this build's own ID and no distinct base layer:
+		// a base/build template, not a pause diff. Prefetching all of it would
+		// pull the whole template image, so it must be skipped.
+		h := buildTestHeader(t, own, []header.BuildMap{
+			{Offset: 0, Length: 3 * testBlockSize, BuildId: own},
+		})
+
+		assert.Nil(t, buildDiffMemoryPrefetchMapping(h))
+	})
+
+	t.Run("nil header returns nil", func(t *testing.T) {
+		t.Parallel()
+
+		assert.Nil(t, buildDiffMemoryPrefetchMapping(nil))
+	})
+
+	t.Run("nil metadata returns nil", func(t *testing.T) {
+		t.Parallel()
+
+		assert.Nil(t, buildDiffMemoryPrefetchMapping(&header.Header{Metadata: nil}))
+	})
+
+	t.Run("zero block size returns nil", func(t *testing.T) {
+		t.Parallel()
+
+		h := &header.Header{Metadata: &header.Metadata{BlockSize: 0, BuildId: own}}
+		assert.Nil(t, buildDiffMemoryPrefetchMapping(h))
+	})
+
+	t.Run("empty diff (no own-build entries) returns nil", func(t *testing.T) {
+		t.Parallel()
+
+		h := buildTestHeader(t, own, []header.BuildMap{
+			{Offset: 0, Length: testBlockSize, BuildId: other},
+		})
+
+		assert.Nil(t, buildDiffMemoryPrefetchMapping(h))
+	})
+
+	t.Run("empty mapping returns nil", func(t *testing.T) {
+		t.Parallel()
+
+		h := buildTestHeader(t, own, nil)
+		assert.Nil(t, buildDiffMemoryPrefetchMapping(h))
+	})
+}
+
+func TestSelectResumePrefetch(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		source          string
+		init, lastCycle bool
+	}{
+		{"off", false, false},
+		{"init", true, false},
+		{"last-cycle", false, true},
+		{"both", true, true},
+		{"", true, false},        // unknown -> init (default)
+		{"garbage", true, false}, // unknown -> init (default)
+	}
+	for _, c := range cases {
+		i, f := selectResumePrefetch(c.source)
+		assert.Equalf(t, c.init, i, "useInit for source %q", c.source)
+		assert.Equalf(t, c.lastCycle, f, "useLastCycle for source %q", c.source)
+	}
+}
+
+func TestCapResumePrefetch(t *testing.T) {
+	t.Parallel()
+
+	const bs = int64(2 * 1024 * 1024)
+	// 5 blocks == 10 MiB.
+	m := &metadata.MemoryPrefetchMapping{Indices: []uint64{0, 1, 2, 3, 4}, BlockSize: bs}
+
+	t.Run("uncapped passes through unchanged", func(t *testing.T) {
+		t.Parallel()
+		assert.Same(t, m, capResumePrefetch(m, -1))
+	})
+
+	t.Run("cap larger than the set passes through unchanged", func(t *testing.T) {
+		t.Parallel()
+		assert.Same(t, m, capResumePrefetch(m, 100))
+	})
+
+	t.Run("cap truncates to the earliest blocks in offset order", func(t *testing.T) {
+		t.Parallel()
+		got := capResumePrefetch(m, 6) // 6 MiB / 2 MiB = 3 blocks
+		require.NotNil(t, got)
+		assert.Equal(t, []uint64{0, 1, 2}, got.Indices)
+		assert.Equal(t, bs, got.BlockSize)
+		assert.Equal(t, []uint64{0, 1, 2, 3, 4}, m.Indices, "original mapping is not mutated")
+	})
+
+	t.Run("zero cap truncates to an empty set", func(t *testing.T) {
+		t.Parallel()
+		got := capResumePrefetch(m, 0)
+		require.NotNil(t, got)
+		assert.Empty(t, got.Indices, "0 MiB keeps no blocks; everything demand-faults")
+		assert.Equal(t, bs, got.BlockSize)
+	})
+
+	t.Run("nil mapping stays nil", func(t *testing.T) {
+		t.Parallel()
+		assert.Nil(t, capResumePrefetch(nil, 4))
+	})
+
+	t.Run("unknown block size passes through unchanged", func(t *testing.T) {
+		t.Parallel()
+		z := &metadata.MemoryPrefetchMapping{Indices: []uint64{0, 1}, BlockSize: 0}
+		assert.Same(t, z, capResumePrefetch(z, 1))
+	})
+}
+
+// TestBuildDiffMemoryPrefetchMappingUnalignedEntry pins the fix for a dedup
+// entry aligned to a smaller granule than Metadata.BlockSize that straddles a
+// block boundary: enumerating by block index must record every block the entry
+// covers, not just the first. BlockSize is 8192 while the own-build entry
+// (offset 4096, length 8192) spans bytes 4096–12287 — blocks 0 and 1 — so
+// byte-stepping from the unaligned offset would drop block 1.
+func TestBuildDiffMemoryPrefetchMappingUnalignedEntry(t *testing.T) {
+	t.Parallel()
+
+	own := uuid.New()
+	other := uuid.New()
+
+	mapping, err := header.NewMapping(4096, []header.BuildMap{
+		{Offset: 0, Length: 4096, BuildId: other}, // base layer, so this reads as a diff
+		{Offset: 4096, Length: 8192, BuildId: own},
+	})
+	require.NoError(t, err)
+
+	h := &header.Header{
+		Metadata: &header.Metadata{BlockSize: 8192, BuildId: own},
+		Mapping:  mapping,
+	}
+
+	got := buildDiffMemoryPrefetchMapping(h)
+	require.NotNil(t, got)
+	assert.Equal(t, []uint64{0, 1}, got.Indices, "the tail block of a cross-boundary entry must not be dropped")
+	assert.Equal(t, int64(8192), got.BlockSize)
+}