diff --git a/packages/orchestrator/pkg/sandbox/resume_prefetch.go b/packages/orchestrator/pkg/sandbox/resume_prefetch.go new file mode 100644 index 0000000000..a5ea8be2ed --- /dev/null +++ b/packages/orchestrator/pkg/sandbox/resume_prefetch.go @@ -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, + } +} diff --git a/packages/orchestrator/pkg/sandbox/resume_prefetch_test.go b/packages/orchestrator/pkg/sandbox/resume_prefetch_test.go new file mode 100644 index 0000000000..9eb98dfe17 --- /dev/null +++ b/packages/orchestrator/pkg/sandbox/resume_prefetch_test.go @@ -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) +} diff --git a/packages/orchestrator/pkg/sandbox/sandbox.go b/packages/orchestrator/pkg/sandbox/sandbox.go index faeaaf5dcf..4580ca2913 100644 --- a/packages/orchestrator/pkg/sandbox/sandbox.go +++ b/packages/orchestrator/pkg/sandbox/sandbox.go @@ -757,7 +757,26 @@ func (f *Factory) ResumeSandbox( return uffd.New(memfile, fcUffdPath), nil }) - // Prefetching + // Prefetching. Derive the prefetch context and register its cancel with the + // cleanup manager *synchronously*, before the goroutine below. execCtx is + // non-cancelable (context.WithoutCancel) and the fetch-only last-cycle path + // has no copy worker to observe uffd close, so without an explicit cancel a + // torn-down sandbox would keep draining a (potentially multi-GiB) diff from + // object storage. Registering here rather than inside the goroutine also + // avoids racing cleanup.Run: a goroutine-side Add could lose the hasRun + // check to a concurrent teardown and never register. + // + // Register as PRIORITY so teardown aborts the fetch first: priority handlers + // run before the normal cleanup list, and (LIFO) this one runs before the + // priority Stop — otherwise the fetchers keep issuing large memfile reads + // through Stop and the rest of the normal cleanup until a late cancel. + prefetchCtx, cancelPrefetch := context.WithCancel(execCtx) + cleanup.AddPriority(ctx, func(context.Context) error { + cancelPrefetch() + + return nil + }) + go func() { memfile, err := t.Memfile(ctx) if err != nil { @@ -771,31 +790,85 @@ func (f *Factory) ResumeSandbox( telemetry.ReportEvent(ctx, "got metadata") - // Start background prefetcher as early as possible if prefetch mapping exists - // Fetching from source starts immediately; copying waits for uffd to be ready - if meta.Prefetch != nil && meta.Prefetch.Memory != nil { - fcUffd, err := uffdPromise.Wait(ctx) - if err != nil { - return - } + // Start background prefetchers as early as possible. Fetching from + // source starts immediately; copying (when prefaulting) waits for uffd. + // + // Up to two independent mappings are replayed on a resume, chosen by the + // resume-prefetch-source flag (see selectResumePrefetch): + // - The init trace (meta.Prefetch.Memory): the build-time + // create-from-template / checkpoint read-hot startup working set. + // Prefaulted, exactly as today. + // - The last-cycle diff: the pages this sandbox's last resume→pause + // cycle wrote — its own pause diff, derived from the memfile header + // (see buildDiffMemoryPrefetchMapping) — a good predictor of the + // next cycle's working set. Replayed FETCH-ONLY: it warms the cache + // and lets the guest fault, because prefaulting a multi-GiB diff + // would load UFFDIO_COPY onto the resume-critical path for no + // workload gain and would regress warm resumes. + // The default source "init" selects only the init trace, preserving + // today's behavior. Pause/resume normally has no init trace + // (SameVersionTemplate drops it), so with source=last-cycle/both the + // last-cycle diff usually runs alone. When both exist, the small init + // trace runs first and last-cycle follows it (a barrier), keeping the + // large last-cycle fetch off the resume-critical path. + source := f.featureFlags.StringFlag(ctx, featureflags.ResumePrefetchSourceFlag, sandboxLDContext(runtime, config)) + useInit, useLastCycle := selectResumePrefetch(source) + + var initMapping *metadata.MemoryPrefetchMapping + if useInit && meta.Prefetch != nil { + initMapping = meta.Prefetch.Memory + } - telemetry.ReportEvent(ctx, "starting prefetcher") - l := logger.L().With(logger.WithSandboxID(runtime.SandboxID), logger.WithTemplateID(runtime.TemplateID), logger.WithTeamID(runtime.TeamID)) + var lastCycleMapping *metadata.MemoryPrefetchMapping + if useLastCycle { + lastCycleMapping = buildDiffMemoryPrefetchMapping(memfile.Header()) + // Bound the last-cycle volume against the shared object-store pool; + // resume-last-cycle-prefetch-max-mib=-1 (default) is uncapped. + maxMiB := f.featureFlags.IntFlag(ctx, featureflags.ResumeLastCyclePrefetchMaxMiBFlag, sandboxLDContext(runtime, config)) + lastCycleMapping = capResumePrefetch(lastCycleMapping, maxMiB) + } - go func() { - p := prefetch.New( - l, - memfile, - fcUffd, - meta.Prefetch.Memory, - f.featureFlags, - ) - err := p.Start(execCtx) - if err != nil { - l.Error(ctx, "failed to start prefetcher", zap.Error(err)) - } - }() + // Record the chosen source and the sizes it resolved to, so a resume can + // be cohorted by prefetch source (guards against flag misconfiguration) + // and the last-cycle set size is visible per resume. + execSpan.SetAttributes( + attribute.String("resume.prefetch.source", source), + attribute.Int("resume.prefetch.init_blocks", initMapping.Count()), + attribute.Int("resume.prefetch.last_cycle_blocks", lastCycleMapping.Count()), + ) + + if initMapping == nil && lastCycleMapping == nil { + return } + + fcUffd, err := uffdPromise.Wait(ctx) + if err != nil { + return + } + + telemetry.ReportEvent(ctx, "starting prefetcher") + l := logger.L().With(logger.WithSandboxID(runtime.SandboxID), logger.WithTemplateID(runtime.TemplateID), logger.WithTeamID(runtime.TeamID)) + + go func() { + // Init trace first, prefaulted (prod behavior). Start blocks until + // its fetch+copy complete, so it acts as a barrier before the + // last-cycle fetch begins. + if initMapping != nil { + p := prefetch.New(l, memfile, fcUffd, initMapping, f.featureFlags) + if err := p.Start(prefetchCtx); err != nil { + l.Error(ctx, "failed to start init prefetcher", zap.Error(err)) + } + } + + // Last-cycle diff, fetch-only. + if lastCycleMapping != nil { + p := prefetch.New(l, memfile, fcUffd, lastCycleMapping, f.featureFlags) + p.Prefault = false + if err := p.Start(prefetchCtx); err != nil { + l.Error(ctx, "failed to start last-cycle prefetcher", zap.Error(err)) + } + } + }() }() // Slot initialization diff --git a/packages/orchestrator/pkg/sandbox/uffd/prefetch/metrics.go b/packages/orchestrator/pkg/sandbox/uffd/prefetch/metrics.go index 3c719d52e4..79b41d6141 100644 --- a/packages/orchestrator/pkg/sandbox/uffd/prefetch/metrics.go +++ b/packages/orchestrator/pkg/sandbox/uffd/prefetch/metrics.go @@ -37,7 +37,22 @@ var durationHistogram = utils.Must(meter.Int64Histogram( metric.WithUnit("ms"), )) +// mappingBlocksHistogram records the size (in blocks) of each prefetch mapping +// replayed on a resume, tagged by mode ("prefault" for the init trace, "fetch" +// for the fetch-only last-cycle diff). It gives the fleet distribution of the +// recorded working-set size; the "fetch" series sizes the last-cycle-prefetch +// rollout (how much a resume prefetches) and flags idle-at-pause sandboxes in +// its bottom bucket. +var mappingBlocksHistogram = utils.Must(meter.Int64Histogram( + "orchestrator.sandbox.uffd.prefetch.mapping_blocks", + metric.WithDescription("Prefetch mapping size in blocks, by mode"), + metric.WithUnit("{block}"), +)) + var ( + modePrefaultAttr = telemetry.PrecomputeAttrs(attribute.String("mode", "prefault")) + modeFetchAttr = telemetry.PrecomputeAttrs(attribute.String("mode", "fetch")) + stageFetchedAttr = telemetry.PrecomputeAttrs(attribute.String("stage", "fetched")) stageFetchSkippedAttr = telemetry.PrecomputeAttrs(attribute.String("stage", "fetch_skipped")) stageCopiedAttr = telemetry.PrecomputeAttrs(attribute.String("stage", "copied")) diff --git a/packages/orchestrator/pkg/sandbox/uffd/prefetch/prefetcher.go b/packages/orchestrator/pkg/sandbox/uffd/prefetch/prefetcher.go index 14e85b1c45..aa32dfa0ba 100644 --- a/packages/orchestrator/pkg/sandbox/uffd/prefetch/prefetcher.go +++ b/packages/orchestrator/pkg/sandbox/uffd/prefetch/prefetcher.go @@ -21,6 +21,7 @@ import ( "github.com/e2b-dev/infra/packages/shared/pkg/logger" "github.com/e2b-dev/infra/packages/shared/pkg/storage" "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" + "github.com/e2b-dev/infra/packages/shared/pkg/units" ) var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/prefetch") @@ -30,6 +31,38 @@ type prefetchData struct { data []byte } +// extent is a run of contiguous block indices fetched from the source in a +// single source.Slice call. Coalescing contiguous indices into one extent +// turns many small sequential reads into fewer, larger ones. The copy phase +// always stays per-block (split out of the fetched extent), since UFFDIO_COPY +// requires page-sized data. +type extent struct { + startIdx uint64 + blocks int +} + +// coalesceIndices groups sorted, deduplicated block indices into maximal +// contiguous runs, each capped at maxBlocks. With maxBlocks<=1 every extent is +// a single block, reproducing the per-block fetch exactly (coalescing off). +func coalesceIndices(indices []uint64, maxBlocks int) []extent { + if maxBlocks < 1 { + maxBlocks = 1 + } + + out := make([]extent, 0, len(indices)) + for i := 0; i < len(indices); { + n := 1 + for i+n < len(indices) && indices[i+n] == indices[i+n-1]+1 && n < maxBlocks { + n++ + } + + out = append(out, extent{startIdx: indices[i], blocks: n}) + i += n + } + + return out +} + // Prefetcher handles background prefetching of memory pages. // It proactively fetches pages that are known to be needed based on the prefetch mapping // collected during template build. @@ -40,10 +73,17 @@ type prefetchData struct { // // Both phases run with their own parallelism limits and don't block each other. type Prefetcher struct { - logger logger.Logger - source block.Slicer - uffd uffd.MemoryBackend - mapping *metadata.MemoryPrefetchMapping + logger logger.Logger + source block.Slicer + uffd uffd.MemoryBackend + mapping *metadata.MemoryPrefetchMapping + // Prefault installs each fetched page into the guest via UFFDIO_COPY (the + // copy phase). New defaults it to true. Set it false for a "fetch-only" + // run: the prefetcher only warms the shared chunk cache and the guest still + // faults (served fast from the warm cache). Fetch-only is used for the + // last-cycle diff, whose multi-GiB prefault would load UFFDIO_COPY onto the + // resume-critical path for no workload gain and would regress warm resumes. + Prefault bool featureFlags *featureflags.Client } @@ -59,6 +99,7 @@ func New( source: source, uffd: uffd, mapping: mapping, + Prefault: true, featureFlags: featureFlags, } } @@ -89,9 +130,10 @@ func (p *Prefetcher) Start(ctx context.Context) error { return nil } - // Get worker counts from feature flags at runtime + // Get worker counts and coalescing config from feature flags at runtime maxFetchWorkers := p.featureFlags.IntFlag(ctx, featureflags.MemoryPrefetchMaxFetchWorkers) maxCopyWorkers := p.featureFlags.IntFlag(ctx, featureflags.MemoryPrefetchMaxCopyWorkers) + coalesceMaxMB := p.featureFlags.IntFlag(ctx, featureflags.MemoryPrefetchCoalesceMaxMB) // cancelRun aborts the whole run early. Copy workers fire it on ErrClosed // (uffd gone: sandbox teardown) so fetch workers stop fetching and @@ -103,11 +145,22 @@ func (p *Prefetcher) Start(ctx context.Context) error { blockSize := p.mapping.BlockSize totalPages := len(indices) + // Record the recorded working-set size for this run, tagged by mode so the + // fetch-only (last-cycle) distribution is separable from the prefaulted + // init trace. Available for rollout dashboards without Tempo sampling. + modeAttr := modeFetchAttr + if p.Prefault { + modeAttr = modePrefaultAttr + } + mappingBlocksHistogram.Record(ctx, int64(totalPages), modeAttr) + span.SetAttributes( attribute.Int64("prefetch.total_pages", int64(totalPages)), attribute.Int64("prefetch.block_size", blockSize), attribute.Int("prefetch.max_fetch_workers", maxFetchWorkers), attribute.Int("prefetch.max_copy_workers", maxCopyWorkers), + attribute.Int("prefetch.coalesce_max_mb", coalesceMaxMB), + attribute.Bool("prefetch.prefault", p.Prefault), ) p.logger.Debug(ctx, "prefetch: starting background prefetch", @@ -115,13 +168,28 @@ func (p *Prefetcher) Start(ctx context.Context) error { zap.Int64("block_size", blockSize), zap.Int("max_fetch_workers", maxFetchWorkers), zap.Int("max_copy_workers", maxCopyWorkers), + zap.Int("coalesce_max_mb", coalesceMaxMB), ) + // Coalesce contiguous blocks into extents. maxBlk<=1 (coalesceMaxMB<=0, + // the default) reproduces the per-block fetch exactly. + maxBlk := 1 + if coalesceMaxMB > 0 && blockSize > 0 { + maxBlk = max(1, int(units.MBToBytes(int64(coalesceMaxMB))/blockSize)) + } + extents := coalesceIndices(indices, maxBlk) + span.SetAttributes(attribute.Int("prefetch.extents", len(extents))) + // Channels for work distribution - // Fetch channel: all offsets to fetch (large buffer so main goroutine doesn't block) - fetchCh := make(chan int64, totalPages) - // Copy channel: offsets ready to copy (fetched successfully) - copyCh := make(chan prefetchData, totalPages) + // Fetch channel: all extents to fetch (large buffer so main goroutine doesn't block) + fetchCh := make(chan extent, len(extents)) + // Copy channel: pages fetched successfully, ready to install. Only created + // when prefaulting; a fetch-only run leaves it nil and skips the copy phase, + // so fetched bytes are dropped after warming the cache (no in-memory pile-up). + var copyCh chan prefetchData + if p.Prefault { + copyCh = make(chan prefetchData, totalPages) + } // Counters for statistics var fetchedCount atomic.Uint64 @@ -139,9 +207,9 @@ func (p *Prefetcher) Start(ctx context.Context) error { // duration below is immune to wall-clock steps. var copyStart atomic.Pointer[time.Time] - // Queue all offsets to fetch in the order they were originally accessed - for _, idx := range indices { - fetchCh <- header.BlockOffset(int64(idx), blockSize) + // Queue all extents to fetch, in offset order. + for _, e := range extents { + fetchCh <- e } close(fetchCh) @@ -152,16 +220,21 @@ func (p *Prefetcher) Start(ctx context.Context) error { }) } - // Start copy coordinator - waits for uffd ready, then spawns copy workers - copyWg.Go(func() { - p.startCopyWorkers(ctx, cancelRun, copyCh, maxCopyWorkers, ©Start, &copiedCount, ©SkippedCount) - }) + // Start copy coordinator - waits for uffd ready, then spawns copy workers. + // Skipped entirely for fetch-only runs. + if p.Prefault { + copyWg.Go(func() { + p.startCopyWorkers(ctx, cancelRun, copyCh, maxCopyWorkers, ©Start, &copiedCount, ©SkippedCount) + }) + } // Wait for fetch workers to complete fetchWg.Wait() fetchDuration := time.Since(runStart) // Close copy channel when all fetch workers are done - close(copyCh) + if copyCh != nil { + close(copyCh) + } // Wait for copy workers to complete copyWg.Wait() @@ -223,10 +296,15 @@ func (p *Prefetcher) startCopyWorkers( copyWorkerWg.Wait() } -// fetchWorker fetches pages from the source to populate the cache +// fetchWorker fetches extents from the source to populate the cache. An +// extent is one contiguous run of blocks fetched in a single source.Slice +// call; with coalescing off (the default) every extent is a single block, +// reproducing the plain per-block fetch. The copy phase always stays +// per-block: an extent is split into page-sized sub-slices, since UFFDIO_COPY +// requires page-sized data. func (p *Prefetcher) fetchWorker( ctx context.Context, - fetchCh <-chan int64, + fetchCh <-chan extent, copyCh chan<- prefetchData, blockSize int64, fetchedCount *atomic.Uint64, @@ -236,29 +314,58 @@ func (p *Prefetcher) fetchWorker( select { case <-ctx.Done(): return - case offset, ok := <-fetchCh: + case e, ok := <-fetchCh: if !ok { return } - // Fetch from source - this populates the cache - data, err := p.source.Slice(ctx, offset, blockSize) + baseOffset := header.BlockOffset(int64(e.startIdx), blockSize) + + // Fetch from source - this populates the cache. A multi-block + // extent is one larger sequential read spanning e.blocks pages. + data, err := p.source.Slice(ctx, baseOffset, blockSize*int64(e.blocks)) if err != nil { - p.logger.Debug(ctx, "prefetch: failed to fetch page", - zap.Int64("offset", offset), + p.logger.Debug(ctx, "prefetch: failed to fetch extent", + zap.Int64("offset", baseOffset), + zap.Int("blocks", e.blocks), zap.Error(err), ) - skippedCount.Add(1) + skippedCount.Add(uint64(e.blocks)) continue } - fetchedCount.Add(1) + // Guard against a short read with a nil error: the per-page slicing + // below indexes up to blockSize*e.blocks and would panic (slice + // bounds out of range) — which crashes the whole orchestrator — if + // the source returned fewer bytes than requested. + if int64(len(data)) < blockSize*int64(e.blocks) { + p.logger.Debug(ctx, "prefetch: short read from source", + zap.Int64("offset", baseOffset), + zap.Int("blocks", e.blocks), + zap.Int("got_bytes", len(data)), + ) + skippedCount.Add(uint64(e.blocks)) - // Queue for copy (non-blocking - channel has enough capacity) - select { - case copyCh <- prefetchData{offset: offset, data: data}: - case <-ctx.Done(): + continue + } + + fetchedCount.Add(uint64(e.blocks)) + + // Queue each page for copy (non-blocking - channel has enough + // capacity). Fetch-only runs pass a nil copyCh: warming the cache is + // the whole job, so the fetched bytes are dropped here. + if copyCh == nil { + continue + } + for b := range e.blocks { + off := baseOffset + int64(b)*blockSize + sub := data[int64(b)*blockSize : int64(b+1)*blockSize] + select { + case copyCh <- prefetchData{offset: off, data: sub}: + case <-ctx.Done(): + return + } } } } diff --git a/packages/orchestrator/pkg/sandbox/uffd/prefetch/prefetcher_test.go b/packages/orchestrator/pkg/sandbox/uffd/prefetch/prefetcher_test.go index a0c15d3e97..f83a1db9cc 100644 --- a/packages/orchestrator/pkg/sandbox/uffd/prefetch/prefetcher_test.go +++ b/packages/orchestrator/pkg/sandbox/uffd/prefetch/prefetcher_test.go @@ -12,7 +12,10 @@ import ( "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd" "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/userfaultfd" + "github.com/e2b-dev/infra/packages/orchestrator/pkg/template/metadata" + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" "github.com/e2b-dev/infra/packages/shared/pkg/logger" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" ) // prefaultResult scripts one Prefault call of fakeBackend. @@ -125,3 +128,234 @@ func TestCopyWorkerCountsOnlyInstalledAsCopied(t *testing.T) { require.EqualValues(t, 2, copied.Load()) require.EqualValues(t, 2, skipped.Load()) } + +// TestCoalesceIndices covers the extent grouping used to turn many small +// per-block fetches into fewer, larger sequential ones. +func TestCoalesceIndices(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + indices []uint64 + maxBlocks int + want []extent + }{ + { + name: "a contiguous run coalesces into one extent", + indices: []uint64{10, 11, 12, 13}, + maxBlocks: 8, + want: []extent{{startIdx: 10, blocks: 4}}, + }, + { + name: "a gap breaks the run into separate extents", + indices: []uint64{10, 11, 20, 21, 22}, + maxBlocks: 8, + want: []extent{{startIdx: 10, blocks: 2}, {startIdx: 20, blocks: 3}}, + }, + { + name: "maxBlocks caps a long run into multiple extents", + indices: []uint64{0, 1, 2, 3, 4, 5}, + maxBlocks: 2, + want: []extent{{startIdx: 0, blocks: 2}, {startIdx: 2, blocks: 2}, {startIdx: 4, blocks: 2}}, + }, + { + name: "maxBlocks==1 is one extent per index (coalescing off)", + indices: []uint64{5, 6, 9}, + maxBlocks: 1, + want: []extent{{startIdx: 5, blocks: 1}, {startIdx: 6, blocks: 1}, {startIdx: 9, blocks: 1}}, + }, + { + name: "maxBlocks<1 is clamped to 1", + indices: []uint64{1, 2}, + maxBlocks: 0, + want: []extent{{startIdx: 1, blocks: 1}, {startIdx: 2, blocks: 1}}, + }, + { + name: "no indices yields no extents", + indices: nil, + maxBlocks: 4, + want: []extent{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got := coalesceIndices(tt.indices, tt.maxBlocks) + require.Equal(t, tt.want, got) + }) + } +} + +// fakeSlicer is a deterministic, in-memory block.Slicer: Slice returns +// `length` bytes where byte i equals byte(off+i), so a fetched extent's data +// can be checked against the offset each page should carry. +type fakeSlicer struct { + blockSize int64 +} + +func (f *fakeSlicer) Slice(_ context.Context, off, length int64) ([]byte, error) { + data := make([]byte, length) + for i := range data { + data[i] = byte(off + int64(i)) + } + + return data, nil +} + +func (f *fakeSlicer) BlockSize() int64 { return f.blockSize } + +// TestFetchWorkerSplitsCoalescedExtentIntoPerPageCopies is the prefetcher-level +// check that coalescing the fetch (one source.Slice call spanning several +// blocks) doesn't change what gets queued for copy: the copy phase must still +// see exactly one page-sized prefetchData per block, each carrying that +// block's own bytes at its own offset — never the whole coalesced extent. +func TestFetchWorkerSplitsCoalescedExtentIntoPerPageCopies(t *testing.T) { + t.Parallel() + + const blockSize = 4096 + + log, err := logger.NewDevelopmentLogger() + require.NoError(t, err) + + p := &Prefetcher{logger: log, source: &fakeSlicer{blockSize: blockSize}} + + const startIdx = 5 + const blocks = 3 + + fetchCh := make(chan extent, 1) + fetchCh <- extent{startIdx: startIdx, blocks: blocks} + close(fetchCh) + + copyCh := make(chan prefetchData, blocks) + + var fetched, skipped atomic.Uint64 + p.fetchWorker(t.Context(), fetchCh, copyCh, blockSize, &fetched, &skipped) + close(copyCh) + + require.EqualValues(t, blocks, fetched.Load(), "one coalesced extent still counts as its full block count") + require.Zero(t, skipped.Load()) + + var got []prefetchData + for d := range copyCh { + got = append(got, d) + } + require.Len(t, got, blocks, "the copy phase gets one entry per block, not one per extent") + + for i, d := range got { + wantOffset := header.BlockOffset(int64(startIdx+i), blockSize) + require.Equal(t, wantOffset, d.offset) + require.Len(t, d.data, blockSize, "copy data must be exactly one page, never a multi-block extent") + require.Equal(t, byte(wantOffset), d.data[0], "each page must carry its own block's bytes") + } +} + +// TestFetchWorkerFetchOnlyQueuesNothing checks the fetch-only path (Prefault +// off, nil copyCh): the worker still fetches every block to warm the cache but +// queues nothing for copy and does not block, so the guest is left to fault the +// (now-warm) pages itself. +func TestFetchWorkerFetchOnlyQueuesNothing(t *testing.T) { + t.Parallel() + + const blockSize = 4096 + + log, err := logger.NewDevelopmentLogger() + require.NoError(t, err) + + p := &Prefetcher{logger: log, source: &fakeSlicer{blockSize: blockSize}} + + fetchCh := make(chan extent, 1) + fetchCh <- extent{startIdx: 2, blocks: 3} + close(fetchCh) + + var fetched, skipped atomic.Uint64 + // nil copyCh == fetch-only. Must return (not deadlock) and count fetches. + p.fetchWorker(t.Context(), fetchCh, nil, blockSize, &fetched, &skipped) + + require.EqualValues(t, 3, fetched.Load(), "fetch-only still fetches every block to warm the cache") + require.Zero(t, skipped.Load()) +} + +// countingBackend records how many times Prefault is called (it must be zero +// for a fetch-only run). +type countingBackend struct { + uffd.MemoryBackend + + prefaults atomic.Int64 +} + +func (b *countingBackend) Prefault(context.Context, int64, []byte) (bool, error) { + b.prefaults.Add(1) + + return true, nil +} + +// TestStartFetchOnlyNeverPrefaults drives the whole Start() with Prefault off +// and asserts the D6 guarantee at the integration seam: the copy phase is +// skipped entirely (copyCh nil, no copy coordinator), so the backend is never +// prefaulted even though every block is fetched to warm the cache. +func TestStartFetchOnlyNeverPrefaults(t *testing.T) { + t.Parallel() + + log, err := logger.NewDevelopmentLogger() + require.NoError(t, err) + + ff, err := featureflags.NewClient() + require.NoError(t, err) + t.Cleanup(func() { _ = ff.Close(t.Context()) }) + + const bs = int64(4096) + be := &countingBackend{} + p := &Prefetcher{ + logger: log, + source: &fakeSlicer{blockSize: bs}, + uffd: be, + mapping: &metadata.MemoryPrefetchMapping{Indices: []uint64{0, 1, 2, 3}, BlockSize: bs}, + Prefault: false, + featureFlags: ff, + } + + require.NoError(t, p.Start(t.Context())) + require.Zero(t, be.prefaults.Load(), "fetch-only run must never call Prefault") +} + +// shortSlicer returns one byte fewer than requested with a nil error, modelling +// a truncated/partial source read. +type shortSlicer struct{ blockSize int64 } + +func (s *shortSlicer) Slice(_ context.Context, _, length int64) ([]byte, error) { + return make([]byte, length-1), nil +} + +func (s *shortSlicer) BlockSize() int64 { return s.blockSize } + +// TestFetchWorkerShortReadDoesNotPanic checks that a short read with a nil +// error is skipped, not sliced into per-page copies — which would panic (slice +// bounds out of range) and crash the orchestrator. +func TestFetchWorkerShortReadDoesNotPanic(t *testing.T) { + t.Parallel() + + const blockSize = 4096 + + log, err := logger.NewDevelopmentLogger() + require.NoError(t, err) + + p := &Prefetcher{logger: log, source: &shortSlicer{blockSize: blockSize}} + + fetchCh := make(chan extent, 1) + fetchCh <- extent{startIdx: 0, blocks: 2} + close(fetchCh) + + copyCh := make(chan prefetchData, 2) + + var fetched, skipped atomic.Uint64 + require.NotPanics(t, func() { + p.fetchWorker(t.Context(), fetchCh, copyCh, blockSize, &fetched, &skipped) + }) + close(copyCh) + + require.Zero(t, fetched.Load(), "a short read is not counted as fetched") + require.EqualValues(t, 2, skipped.Load(), "both blocks of the extent are skipped") + require.Empty(t, copyCh, "nothing is queued for copy on a short read") +} diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index 1c4f02ffe4..b93150ffa2 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -318,6 +318,36 @@ var ( // Copy uses uffd syscalls, so we limit parallelism to avoid overwhelming the system. MemoryPrefetchMaxCopyWorkers = NewIntFlag("memory-prefetch-max-copy-workers", 8) + // MemoryPrefetchCoalesceMaxMB caps how many contiguous prefetch blocks are + // merged into a single source.Slice fetch (in MiB of extent size). 0 + // disables coalescing: every block is fetched individually, matching + // today's behavior. The copy phase is unaffected either way — it always + // installs one page at a time, since UFFDIO_COPY requires page-sized data. + MemoryPrefetchCoalesceMaxMB = NewIntFlag("memory-prefetch-coalesce-max-mb", 0) + + // ResumePrefetchSourceFlag selects which trace the resume prefetcher + // replays: + // "init" — only the build-time / harvested read-hot init trace + // (meta.Prefetch.Memory), prefaulted. Preserves today's + // behavior, so this is the default and a no-op-equivalent. + // "last-cycle" — only the sandbox's own pause diff (the pages the last + // resume→pause cycle wrote), derived from the memfile header + // and replayed fetch-only. + // "both" — init first (prefaulted), then last-cycle (fetch-only) behind + // a barrier, so the large last-cycle fetch stays off the + // resume-critical path. + // "off" — kill switch, no resume prefetch. + // Unknown values fall back to "init". + ResumePrefetchSourceFlag = NewStringFlag("resume-prefetch-source", "init") + + // ResumeLastCyclePrefetchMaxMiBFlag caps how much of the last-cycle diff a single + // resume prefetches, in MiB. -1 (the default, negative = no limit per the + // codebase convention) is uncapped; the recorded diff is small by + // construction, so this exists to throttle the heavy-churn tail against the + // shared object-store pool without a redeploy. A non-negative N keeps the + // first N MiB of blocks in offset order and leaves the rest to demand-fault. + ResumeLastCyclePrefetchMaxMiBFlag = NewIntFlag("resume-last-cycle-prefetch-max-mib", -1) + // PauseResumePrefetchHarvestFlag makes the orchestrator, after a pause // snapshot is durable, run a throwaway warm resume of the just-written // artifact (driven by envd /init, workload frozen, egress denied) to record