From 7bab357a90d5e795459b4dcff4005531dcd794f4 Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Wed, 1 Jul 2026 15:48:01 +0100 Subject: [PATCH 1/9] feat(orch): serve resume from in-flight memfd during dedup drain A resume that overlaps an in-flight pause of the same sandbox blocks for the full dedup duration: the deduped diff is the only readable diff source and it is not readable until the background dedup drain finishes, so UFFD memory serving in load-snapshot waits on DedupedMemfdCache.Wait. Give DedupedMemfdCache an in-flight read path. After the dedup compare publishes the metadata, expose the still-mapped memfd together with a packed-to-absolute offset index built from the deduped dirty set. While the drain runs, ReadAt/Slice translate the packed diff offset and copy straight from the memfd instead of blocking on the drain; once the drain completes, reads fall through to the drained cache. An RWMutex guards the memfd so the drain's close cannot unmap it under an in-flight reader, and the done barrier makes the memfd-to-cache handover race-free. Gate the behaviour behind memfd-dedup-inflight-serve (default off); when off the prior wait-for-drain behaviour is preserved. Dedup still runs to completion and the uploaded artifact is unchanged. Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 (1M context) --- .../orchestrator/pkg/sandbox/block/memfd.go | 134 +++++++++++++++++- .../pkg/sandbox/block/memfd_test.go | 2 +- .../orchestrator/pkg/sandbox/fc/memory.go | 3 +- packages/orchestrator/pkg/sandbox/sandbox.go | 3 + packages/shared/pkg/featureflags/flags.go | 6 + 5 files changed, 142 insertions(+), 6 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/block/memfd.go b/packages/orchestrator/pkg/sandbox/block/memfd.go index 7bc59793f0..5ea0794031 100644 --- a/packages/orchestrator/pkg/sandbox/block/memfd.go +++ b/packages/orchestrator/pkg/sandbox/block/memfd.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "os" + "sync" "time" "github.com/RoaringBitmap/roaring/v2" @@ -232,10 +233,23 @@ func (m *MemfdCache) Size() (int64, error) { return m.cache.Size() } // DedupedMemfdCache runs compare+drain on a goroutine; metaOut resolves // after compare, reads against the cache block on done until drain finishes. +// +// When inflight serving is enabled, reads are served directly from the +// still-mapped memfd during the drain window instead of blocking on done, so a +// resume overlapping the pause is not delayed by the dedup drain. type DedupedMemfdCache struct { outPath string cancel context.CancelFunc done *utils.SetOnce[*Cache] + + inflight bool + // mu guards memfd + index. memfd is non-nil only between "compare done" + // (published in runDedup) and "drain done" (closed under the write lock so + // it can't be unmapped under an in-flight reader). index translates a + // packed diff-storage offset back to the absolute memfd offset. + mu sync.RWMutex + memfd *Memfd + index packedIndex } func NewCacheFromMemfdDeduped( @@ -250,21 +264,74 @@ func NewCacheFromMemfdDeduped( budget DedupBudget, inputEmpty *roaring.Bitmap, metaOut *utils.SetOnce[*header.DiffMetadata], + inflightServe bool, ) (*DedupedMemfdCache, error) { if blockSize%header.PageSize != 0 { return nil, fmt.Errorf("diff block size %d not a multiple of dedup page size %d", blockSize, header.PageSize) } drainCtx, cancel := context.WithCancel(context.WithoutCancel(ctx)) d := &DedupedMemfdCache{ - outPath: outPath, - cancel: cancel, - done: utils.NewSetOnce[*Cache](), + outPath: outPath, + cancel: cancel, + done: utils.NewSetOnce[*Cache](), + inflight: inflightServe, } go d.runDedup(drainCtx, base, blockSize, memfd, dirty, bestEffort, directIO, budget, inputEmpty, metaOut) return d, nil } +// packedSeg maps a contiguous run of the packed diff artifact back to its +// absolute (device) memfd offset. absStart+delta is the memfd offset for a +// read at packedStart+delta. +type packedSeg struct { + packedStart int64 + absStart int64 + length int64 +} + +// packedIndex is the ascending-by-packedStart list of dirty runs. It mirrors +// the packing done by dedupDrain (BitsetRanges over pageDirty at PageSize) and +// the BuildStorageOffset assignment in header.CreateMapping, so a packed +// offset resolves to the same absolute offset the deduped header would map to. +type packedIndex []packedSeg + +func buildPackedIndex(pageDirty *roaring.Bitmap) packedIndex { + var idx packedIndex + var packed int64 + for r := range BitsetRanges(pageDirty, header.PageSize) { + idx = append(idx, packedSeg{packedStart: packed, absStart: r.Start, length: r.Size}) + packed += r.Size + } + + return idx +} + +// translate maps [off, off+length) in packed space to an absolute memfd offset. +// It only succeeds when the whole range lies in a single dirty run (which is +// always the case for a single build.File read segment); otherwise the caller +// falls back to waiting for the drained cache. +func (idx packedIndex) translate(off, length int64) (int64, bool) { + lo, hi := 0, len(idx) + for lo < hi { + mid := (lo + hi) / 2 + if idx[mid].packedStart+idx[mid].length <= off { + lo = mid + 1 + } else { + hi = mid + } + } + if lo >= len(idx) { + return 0, false + } + seg := idx[lo] + if off < seg.packedStart || off+length > seg.packedStart+seg.length { + return 0, false + } + + return seg.absStart + (off - seg.packedStart), true +} + func (d *DedupedMemfdCache) runDedup( ctx context.Context, base ReadonlyDevice, @@ -307,10 +374,25 @@ func (d *DedupedMemfdCache) runDedup( attribute.Int64("dedup.header_empty_pages", int64(plan.pageEmpty.GetCardinality()))) logSetOnceErr(ctx, "dedup metaOut", metaOut.SetValue(meta)) + // Publish the memfd + packed→absolute index so a resume overlapping this + // pause serves dirty pages straight from the memfd during the drain window + // instead of blocking on done. Reads take mu.RLock; the close below takes + // mu.Lock, so the memfd is never unmapped under an in-flight reader. + if d.inflight { + d.mu.Lock() + d.index = buildPackedIndex(plan.pageDirty) + d.memfd = memfd + d.mu.Unlock() + } + writeStart := time.Now() cache, err := dedupDrain(ctx, src, plan.pageDirty, blockSize, d.outPath, directIO) writeDur := time.Since(writeStart) - if closeErr := memfd.Close(); closeErr != nil { + d.mu.Lock() + closeErr := memfd.Close() + d.memfd = nil + d.mu.Unlock() + if closeErr != nil { logger.L().Warn(ctx, "close memfd after dedup drain", zap.Error(closeErr)) } @@ -331,7 +413,44 @@ func (d *DedupedMemfdCache) Wait(ctx context.Context) (*Cache, error) { return d.done.WaitWithContext(ctx) } +// tryInflightRead fills b from the memfd if inflight serving is active and the +// drain has not yet finished. Returns ok=false to fall back to the drained +// cache (inflight disabled, drain done, memfd already closed, or the range +// can't be resolved from a single dirty run). +func (d *DedupedMemfdCache) tryInflightRead(b []byte, off int64) (int, bool) { + if !d.inflight { + return 0, false + } + // Drain finished: the cache is authoritative, use it. + select { + case <-d.done.Done: + return 0, false + default: + } + + d.mu.RLock() + defer d.mu.RUnlock() + if d.memfd == nil || d.index == nil { + return 0, false + } + absOff, ok := d.index.translate(off, int64(len(b))) + if !ok { + return 0, false + } + src, err := d.memfd.Slice(absOff, int64(len(b))) + if err != nil { + return 0, false + } + copy(b, src) + + return len(b), true +} + func (d *DedupedMemfdCache) ReadAt(b []byte, off int64) (int, error) { + if n, ok := d.tryInflightRead(b, off); ok { + return n, nil + } + c, err := d.Wait(context.Background()) if err != nil { return 0, err @@ -341,6 +460,13 @@ func (d *DedupedMemfdCache) ReadAt(b []byte, off int64) (int, error) { } func (d *DedupedMemfdCache) Slice(off, length int64) ([]byte, error) { + if d.inflight { + buf := make([]byte, length) + if _, ok := d.tryInflightRead(buf, off); ok { + return buf, nil + } + } + c, err := d.Wait(context.Background()) if err != nil { return nil, err diff --git a/packages/orchestrator/pkg/sandbox/block/memfd_test.go b/packages/orchestrator/pkg/sandbox/block/memfd_test.go index d7e4a66278..5371af7461 100644 --- a/packages/orchestrator/pkg/sandbox/block/memfd_test.go +++ b/packages/orchestrator/pkg/sandbox/block/memfd_test.go @@ -227,7 +227,7 @@ func TestNewCacheFromMemfdDeduped_DetachesCompareAndDrain(t *testing.T) { metaOut := utils.NewSetOnce[*header.DiffMetadata]() cache, err := NewCacheFromMemfdDeduped( ctx, &fakeOriginalDevice{data: baseData}, pageSize, cachePath, memfd, dirty, false, false, - DedupBudget{}, nil, metaOut, + DedupBudget{}, nil, metaOut, false, ) require.NoError(t, err) t.Cleanup(func() { _ = cache.Close() }) diff --git a/packages/orchestrator/pkg/sandbox/fc/memory.go b/packages/orchestrator/pkg/sandbox/fc/memory.go index 58a63488f3..8583c1ce75 100644 --- a/packages/orchestrator/pkg/sandbox/fc/memory.go +++ b/packages/orchestrator/pkg/sandbox/fc/memory.go @@ -83,6 +83,7 @@ func (p *Process) ExportMemory( dedupBudget block.DedupBudget, inputEmpty *roaring.Bitmap, metaOut *utils.SetOnce[*header.DiffMetadata], + dedupInflightServe bool, ) (_ block.DiffSource, e error) { // Resolve metaOut on every sync error so Wait-ers don't hang. Success paths // resolve it inline; the memfd-dedup goroutine owns metaOut after this returns. @@ -99,7 +100,7 @@ func (p *Process) ExportMemory( if memfd != nil { if originalMemfile != nil { return block.NewCacheFromMemfdDeduped(ctx, originalMemfile, blockSize, cachePath, memfd, include, - dedupBestEffort, dedupDirectIO, dedupBudget, inputEmpty, metaOut) + dedupBestEffort, dedupDirectIO, dedupBudget, inputEmpty, metaOut, dedupInflightServe) } var ( src block.DiffSource diff --git a/packages/orchestrator/pkg/sandbox/sandbox.go b/packages/orchestrator/pkg/sandbox/sandbox.go index b7c30b2fbb..dd503c16a0 100644 --- a/packages/orchestrator/pkg/sandbox/sandbox.go +++ b/packages/orchestrator/pkg/sandbox/sandbox.go @@ -1494,6 +1494,7 @@ func (s *Sandbox) processMemorySnapshot(ctx context.Context, buildID uuid.UUID) dedupBestEffort, dedupDirectIO, dedupBudget, + s.featureFlags.BoolFlag(ctx, featureflags.MemfdDedupInflightServeFlag, sandboxLDContext(s.Runtime, s.Config)), ) if err != nil { return MemorySnapshot{}, fmt.Errorf("error while post processing: %w", err) @@ -1537,6 +1538,7 @@ func pauseProcessMemory( dedupBestEffort bool, dedupDirectIO bool, dedupBudget block.DedupBudget, + dedupInflightServe bool, ) (d build.Diff, h *DiffHeader, e error) { ctx, span := tracer.Start(ctx, "process-memory") defer span.End() @@ -1547,6 +1549,7 @@ func pauseProcessMemory( cache, err := fc.ExportMemory( ctx, diffMetadata.Dirty, memfileDiffPath, diffMetadata.BlockSize, memfd, bgCopy, originalMemfile, dedupBestEffort, dedupDirectIO, dedupBudget, diffMetadata.Empty, metaOut, + dedupInflightServe, ) if err != nil { return nil, nil, fmt.Errorf("failed to export memory: %w", err) diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index 3a68f14f83..b5b4201025 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -160,6 +160,12 @@ var ( "fetchRunWindowPages": 0, })) + // MemfdDedupInflightServeFlag lets a resume that overlaps an in-flight + // memfile dedup serve dirty pages straight from the still-mapped memfd + // instead of blocking on the dedup drain. Only affects the memfd-dedup + // path; off restores the prior wait-for-drain behavior. + MemfdDedupInflightServeFlag = NewBoolFlag("memfd-dedup-inflight-serve", false) + // PeerToPeerChunkTransferFlag enables peer-to-peer chunk routing. PeerToPeerChunkTransferFlag = NewBoolFlag("peer-to-peer-chunk-transfer", false) // PeerToPeerAsyncCheckpointFlag makes Checkpoint upload fire-and-forget instead From 05bc6a9456174126c2c7727437867df40ee056e5 Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Wed, 1 Jul 2026 15:48:25 +0100 Subject: [PATCH 2/9] test(orch): cover in-flight memfd serving during dedup drain Add unit tests for the packed-to-absolute offset index (single-run translation plus boundary and out-of-range rejection) and for DedupedMemfdCache serving dirty pages from the memfd while the drain is in progress, then bypassing the in-flight path once the drain resolves. Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 (1M context) --- .../pkg/sandbox/block/memfd_test.go | 70 +++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/packages/orchestrator/pkg/sandbox/block/memfd_test.go b/packages/orchestrator/pkg/sandbox/block/memfd_test.go index 5371af7461..3d335411da 100644 --- a/packages/orchestrator/pkg/sandbox/block/memfd_test.go +++ b/packages/orchestrator/pkg/sandbox/block/memfd_test.go @@ -248,3 +248,73 @@ func TestNewCacheFromMemfdDeduped_DetachesCompareAndDrain(t *testing.T) { expected = append(expected, srcData[pageSize*4:pageSize*5]...) require.Equal(t, expected, got) } + +// buildPackedIndex.translate maps a packed diff offset back to the absolute +// memfd offset for the dirty run it belongs to, and refuses ranges that cross +// a run boundary or fall outside any run. +func TestPackedIndexTranslate(t *testing.T) { + t.Parallel() + + ps := int64(header.PageSize) + // Dirty pages 0, 2, 5 pack contiguously as seg0=[0,ps)->abs 0, + // seg1=[ps,2ps)->abs 2ps, seg2=[2ps,3ps)->abs 5ps. + dirty := roaring.New() + dirty.AddMany([]uint32{0, 2, 5}) + idx := buildPackedIndex(dirty) + + for _, tc := range []struct{ packed, abs int64 }{{0, 0}, {ps, 2 * ps}, {2 * ps, 5 * ps}} { + abs, ok := idx.translate(tc.packed, ps) + require.True(t, ok) + require.Equal(t, tc.abs, abs) + } + + // A range spanning two runs can't be served from one contiguous memfd span. + _, ok := idx.translate(0, 2*ps) + require.False(t, ok) + // Past the last run. + _, ok = idx.translate(3*ps, ps) + require.False(t, ok) +} + +// While the drain is in progress (done unresolved), reads are served from the +// still-mapped memfd via the packed→absolute index; once done resolves the +// inflight path is bypassed in favor of the drained cache. +func TestDedupedMemfdCache_InflightServesFromMemfd(t *testing.T) { + t.Parallel() + + ps := int64(header.PageSize) + memfd, data := newTestMemfd(t, ps*6) + t.Cleanup(func() { _ = memfd.Close() }) + + // Non-adjacent dirty set so packed offsets differ from absolute offsets. + dirty := roaring.New() + dirty.AddMany([]uint32{0, 2, 5}) + + // Construct the post-compare, pre-drain state directly (no goroutine): the + // memfd + index are published and done is unresolved. + d := &DedupedMemfdCache{ + done: utils.NewSetOnce[*Cache](), + inflight: true, + memfd: memfd, + index: buildPackedIndex(dirty), + } + + // ReadAt at packed offsets resolves to the right absolute memfd pages. + for i, srcPage := range []int64{0, 2, 5} { + got := make([]byte, ps) + n, err := d.ReadAt(got, int64(i)*ps) + require.NoError(t, err) + require.Equal(t, int(ps), n) + require.Equal(t, data[srcPage*ps:(srcPage+1)*ps], got) + } + + // Slice takes the same path and returns a copy of the memfd bytes. + s, err := d.Slice(ps, ps) // packed page 1 -> absolute page 2 + require.NoError(t, err) + require.Equal(t, data[2*ps:3*ps], s) + + // Once the drain resolves done, the inflight path is bypassed. + require.NoError(t, d.done.SetValue(nil)) + _, ok := d.tryInflightRead(make([]byte, ps), 0) + require.False(t, ok) +} From 07c54f3c5a954d9dc2d7e036af21047d7bbd82f9 Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Wed, 1 Jul 2026 16:07:58 +0100 Subject: [PATCH 3/9] test(orch): stress in-flight memfd serving across the dedup handover Drive a real dedup goroutine with in-flight serving enabled and hammer tryInflightRead from many goroutines across the drain window and the memfd-to-cache handover: assert every served read returns the correct bytes and, under -race, that the drain's memfd close never unmaps beneath an in-flight reader. Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 (1M context) --- .../pkg/sandbox/block/memfd_test.go | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/packages/orchestrator/pkg/sandbox/block/memfd_test.go b/packages/orchestrator/pkg/sandbox/block/memfd_test.go index 3d335411da..bd1c165f07 100644 --- a/packages/orchestrator/pkg/sandbox/block/memfd_test.go +++ b/packages/orchestrator/pkg/sandbox/block/memfd_test.go @@ -6,8 +6,11 @@ import ( "bytes" "context" "crypto/rand" + "fmt" "io" "os" + "runtime" + "sync" "testing" "github.com/RoaringBitmap/roaring/v2" @@ -318,3 +321,77 @@ func TestDedupedMemfdCache_InflightServesFromMemfd(t *testing.T) { _, ok := d.tryInflightRead(make([]byte, ps), 0) require.False(t, ok) } + +// End-to-end: drive a real dedup goroutine with inflight serving on and hammer +// tryInflightRead from many goroutines across the drain window and the +// memfd->cache handover. Every served read must return the correct bytes, and +// under -race the drain's memfd close must never unmap beneath an in-flight +// reader. Base is all zeros so every (random) source page stays in the diff: +// the deduped dirty set is the full range and packed offsets equal absolute. +func TestDedupedMemfdCache_InflightConcurrentDrainRace(t *testing.T) { + t.Parallel() + + ps := int64(header.PageSize) + const numPages = 4096 + size := ps * numPages + + memfd, srcData := newTestMemfd(t, size) + base := &fakeOriginalDevice{data: make([]byte, size)} + + dirty := roaring.New() + dirty.AddRange(0, numPages) + + metaOut := utils.NewSetOnce[*header.DiffMetadata]() + cache, err := NewCacheFromMemfdDeduped( + t.Context(), base, ps, t.TempDir()+"/dedup-concurrent", memfd, dirty, + false, false, DedupBudget{}, nil, metaOut, true, + ) + require.NoError(t, err) + t.Cleanup(func() { _ = cache.Close() }) + + var wg sync.WaitGroup + stop := make(chan struct{}) + errCh := make(chan error, 16) + for g := range 8 { + wg.Add(1) + go func(seed int) { + defer wg.Done() + buf := make([]byte, ps) + for i := seed; ; i++ { + select { + case <-stop: + return + default: + } + page := int64(i % numPages) + n, ok := cache.tryInflightRead(buf, page*ps) + if ok && (n != int(ps) || !bytes.Equal(buf, srcData[page*ps:(page+1)*ps])) { + errCh <- fmt.Errorf("inflight page %d mismatch", page) + + return + } + runtime.Gosched() + } + }(g) + } + + // Let compare+drain complete (memfd closes, done resolves) under reader load. + _, err = cache.Wait(t.Context()) + require.NoError(t, err) + + // Post-handover: every page reads correctly from the drained cache. + buf := make([]byte, ps) + for page := range int64(numPages) { + _, rErr := cache.ReadAt(buf, page*ps) + require.NoError(t, rErr) + require.Equal(t, srcData[page*ps:(page+1)*ps], buf) + } + + close(stop) + wg.Wait() + select { + case e := <-errCh: + t.Fatal(e) + default: + } +} From d25765f6e1f22f60fac194344a3d8d47467e107b Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Wed, 1 Jul 2026 17:51:16 +0100 Subject: [PATCH 4/9] feat(shared): add provisional identity-offset diff-header builder Add createIdentityMapping and DiffMetadata.ToProvisionalDiffHeader, which build a local-only memfile header that attributes every dirty page to a given build at an identity storage offset (its device offset) composed over the parent header. Unlike ToDiffHeader it needs no dedup metadata, so a pause can produce it immediately and a concurrent resume can serve from the still-mapped memfd without waiting for the dedup compare. The uploaded artifact continues to use the deduped ToDiffHeader output; this header is never persisted. Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 (1M context) --- packages/shared/pkg/storage/header/mapping.go | 26 +++++++++++++++ .../shared/pkg/storage/header/metadata.go | 32 +++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/packages/shared/pkg/storage/header/mapping.go b/packages/shared/pkg/storage/header/mapping.go index 7d2cad0cf6..e10c52f16f 100644 --- a/packages/shared/pkg/storage/header/mapping.go +++ b/packages/shared/pkg/storage/header/mapping.go @@ -43,6 +43,32 @@ func CreateMapping( return mappings } +// createIdentityMapping is like CreateMapping but sets each range's +// BuildStorageOffset equal to its device Offset instead of packing the ranges +// contiguously. It is for a build whose data is addressed by absolute device +// offset (e.g. a still-mapped memfd) rather than a compacted diff artifact — +// used to build a provisional, local-only header while dedup is still running. +func createIdentityMapping( + buildId *uuid.UUID, + dirty *roaring.Bitmap, + blockSize int64, +) []BuildMap { + var mappings []BuildMap + + for start, endExcl := range dirty.Ranges() { + blockLength := int64(endExcl) - int64(start) + off := uint64(BlockOffset(int64(start), blockSize)) + mappings = append(mappings, BuildMap{ + Offset: off, + BuildId: *buildId, + Length: uint64(BlockOffset(blockLength, blockSize)), + BuildStorageOffset: off, + }) + } + + return mappings +} + // MergeMappings merges two sets of mappings. // // The mapping are stored in a sorted order. diff --git a/packages/shared/pkg/storage/header/metadata.go b/packages/shared/pkg/storage/header/metadata.go index 00e0bbde97..6faea94328 100644 --- a/packages/shared/pkg/storage/header/metadata.go +++ b/packages/shared/pkg/storage/header/metadata.go @@ -99,6 +99,38 @@ func NewDiffMetadata(blockSize int64, dirty, empty *roaring.Bitmap) *DiffMetadat } } +// ToProvisionalDiffHeader builds a local-only header that describes the memfd +// directly — all dirty pages attributed to buildID with identity storage +// offsets (see createIdentityMapping), composed over originalHeader. Unlike +// ToDiffHeader it needs no dedup metadata, so it is available at pause time and +// lets a resume serve immediately. It must never be uploaded (the uploaded +// artifact always uses the deduped ToDiffHeader output). +func (d *DiffMetadata) ToProvisionalDiffHeader( + ctx context.Context, + originalHeader *Header, + buildID uuid.UUID, +) (*Header, error) { + _, span := tracer.Start(ctx, "to provisional diff-header") + defer span.End() + + dirtyMappings := createIdentityMapping(&buildID, d.Dirty, d.BlockSize) + emptyMappings := CreateMapping(&ignoreBuildID, d.Empty, d.BlockSize) + diffMapping := MergeMappings(dirtyMappings, emptyMappings) + + m := NormalizeMappings(MergeMappings(originalHeader.Mapping.Slice(), diffMapping)) + metadata := originalHeader.Metadata.NextGeneration(buildID) + + h, err := newDiffHeader(metadata, m, originalHeader.Builds) + if err != nil { + return nil, fmt.Errorf("failed to create provisional header: %w", err) + } + if err := h.Mapping.Validate(h.Metadata.Size, PageSize); err != nil { + return nil, fmt.Errorf("invalid provisional header mappings: %w", err) + } + + return h, nil +} + func (d *DiffMetadata) toDiffMapping( ctx context.Context, buildID uuid.UUID, From 19335be025dd2a44a1fe37c4fc9b4dbd89a46ec2 Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Wed, 1 Jul 2026 17:51:16 +0100 Subject: [PATCH 5/9] test(shared): cover the provisional diff-header builder Assert ToProvisionalDiffHeader maps dirty pages to the new build at identity storage offsets and leaves clean pages on the parent. Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 (1M context) --- .../pkg/storage/header/provisional_test.go | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 packages/shared/pkg/storage/header/provisional_test.go diff --git a/packages/shared/pkg/storage/header/provisional_test.go b/packages/shared/pkg/storage/header/provisional_test.go new file mode 100644 index 0000000000..c0dc6d2aa3 --- /dev/null +++ b/packages/shared/pkg/storage/header/provisional_test.go @@ -0,0 +1,48 @@ +package header + +import ( + "testing" + + "github.com/RoaringBitmap/roaring/v2" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// ToProvisionalDiffHeader attributes every dirty page to the new build with an +// identity storage offset (= its device offset), and leaves clean pages mapped +// to the parent — all without any dedup metadata. +func TestToProvisionalDiffHeader_IdentityOffsetsOverParent(t *testing.T) { + t.Parallel() + + const numPages = 6 + parent := uuid.New() + child := uuid.New() + + base, err := NewHeader( + &Metadata{Version: MetadataVersionV4, BlockSize: PageSize, Size: numPages * PageSize, BuildId: parent, BaseBuildId: parent}, + []BuildMap{{Offset: 0, Length: numPages * PageSize, BuildId: parent, BuildStorageOffset: 0}}, + ) + require.NoError(t, err) + + dirty := roaring.New() + dirty.AddMany([]uint32{1, 3}) + dm := &DiffMetadata{Dirty: dirty, Empty: roaring.New(), BlockSize: PageSize} + + h, err := dm.ToProvisionalDiffHeader(t.Context(), base, child) + require.NoError(t, err) + + // Dirty pages resolve to the child build at an identity storage offset. + for _, page := range []int64{1, 3} { + m, err := h.GetShiftedMapping(t.Context(), page*PageSize) + require.NoError(t, err) + require.Equal(t, child, m.BuildId, "page %d should map to child", page) + require.Equal(t, uint64(page*PageSize), m.Offset, "page %d storage offset must be identity", page) + } + + // Clean pages still resolve to the parent. + for _, page := range []int64{0, 2, 4, 5} { + m, err := h.GetShiftedMapping(t.Context(), page*PageSize) + require.NoError(t, err) + require.Equal(t, parent, m.BuildId, "page %d should map to parent", page) + } +} From 659f9a1e0d6fd75c16a6536d2a31184faf2efd55 Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Wed, 1 Jul 2026 22:32:25 +0100 Subject: [PATCH 6/9] feat(orch): add memfd identity-offset source for provisional-header serving Add DedupedMemfdCache.ServeMemfd (identity device-offset reads guarded by the existing memfd RWMutex) and MemfdIdentitySource, a local diff source that serves dirty pages from the still-mapped memfd while dedup runs. This is what a provisional (pre-dedup) memfile header reads from, so a resume overlapping the pause can serve dirty pages immediately instead of waiting for the deduped header. Publish the memfd at construction and free it exactly once under the write lock (releaseMemfd) so serving can't race the drain's memfd close. Composes with the in-flight drain serving rather than replacing it. Not yet wired into the pause/resume path. Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 (1M context) --- .../orchestrator/pkg/sandbox/block/memfd.go | 157 ++++++++++++++++-- 1 file changed, 141 insertions(+), 16 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/block/memfd.go b/packages/orchestrator/pkg/sandbox/block/memfd.go index 5ea0794031..590654b131 100644 --- a/packages/orchestrator/pkg/sandbox/block/memfd.go +++ b/packages/orchestrator/pkg/sandbox/block/memfd.go @@ -243,15 +243,28 @@ type DedupedMemfdCache struct { done *utils.SetOnce[*Cache] inflight bool - // mu guards memfd + index. memfd is non-nil only between "compare done" - // (published in runDedup) and "drain done" (closed under the write lock so - // it can't be unmapped under an in-flight reader). index translates a + // swapped is resolved by MarkSwapped once the local template has swapped + // off the provisional header onto the deduped one. runDedup waits on it + // (bounded) before releasing the memfd, so a provisional read can't hit a + // released memfd during the compare→swap window (which outlives the drain + // when the dirty set is small and the parent header is fragmented). + swapped *utils.SetOnce[struct{}] + // mu guards memfd + index. memfd is non-nil from construction until it is + // released — after the drain, or, when a provisional header serves from it, + // after the swap (MarkSwapped) or a grace timeout — closed under the write + // lock so it can't be unmapped under an in-flight reader. index translates a // packed diff-storage offset back to the absolute memfd offset. mu sync.RWMutex memfd *Memfd index packedIndex } +// memfdSwapGrace bounds how long runDedup keeps the memfd mapped waiting for the +// provisional→deduped header swap, so a failed or absent swap can't leak the +// mapping. The swap normally completes before the drain does (it only waits on +// the compare), so this grace is a backstop, not the common path. +const memfdSwapGrace = 30 * time.Second + func NewCacheFromMemfdDeduped( ctx context.Context, base ReadonlyDevice, @@ -274,7 +287,13 @@ func NewCacheFromMemfdDeduped( outPath: outPath, cancel: cancel, done: utils.NewSetOnce[*Cache](), + swapped: utils.NewSetOnce[struct{}](), inflight: inflightServe, + // Publish the memfd up front (guarded by mu) so a provisional-header + // resume can serve dirty pages via ServeMemfd during the compare window, + // before metaOut/the deduped header resolves. runDedup frees it under mu + // after the drain. + memfd: memfd, } go d.runDedup(drainCtx, base, blockSize, memfd, dirty, bestEffort, directIO, budget, inputEmpty, metaOut) @@ -353,7 +372,7 @@ func (d *DedupedMemfdCache) runDedup( compareDur := time.Since(compareStart) if err != nil { logSetOnceErr(ctx, "dedup metaOut", metaOut.SetError(err)) - logSetOnceErr(ctx, "dedup done", d.done.SetError(errors.Join(err, memfd.Close()))) + logSetOnceErr(ctx, "dedup done", d.done.SetError(errors.Join(err, d.releaseMemfd()))) return } @@ -372,12 +391,13 @@ func (d *DedupedMemfdCache) runDedup( // Whole-VM empty set recorded in the header (scan zeros + inputEmpty). telemetry.SetAttributes(ctx, attribute.Int64("dedup.header_empty_pages", int64(plan.pageEmpty.GetCardinality()))) - logSetOnceErr(ctx, "dedup metaOut", metaOut.SetValue(meta)) - // Publish the memfd + packed→absolute index so a resume overlapping this - // pause serves dirty pages straight from the memfd during the drain window - // instead of blocking on done. Reads take mu.RLock; the close below takes - // mu.Lock, so the memfd is never unmapped under an in-flight reader. + // Publish the packed→absolute index BEFORE resolving metaOut. metaOut + // resolving unblocks the deduped-header build and the subsequent SwapHeader, + // after which reads route to this diff and rely on the index; installing it + // first guarantees tryInflightRead never sees a nil index post-swap. Reads + // take mu.RLock; the drain's close below takes mu.Lock, so the memfd is never + // unmapped under an in-flight reader. if d.inflight { d.mu.Lock() d.index = buildPackedIndex(plan.pageDirty) @@ -385,19 +405,38 @@ func (d *DedupedMemfdCache) runDedup( d.mu.Unlock() } + logSetOnceErr(ctx, "dedup metaOut", metaOut.SetValue(meta)) + writeStart := time.Now() cache, err := dedupDrain(ctx, src, plan.pageDirty, blockSize, d.outPath, directIO) writeDur := time.Since(writeStart) - d.mu.Lock() - closeErr := memfd.Close() - d.memfd = nil - d.mu.Unlock() - if closeErr != nil { - logger.L().Warn(ctx, "close memfd after dedup drain", zap.Error(closeErr)) - } recordDedupAttrs(ctx, plan, scanEmptyPages, compareDur, writeDur) + // Resolve the drained cache immediately so upload and post-swap reads don't + // wait. Then, when inflight serving is active, keep the memfd mapped until + // the local template swaps off the provisional header (MarkSwapped) so a + // provisional read never hits a released memfd. The swap only waits on the + // compare, so it has usually already fired by now; ctx and the grace bound + // the wait so a failed or absent swap can't leak the mapping. logSetOnceErr(ctx, "dedup done", d.done.SetResult(cache, err)) + if d.inflight { + select { + case <-d.swapped.Done: + case <-ctx.Done(): + case <-time.After(memfdSwapGrace): + logger.L().Warn(ctx, "memfd swap grace elapsed; releasing memfd without a swap signal") + } + } + if closeErr := d.releaseMemfd(); closeErr != nil { + logger.L().Warn(ctx, "close memfd after dedup drain", zap.Error(closeErr)) + } +} + +// MarkSwapped signals that the local template has swapped off the provisional +// header, so runDedup can release the memfd it was serving provisional reads +// from. Idempotent; safe to call when no provisional header was ever built. +func (d *DedupedMemfdCache) MarkSwapped() { + _ = d.swapped.SetValue(struct{}{}) } // logSetOnceErr warns on a SetOnce.SetValue/SetError failure (i.e. a @@ -475,6 +514,40 @@ func (d *DedupedMemfdCache) Slice(off, length int64) ([]byte, error) { return c.Slice(off, length) } +// releaseMemfd closes the memfd exactly once, under the write lock so it can't +// be unmapped beneath an in-flight ServeMemfd/tryInflightRead reader. +func (d *DedupedMemfdCache) releaseMemfd() error { + d.mu.Lock() + defer d.mu.Unlock() + if d.memfd == nil { + return nil + } + err := d.memfd.Close() + d.memfd = nil + + return err +} + +// ServeMemfd copies [off, off+len(b)) from the still-mapped memfd using +// identity (device) addressing, for a provisional local header that attributes +// dirty pages to the memfd at their device offset. Returns BytesNotAvailableError +// once the memfd has been released (drain finished), so the caller falls back to +// the drained cache via the swapped-in deduped header. +func (d *DedupedMemfdCache) ServeMemfd(b []byte, off int64) (int, error) { + d.mu.RLock() + defer d.mu.RUnlock() + if d.memfd == nil { + return 0, BytesNotAvailableError{} + } + src, err := d.memfd.Slice(off, int64(len(b))) + if err != nil { + return 0, err + } + copy(b, src) + + return len(b), nil +} + func (d *DedupedMemfdCache) Close() error { d.cancel() c, _ := d.done.Wait() @@ -486,6 +559,58 @@ func (d *DedupedMemfdCache) Close() error { return nil } +// MemfdIdentitySource is the provisional local diff source: it serves dirty +// pages from a DedupedMemfdCache's still-mapped memfd at identity (device) +// offsets while dedup runs. It backs a distinct provisional build id; once the +// deduped header is swapped in, reads route to the real build id instead and +// this source is dropped. It is never uploaded. +type MemfdIdentitySource struct { + d *DedupedMemfdCache + size int64 +} + +func NewMemfdIdentitySource(d *DedupedMemfdCache, size int64) *MemfdIdentitySource { + return &MemfdIdentitySource{d: d, size: size} +} + +func (s *MemfdIdentitySource) ReadAt(b []byte, off int64) (int, error) { return s.d.ServeMemfd(b, off) } + +func (s *MemfdIdentitySource) Slice(off, length int64) ([]byte, error) { + b := make([]byte, length) + if _, err := s.d.ServeMemfd(b, off); err != nil { + return nil, err + } + + return b, nil +} + +// IsCached reports the range as resident while the memfd is mapped. It satisfies +// the CachePeeker contract for callers that hold a *MemfdIdentitySource directly +// (e.g. the block-level tests). Note the wrapping *build.localDiff does not +// promote this method, so build.File.IsCached does not reach it today. +func (s *MemfdIdentitySource) IsCached(_ context.Context, off, length int64) bool { + s.d.mu.RLock() + defer s.d.mu.RUnlock() + + return s.d.memfd != nil && off >= 0 && off+length <= s.size +} + +func (s *MemfdIdentitySource) Size() (int64, error) { return s.size, nil } +func (s *MemfdIdentitySource) BlockSize() int64 { return header.PageSize } +func (s *MemfdIdentitySource) Close() error { return nil } + +// FileSize reports the on-disk cache footprint, which is zero: this source wraps +// a still-mapped memfd, not a file in the cache directory, and evicting it frees +// no disk bytes. Returning the logical size here would inflate the DiffStore's +// disk-eviction accounting by ~guest-RAM bytes for a phantom entry. +func (s *MemfdIdentitySource) FileSize(context.Context) (int64, error) { return 0, nil } + +// Path has no meaning for a memfd-backed source; it is local-only and never +// uploaded, so the upload path (which needs a file path) must never reach it. +func (s *MemfdIdentitySource) Path(context.Context) (string, error) { + return "", errors.New("provisional memfd source has no path") +} + func (d *DedupedMemfdCache) Path(ctx context.Context) (string, error) { c, err := d.Wait(ctx) if err != nil { From 2043ff4eb5cd50ca1710c22749fa7dd26d04e784 Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Wed, 1 Jul 2026 22:32:45 +0100 Subject: [PATCH 7/9] test(orch): cover the memfd identity-offset source Assert MemfdIdentitySource serves dirty pages from the memfd at identity offsets while mapped, and reports BytesNotAvailableError once the memfd is released. Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 (1M context) --- .../pkg/sandbox/block/memfd_test.go | 78 +++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/packages/orchestrator/pkg/sandbox/block/memfd_test.go b/packages/orchestrator/pkg/sandbox/block/memfd_test.go index bd1c165f07..bc8b59ce6a 100644 --- a/packages/orchestrator/pkg/sandbox/block/memfd_test.go +++ b/packages/orchestrator/pkg/sandbox/block/memfd_test.go @@ -6,12 +6,14 @@ import ( "bytes" "context" "crypto/rand" + "errors" "fmt" "io" "os" "runtime" "sync" "testing" + "time" "github.com/RoaringBitmap/roaring/v2" "github.com/stretchr/testify/require" @@ -328,6 +330,54 @@ func TestDedupedMemfdCache_InflightServesFromMemfd(t *testing.T) { // under -race the drain's memfd close must never unmap beneath an in-flight // reader. Base is all zeros so every (random) source page stays in the diff: // the deduped dirty set is the full range and packed offsets equal absolute. +// TestDedupedMemfdCache_MemfdHeldUntilSwap covers the swap/release ordering: with +// inflight serving active, the memfd stays mapped past the drain — past the point +// the drained cache resolves — until MarkSwapped fires, so a provisional read +// served via ServeMemfd never hits a released memfd during the compare→swap +// window (which outlives the drain for a small dirty set + fragmented parent). +// Before the fix the memfd was closed immediately after the drain. +func TestDedupedMemfdCache_MemfdHeldUntilSwap(t *testing.T) { + t.Parallel() + + ps := int64(header.PageSize) + const numPages = 64 + size := ps * numPages + + memfd, srcData := newTestMemfd(t, size) + base := &fakeOriginalDevice{data: make([]byte, size)} // all-zero base → every dirty page kept + dirty := roaring.New() + dirty.AddRange(0, numPages) + + metaOut := utils.NewSetOnce[*header.DiffMetadata]() + cache, err := NewCacheFromMemfdDeduped( + t.Context(), base, ps, t.TempDir()+"/dedup-held", memfd, dirty, + false, false, DedupBudget{}, nil, metaOut, true, + ) + require.NoError(t, err) + t.Cleanup(func() { _ = cache.Close() }) + + // The drained cache resolves once the drain completes... + _, err = cache.Wait(t.Context()) + require.NoError(t, err) + + // ...but the memfd is still mapped (held for a pending provisional swap), so a + // provisional identity read still returns the page bytes. + buf := make([]byte, ps) + n, err := cache.ServeMemfd(buf, 3*ps) + require.NoError(t, err) + require.Equal(t, int(ps), n) + require.Equal(t, srcData[3*ps:4*ps], buf) + + // Once the swap is signaled the memfd is released and provisional reads fail. + cache.MarkSwapped() + require.Eventually(t, func() bool { + _, e := cache.ServeMemfd(buf, 3*ps) + var bna BytesNotAvailableError + + return errors.As(e, &bna) + }, time.Second, time.Millisecond) +} + func TestDedupedMemfdCache_InflightConcurrentDrainRace(t *testing.T) { t.Parallel() @@ -395,3 +445,31 @@ func TestDedupedMemfdCache_InflightConcurrentDrainRace(t *testing.T) { default: } } + +// MemfdIdentitySource serves dirty pages from the memfd at identity offsets +// while it is mapped, and reports BytesNotAvailableError once it is released so +// the caller falls back to the drained cache (via the swapped-in deduped header). +func TestMemfdIdentitySource_ServesThenReleases(t *testing.T) { + t.Parallel() + + ps := int64(header.PageSize) + memfd, data := newTestMemfd(t, ps*4) + d := &DedupedMemfdCache{done: utils.NewSetOnce[*Cache](), memfd: memfd} + src := NewMemfdIdentitySource(d, ps*4) + + for _, page := range []int64{0, 2, 3} { + b := make([]byte, ps) + n, err := src.ReadAt(b, page*ps) + require.NoError(t, err) + require.Equal(t, int(ps), n) + require.Equal(t, data[page*ps:(page+1)*ps], b) + } + 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 + require.ErrorAs(t, err, &bna) + require.False(t, src.IsCached(t.Context(), 0, ps)) +} From a3922e4eaeb2803c1318b8ff797ee7db414d4b69 Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Thu, 2 Jul 2026 14:21:20 +0100 Subject: [PATCH 8/9] feat(orch): serve resume from a provisional header during dedup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A resume that overlaps an in-flight pause of the same sandbox can block in the storage-template-memfile span: the local memfile template can't be built until the deduped header resolves, which only happens after dedup's compare+defrag. The in-flight memfd serving added earlier removes the drain wait but not this header wait, because it is upstream of the diff source. At pause, build a provisional local header (identity storage offsets, dirty pages attributed to a fresh provisional build id) from the dirty bitmap and the parent header, and register a memfd-backed diff source for that build id. The local template is built from the provisional header immediately, so a concurrent resume serves dirty pages straight from the still-mapped memfd instead of blocking. Once the deduped header resolves it is swapped in (SwapHeader), routing dirty pages to the compacted deduped diff; the distinct build id makes the swap race-free since build.File resolves the source from the same header snapshot as the offset. The swap is conditional (compare-and-swap from the still-provisional header) so a late swap can't clobber a newer header the upload may have finalized first. The provisional header is local-only: the upload continues to use the deduped header, so the persisted artifact is unchanged. A pause must never parent a snapshot off the provisional header: it maps dirty pages to a synthetic build id with no storage object, so an uploaded header that inherited those mappings (from a resume paused again before the swap) would be unreadable on a cold or cross-node resume. build.File carries the deduped header as a durable-header future, and processMemorySnapshot resolves the parent via DurableHeader (waiting for the deduped header if a swap is still pending) rather than the live, possibly-provisional header. The peer-to-peer header source serves the durable header for the same reason, so a cross-node resume in the window receives the deduped mappings (waiting if a swap is pending) instead of the synthetic provisional build id it cannot resolve. Scheduling metadata (storageTemplate.SchedulingMetadata) reads the durable header too, but NON-BLOCKING (DurableHeaderNow): it runs on the create/resume response path, so while the deduped header is still pending it reports rootfs-only metadata rather than block the response for the whole dedup — and never carries the provisional build id. The dedup goroutine holds the memfd until the swap, bounded by a grace timeout; the swap goroutine signals on every exit (success or error), and buildProvisionalMemfile signals up front when it declines to build a provisional source, so the fallback releases at drain-time and the grace is only a true last-resort backstop — a counter (orchestrator.memfd.swap_grace_elapsed) surfaces it if it ever fires. Both the main memfile diff and the provisional diff are pinned in the build store for the provisional window: they share a DedupedMemfdCache (and memfd), and resume reads refresh only the provisional entry, so disk-pressure eviction of either would break the in-flight serve — evicting the main entry Closes it (cancelling dedup, tearing down the memfd), evicting the provisional entry makes a still-provisional read miss the store and fetch a never-uploaded id. Both are unpinned after the swap. Gated by memfd-dedup-inflight-serve (default off); composes with the in-flight drain serving (provisional covers the compare window, in-flight the drain). Signed-off-by: Nikita Kalyazin Co-Authored-By: Claude Opus 4.8 (1M context) --- .../orchestrator/pkg/sandbox/block/memfd.go | 22 ++- .../orchestrator/pkg/sandbox/build/build.go | 57 +++++++- .../orchestrator/pkg/sandbox/build/cache.go | 25 ++++ packages/orchestrator/pkg/sandbox/sandbox.go | 105 ++++++++++++-- .../pkg/sandbox/template/cache.go | 128 +++++++++++++++++- .../pkg/sandbox/template/peerserver/header.go | 16 +++ .../pkg/sandbox/template/storage.go | 17 +++ .../pkg/sandbox/template/storage_template.go | 58 ++++++-- packages/orchestrator/pkg/server/sandboxes.go | 3 + .../template/build/layer/layer_executor.go | 3 + packages/shared/pkg/featureflags/flags.go | 7 +- 11 files changed, 409 insertions(+), 32 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/block/memfd.go b/packages/orchestrator/pkg/sandbox/block/memfd.go index 590654b131..30fb71d98e 100644 --- a/packages/orchestrator/pkg/sandbox/block/memfd.go +++ b/packages/orchestrator/pkg/sandbox/block/memfd.go @@ -11,7 +11,9 @@ import ( "time" "github.com/RoaringBitmap/roaring/v2" + "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" "go.opentelemetry.io/otel/trace" "go.uber.org/zap" "golang.org/x/sys/unix" @@ -265,6 +267,15 @@ type DedupedMemfdCache struct { // the compare), so this grace is a backstop, not the common path. const memfdSwapGrace = 30 * time.Second +// swapGraceElapsedCounter counts memfd releases that fell back to the grace +// timeout because no swap signal arrived. Expected to be ~0 (the swap normally +// signals before the drain finishes); a nonzero rate means swaps are being lost +// — e.g. pauses aborting before AddSnapshot spawns the swap goroutine — and the +// memfd (guest-RAM-sized) is being held the full grace on those pauses. +var swapGraceElapsedCounter = utils.Must(otel.Meter("github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block"). + Int64Counter("orchestrator.memfd.swap_grace_elapsed", + metric.WithDescription("Dedup memfd releases that timed out on the swap grace without a swap signal"))) + func NewCacheFromMemfdDeduped( ctx context.Context, base ReadonlyDevice, @@ -413,17 +424,20 @@ func (d *DedupedMemfdCache) runDedup( recordDedupAttrs(ctx, plan, scanEmptyPages, compareDur, writeDur) // Resolve the drained cache immediately so upload and post-swap reads don't - // wait. Then, when inflight serving is active, keep the memfd mapped until - // the local template swaps off the provisional header (MarkSwapped) so a + // wait. Then, when inflight serving is active, keep the memfd mapped until the + // local template swaps off the provisional header (MarkSwapped) so a // provisional read never hits a released memfd. The swap only waits on the - // compare, so it has usually already fired by now; ctx and the grace bound - // the wait so a failed or absent swap can't leak the mapping. + // compare, so it has usually already fired by now; ctx and the grace bound the + // wait so a failed or absent swap can't leak the mapping. When no provisional + // source is built, buildProvisionalMemfile calls MarkSwapped up front, so this + // returns immediately and releases at drain-time like the non-inflight path. logSetOnceErr(ctx, "dedup done", d.done.SetResult(cache, err)) if d.inflight { select { case <-d.swapped.Done: case <-ctx.Done(): case <-time.After(memfdSwapGrace): + swapGraceElapsedCounter.Add(ctx, 1) logger.L().Warn(ctx, "memfd swap grace elapsed; releasing memfd without a swap signal") } } diff --git a/packages/orchestrator/pkg/sandbox/build/build.go b/packages/orchestrator/pkg/sandbox/build/build.go index 6925a74ab8..92357b80b1 100644 --- a/packages/orchestrator/pkg/sandbox/build/build.go +++ b/packages/orchestrator/pkg/sandbox/build/build.go @@ -52,7 +52,12 @@ var ( ) type File struct { - header atomic.Pointer[header.Header] + header atomic.Pointer[header.Header] + // durable, when set, is the future for the header this file will eventually + // settle on (the deduped header while a provisional header is being served). + // DurableHeader waits on it so a pause never parents a snapshot off a + // provisional (local-only) header — see SetDurableHeader. + durable atomic.Pointer[utils.SetOnce[*header.Header]] store *DiffStore fileType DiffType persistence storage.StorageProvider @@ -91,6 +96,56 @@ func (b *File) SwapHeader(h *header.Header) { b.header.Store(h) } +// SwapHeaderIfCurrent atomically replaces the header with next only if it is +// still old, reporting whether it did. The provisional→deduped swap uses this so +// that, if it runs late, it can't clobber a newer header already installed by +// another writer (e.g. Upload.publish finalizing the build, which swaps +// unconditionally) with the older, still-incomplete deduped header. +func (b *File) SwapHeaderIfCurrent(old, next *header.Header) bool { + return b.header.CompareAndSwap(old, next) +} + +// SetDurableHeader records the future for the durable header this file will +// settle on, so DurableHeader can return it instead of the currently-installed +// (possibly provisional) header. Set once, before the file is shared. +func (b *File) SetDurableHeader(f *utils.SetOnce[*header.Header]) { + b.durable.Store(f) +} + +// DurableHeader returns the header safe to persist as a parent: the durable +// future's value if one was set (waiting for it), otherwise the currently +// installed header. Callers building an *uploaded* snapshot header (a Pause) +// must use this, never Header(), so they never inherit a provisional header's +// synthetic build-id mappings, which have no storage object. +func (b *File) DurableHeader(ctx context.Context) (*header.Header, error) { + if f := b.durable.Load(); f != nil { + return f.WaitWithContext(ctx) + } + + return b.Header(), nil +} + +// DurableHeaderNow is the non-blocking form of DurableHeader: it returns +// (header, true) when the durable header is already known — no provisional swap +// pending, or the deduped future has resolved — and (nil, false) while a swap is +// still pending. Latency-sensitive callers (e.g. scheduling metadata on the +// create/resume response path) use it so they neither block on the dedup nor +// observe the provisional header's synthetic build id. +func (b *File) DurableHeaderNow() (*header.Header, bool) { + f := b.durable.Load() + if f == nil { + return b.Header(), true + } + select { + case <-f.Done: + h, err := f.Wait() + + return h, err == nil + default: + return nil, false + } +} + // ReadAt times file.read_at around readAt; Slice composes via readAt directly to // avoid double counting. func (b *File) ReadAt(ctx context.Context, p []byte, off int64) (n int, err error) { diff --git a/packages/orchestrator/pkg/sandbox/build/cache.go b/packages/orchestrator/pkg/sandbox/build/cache.go index d6ce891314..580a14c0e5 100644 --- a/packages/orchestrator/pkg/sandbox/build/cache.go +++ b/packages/orchestrator/pkg/sandbox/build/cache.go @@ -53,6 +53,12 @@ type DiffStore struct { pdDelay time.Duration insertionTimes sync.Map // map[DiffStoreKey]time.Time — tracks when each diff was cached + + // pinned entries are skipped by disk-pressure eviction (TTL eviction still + // applies). Used to protect a diff whose Close would tear down state another + // live entry depends on — e.g. the memfile diff whose DedupedMemfdCache is + // also serving an in-flight provisional resume. + pinned sync.Map // map[DiffStoreKey]struct{} } func NewDiffStore( @@ -284,6 +290,12 @@ func (s *DiffStore) deleteOldestFromCache(ctx context.Context) (suc bool, e erro return true } + // Skip pinned entries (e.g. a memfile diff still backing an in-flight + // provisional resume); closing them would tear down shared state. + if s.isPinned(item.Key()) { + return true + } + sfSize, err := item.Value().FileSize(ctx) if err != nil { logger.L().Warn(ctx, "failed to get size of deleted item from cache", zap.Error(err)) @@ -324,6 +336,19 @@ func (s *DiffStore) isBeingDeleted(key DiffStoreKey) bool { return f } +// Pin protects a cached entry from disk-pressure eviction (TTL eviction still +// applies). Idempotent; pair every Pin with an Unpin. +func (s *DiffStore) Pin(key DiffStoreKey) { s.pinned.Store(key, struct{}{}) } + +// Unpin lifts a Pin, making the entry eligible for disk-pressure eviction again. +func (s *DiffStore) Unpin(key DiffStoreKey) { s.pinned.Delete(key) } + +func (s *DiffStore) isPinned(key DiffStoreKey) bool { + _, ok := s.pinned.Load(key) + + return ok +} + func (s *DiffStore) scheduleDelete(ctx context.Context, key DiffStoreKey, dSize int64) { s.pdMu.Lock() defer s.pdMu.Unlock() diff --git a/packages/orchestrator/pkg/sandbox/sandbox.go b/packages/orchestrator/pkg/sandbox/sandbox.go index dd503c16a0..fa1f6b962a 100644 --- a/packages/orchestrator/pkg/sandbox/sandbox.go +++ b/packages/orchestrator/pkg/sandbox/sandbox.go @@ -1435,6 +1435,17 @@ func (s *Sandbox) Pause( type MemorySnapshot struct { Diff build.Diff DiffHeader *DiffHeader + // ProvisionalDiffHeader + ProvisionalDiff, when non-nil, let the local + // template serve immediately from the still-mapped memfd via a distinct + // provisional build id while dedup runs, instead of blocking a concurrent + // resume in storage-template-memfile on the deduped header. They feed only + // the local AddSnapshot path; the upload still uses DiffHeader (deduped). + ProvisionalDiffHeader *header.Header + ProvisionalDiff build.Diff + // ProvisionalSwapDone, when non-nil, is invoked by the AddSnapshot swap + // goroutine once it has swapped the deduped header in; it lets the dedup + // goroutine release the memfd the provisional source was serving from. + ProvisionalSwapDone func() // BlockSize is captured synchronously at Pause time because NewUpload's // compression validation needs it before the async dedup header resolves; // the dedup memfile path produces a page-granular Diff.BlockSize() that @@ -1456,7 +1467,21 @@ func (s *Sandbox) processMemorySnapshot(ctx context.Context, buildID uuid.UUID) if err != nil { return MemorySnapshot{}, fmt.Errorf("failed to get original memfile: %w", err) } + // Parent off the durable (deduped) header, never a provisional (local-only) + // one: a provisional header maps dirty pages to a synthetic build id with no + // storage object, so an uploaded header inheriting those mappings would be + // unreadable on a cold or cross-node resume. DurableHeader waits for the + // deduped header if a provisional swap is still pending; devices without one + // return their current header immediately. memfileHeader := originalMemfile.Header() + if dh, ok := originalMemfile.(interface { + DurableHeader(ctx context.Context) (*header.Header, error) + }); ok { + memfileHeader, err = dh.DurableHeader(ctx) + if err != nil { + return MemorySnapshot{}, fmt.Errorf("failed to resolve durable memfile header: %w", err) + } + } memfileDiffMetadata, err := s.Resources.memory.DiffMetadata(ctx, s.process) if err != nil { @@ -1481,7 +1506,7 @@ func (s *Sandbox) processMemorySnapshot(ctx context.Context, buildID uuid.UUID) } } - memfileDiff, memfileDiffHeader, err := pauseProcessMemory( + memfileDiff, memfileDiffHeader, provMemfileHeader, provMemfileDiff, provMemfileSwapDone, err := pauseProcessMemory( ctx, buildID, memfileHeader, @@ -1501,11 +1526,14 @@ func (s *Sandbox) processMemorySnapshot(ctx context.Context, buildID uuid.UUID) } return MemorySnapshot{ - Diff: memfileDiff, - DiffHeader: memfileDiffHeader, - BlockSize: memfileHeader.Metadata.BlockSize, - header: memfileHeader, - newBytes: memfileDiffMetadata.Dirty.GetCardinality() * uint64(memfileDiffMetadata.BlockSize), + Diff: memfileDiff, + DiffHeader: memfileDiffHeader, + ProvisionalDiffHeader: provMemfileHeader, + ProvisionalDiff: provMemfileDiff, + ProvisionalSwapDone: provMemfileSwapDone, + BlockSize: memfileHeader.Metadata.BlockSize, + header: memfileHeader, + newBytes: memfileDiffMetadata.Dirty.GetCardinality() * uint64(memfileDiffMetadata.BlockSize), }, nil } @@ -1539,7 +1567,7 @@ func pauseProcessMemory( dedupDirectIO bool, dedupBudget block.DedupBudget, dedupInflightServe bool, -) (d build.Diff, h *DiffHeader, e error) { +) (d build.Diff, h *DiffHeader, provisionalHeader *header.Header, provisionalDiff build.Diff, provisionalSwapDone func(), e error) { ctx, span := tracer.Start(ctx, "process-memory") defer span.End() @@ -1552,7 +1580,7 @@ func pauseProcessMemory( dedupInflightServe, ) if err != nil { - return nil, nil, fmt.Errorf("failed to export memory: %w", err) + return nil, nil, nil, nil, nil, fmt.Errorf("failed to export memory: %w", err) } diff, err := build.NewLocalDiffFromCache( @@ -1560,9 +1588,16 @@ func pauseProcessMemory( cache, ) if err != nil { - return nil, nil, fmt.Errorf("failed to create local diff from cache: %w", errors.Join(err, cache.Close())) + return nil, nil, nil, nil, nil, fmt.Errorf("failed to create local diff from cache: %w", errors.Join(err, cache.Close())) } + // Provisional local header: while the deduped header is still + // being computed, let a same-node resume serve dirty pages from the memfd via + // a distinct provisional build id at identity offsets. Gated on the memfd + // dedup path + the inflight-serve flag; best-effort (fall back to the deduped + // header on any error). The upload always uses the deduped header below. + provisionalHeader, provisionalDiff, provisionalSwapDone = buildProvisionalMemfile(ctx, cache, dedupInflightServe, originalMemfile, originalHeader, diffMetadata) + // Build the diff header on a goroutine so Pause returns without waiting // on memfd-dedup compare. ExportMemory resolves metaOut sync for every // other path, so Wait there is non-blocking; the goroutine is harmless. @@ -1589,7 +1624,57 @@ func pauseProcessMemory( setHeader(meta.ToDiffHeader(ctx, originalHeader, buildID)) }() - return diff, headerOut, nil + return diff, headerOut, provisionalHeader, provisionalDiff, provisionalSwapDone, nil +} + +// buildProvisionalMemfile builds the provisional local header + its memfd-backed +// diff source, plus a swap-done callback the AddSnapshot swap goroutine invokes +// once it has swapped the deduped header in (it lets the dedup goroutine release +// the memfd). Returns (nil, nil, nil) — falling back to the deduped header — +// when the path doesn't apply or on any error, so it never blocks a pause. The +// provisional source is keyed by a fresh build id so a header swap to the +// deduped header (after dedup) is race-free (see MemfdIdentitySource). +func buildProvisionalMemfile( + ctx context.Context, + cache block.DiffSource, + enabled bool, + originalMemfile block.ReadonlyDevice, + originalHeader *header.Header, + diffMetadata *header.DiffMetadata, +) (*header.Header, build.Diff, func()) { + dc, ok := cache.(*block.DedupedMemfdCache) + if !ok { + return nil, nil, nil + } + // From here dc != nil. On every path that declines to build a provisional + // source, signal MarkSwapped so runDedup's inflight memfd-hold releases at + // drain-time instead of waiting out the swap grace — nothing will serve from + // the memfd, so there's no reason to hold it. + if !enabled || originalMemfile == nil || originalHeader == nil || diffMetadata == nil { + dc.MarkSwapped() + + return nil, nil, nil + } + + provisionalBuildID := uuid.New() + provisionalHeader, err := diffMetadata.ToProvisionalDiffHeader(ctx, originalHeader, provisionalBuildID) + if err != nil { + logger.L().Warn(ctx, "build provisional memfile header; using deduped header", zap.Error(err)) + dc.MarkSwapped() + + return nil, nil, nil + } + + provisionalSource := block.NewMemfdIdentitySource(dc, int64(originalHeader.Metadata.Size)) + provisionalDiff, err := build.NewLocalDiffFromCache(build.GetDiffStoreKey(provisionalBuildID.String(), build.Memfile), provisionalSource) + if err != nil { + logger.L().Warn(ctx, "build provisional memfile diff; using deduped header", zap.Error(err)) + dc.MarkSwapped() + + return nil, nil, nil + } + + return provisionalHeader, provisionalDiff, dc.MarkSwapped } func pauseProcessRootfs( diff --git a/packages/orchestrator/pkg/sandbox/template/cache.go b/packages/orchestrator/pkg/sandbox/template/cache.go index 2296542fb1..62f3edf20d 100644 --- a/packages/orchestrator/pkg/sandbox/template/cache.go +++ b/packages/orchestrator/pkg/sandbox/template/cache.go @@ -195,6 +195,7 @@ func (c *Cache) GetTemplate( c.blockMetrics, nil, nil, + nil, ) if err != nil { return nil, fmt.Errorf("failed to create template cache from storage: %w", err) @@ -224,12 +225,36 @@ func (c *Cache) AddSnapshot( localMetafile File, memfileDiff build.Diff, rootfsDiff build.Diff, + // provisionalMemfileHeader/Diff: when non-nil the local memfile + // template is built from the provisional header (which attributes dirty pages + // to provisionalMemfileDiff's build id at identity offsets) so a concurrent + // resume serves immediately from the memfd; once memfileHeader (the deduped + // header) resolves it is swapped in. When nil, the template is built directly + // from memfileHeader as before. The upload always uses memfileHeader. + provisionalMemfileHeader *header.Header, + provisionalMemfileDiff build.Diff, + // provisionalSwapDone, when non-nil, is called once the deduped header has + // been swapped in; it lets the dedup goroutine release the memfd the + // provisional source was serving from. + provisionalSwapDone func(), ) error { switch memfileDiff.(type) { case *build.NoDiff: default: c.buildStore.Add(memfileDiff) } + if provisionalMemfileDiff != nil { + if _, ok := provisionalMemfileDiff.(*build.NoDiff); !ok { + // This provisional entry must stay resident until the SwapHeader below + // (the provisional window). It is keyed by a synthetic build id with no + // GCS object, so if it were evicted mid-window a resume read routed to + // it would miss and couldn't be rebuilt (createDiff has nothing to + // fetch), failing the read. It is pinned below (with the main memfile + // diff) so disk-pressure eviction skips it for the window; TTL eviction + // (hours) can't fire within the window (seconds). + c.buildStore.Add(provisionalMemfileDiff) + } + } switch rootfsDiff.(type) { case *build.NoDiff: @@ -237,15 +262,30 @@ func (c *Cache) AddSnapshot( c.buildStore.Add(rootfsDiff) } + // Build the local template from the provisional header (resolved now) so + // Memfile() doesn't block on dedup; fall back to the deduped header future. + // When serving a provisional header, pass the deduped header future as the + // memfile's durable header so a pause parents off it, never the provisional + // header (whose synthetic build id has no storage object). It is applied at + // construction — before the memfile device is published — so no reader can + // observe the device with the durable header unset. + localMemfileHeader := memfileHeader + var durableMemfileHeader *utils.SetOnce[*header.Header] + if provisionalMemfileHeader != nil { + localMemfileHeader = resolvedHeader(provisionalMemfileHeader) + durableMemfileHeader = memfileHeader + } + storageTemplate, err := newTemplateFromStorage( c.config.BuilderConfig, buildId, - memfileHeader, + localMemfileHeader, rootfsHeader, c.persistence, c.blockMetrics, localSnapfile, localMetafile, + durableMemfileHeader, ) if err != nil { return fmt.Errorf("failed to create template cache from storage: %w", err) @@ -253,6 +293,92 @@ func (c *Cache) AddSnapshot( c.getTemplateWithFetch(ctx, storageTemplate, 0) + // Swap the provisional header for the deduped one once dedup finishes, so + // subsequent reads route dirty pages to the (compacted) deduped diff and the + // provisional memfd source is no longer referenced. The durable header was + // wired in at construction above. + if provisionalMemfileHeader != nil { + // Pin both the main memfile diff and the provisional diff for the window. + // They share a DedupedMemfdCache/memfd, but resume reads refresh only the + // provisional entry, so disk-pressure eviction of either would break the + // in-flight provisional serve: evicting the main entry Closes it, which + // cancels dedup and tears down the shared memfd; evicting the provisional + // entry makes a still-provisional-header read miss the store and fall + // through to a storage fetch for the synthetic build id (never uploaded). + // Pinning skips both in disk-pressure eviction (TTL still applies); + // unpinned after the swap. + var pinnedKeys []build.DiffStoreKey + for _, d := range []build.Diff{memfileDiff, provisionalMemfileDiff} { + if d == nil { + continue + } + if _, isNoDiff := d.(*build.NoDiff); isNoDiff { + continue + } + key := d.CacheKey() + c.buildStore.Pin(key) + pinnedKeys = append(pinnedKeys, key) + } + + swapCtx := context.WithoutCancel(ctx) + go func() { + // Signal the dedup goroutine on every exit (success or the error + // returns below) so it releases the memfd promptly. On an error the + // swap can't happen and the resume is already broken, so nothing needs + // the memfd kept mapped; without this the dedup goroutine would wait out + // the full grace before releasing. Unpin the main diff on every exit too. + if provisionalSwapDone != nil { + defer provisionalSwapDone() + } + defer func() { + for _, key := range pinnedKeys { + c.buildStore.Unpin(key) + } + }() + + deduped, err := memfileHeader.Wait() + if err != nil { + logger.L().Warn(swapCtx, "provisional memfile header swap: deduped header failed", zap.Error(err)) + + return + } + mem, err := storageTemplate.Memfile(swapCtx) + if err != nil { + logger.L().Warn(swapCtx, "provisional memfile header swap: get memfile", zap.Error(err)) + + return + } + if mem == nil { + logger.L().Warn(swapCtx, "provisional memfile header swap: memfile is nil") + + return + } + // Swap only if the header is still the provisional one. Upload.publish + // (and the P2P poll path) install a finalized header unconditionally; + // if this goroutine runs late we must not clobber that newer header + // with the older, still-incomplete deduped one. Either way the + // provisional header is no longer needed, so release + drop below. + if cas, ok := mem.(interface { + SwapHeaderIfCurrent(old, next *header.Header) bool + }); ok { + if !cas.SwapHeaderIfCurrent(provisionalMemfileHeader, deduped) { + logger.L().Info(swapCtx, "provisional memfile header swap: header already advanced; skipping") + } + } else { + mem.SwapHeader(deduped) + } + + // Reads now route off the provisional build id; the deferred signal + // above lets the dedup goroutine release the memfd. The provisional + // store entry is intentionally NOT deleted here: a reader that planned + // on the provisional header but has not yet hit the store would miss and + // fall through to a storage fetch for the synthetic build id (never + // uploaded). It is harmless to leave — no reads route to it post-swap, + // it reports FileSize 0 so it doesn't skew disk eviction, and the store + // TTL reclaims it (its Close is a no-op). + }() + } + return nil } diff --git a/packages/orchestrator/pkg/sandbox/template/peerserver/header.go b/packages/orchestrator/pkg/sandbox/template/peerserver/header.go index d4d56d0267..721fb6a2f8 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerserver/header.go +++ b/packages/orchestrator/pkg/sandbox/template/peerserver/header.go @@ -32,7 +32,23 @@ func (f *headerSource) Stream(ctx context.Context, sender Sender) error { return fmt.Errorf("get device: %w", err) } + // Serve the durable (deduped) header, not the live one. While a provisional + // header is being served locally, device.Header() maps dirty pages to a + // synthetic build id backed only by this node's memfd, which a peer can't + // resolve (it's absent from the peer registry and object storage). + // DurableHeader waits for the deduped header if a swap is still pending; + // devices without one return their current header immediately. h := device.Header() + if dh, ok := device.(interface { + DurableHeader(ctx context.Context) (*header.Header, error) + }); ok { + h, err = dh.DurableHeader(ctx) + if err != nil { + span.RecordError(err) + + return fmt.Errorf("resolve durable header: %w", err) + } + } if h == nil { return ErrNotAvailable } diff --git a/packages/orchestrator/pkg/sandbox/template/storage.go b/packages/orchestrator/pkg/sandbox/template/storage.go index 952d642655..2e2e26ddd6 100644 --- a/packages/orchestrator/pkg/sandbox/template/storage.go +++ b/packages/orchestrator/pkg/sandbox/template/storage.go @@ -15,6 +15,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/utils" ) const ( @@ -156,6 +157,22 @@ func (d *Storage) SwapHeader(h *header.Header) { d.source.SwapHeader(h) } +func (d *Storage) SwapHeaderIfCurrent(old, next *header.Header) bool { + return d.source.SwapHeaderIfCurrent(old, next) +} + +func (d *Storage) SetDurableHeader(f *utils.SetOnce[*header.Header]) { + d.source.SetDurableHeader(f) +} + +func (d *Storage) DurableHeader(ctx context.Context) (*header.Header, error) { + return d.source.DurableHeader(ctx) +} + +func (d *Storage) DurableHeaderNow() (*header.Header, bool) { + return d.source.DurableHeaderNow() +} + func (d *Storage) Close() error { return nil } diff --git a/packages/orchestrator/pkg/sandbox/template/storage_template.go b/packages/orchestrator/pkg/sandbox/template/storage_template.go index dc400fe392..49a939352e 100644 --- a/packages/orchestrator/pkg/sandbox/template/storage_template.go +++ b/packages/orchestrator/pkg/sandbox/template/storage_template.go @@ -36,8 +36,13 @@ type storageTemplate struct { memfileHeader *utils.SetOnce[*header.Header] rootfsHeader *utils.SetOnce[*header.Header] - localSnapfile File - localMetafile File + // durableMemfileHeader, when non-nil, is the header the memfile will settle + // on (the deduped header while a provisional header is served); Fetch wires + // it into the memfile device as its durable header before publishing the + // device, so a pause parents off it rather than the provisional header. + durableMemfileHeader *utils.SetOnce[*header.Header] + localSnapfile File + localMetafile File metrics blockmetrics.Metrics persistence storage.StorageProvider @@ -52,6 +57,7 @@ func newTemplateFromStorage( metrics blockmetrics.Metrics, localSnapfile File, localMetafile File, + durableMemfileHeader *utils.SetOnce[*header.Header], ) (*storageTemplate, error) { paths, err := storage.Paths{ BuildID: buildId, @@ -61,17 +67,18 @@ func newTemplateFromStorage( } return &storageTemplate{ - paths: paths, - localSnapfile: localSnapfile, - localMetafile: localMetafile, - memfileHeader: memfileHeader, - rootfsHeader: rootfsHeader, - metrics: metrics, - persistence: persistence, - memfile: utils.NewSetOnce[block.ReadonlyDevice](), - rootfs: utils.NewSetOnce[block.ReadonlyDevice](), - snapfile: utils.NewSetOnce[File](), - metafile: utils.NewSetOnce[File](), + paths: paths, + localSnapfile: localSnapfile, + localMetafile: localMetafile, + memfileHeader: memfileHeader, + rootfsHeader: rootfsHeader, + durableMemfileHeader: durableMemfileHeader, + metrics: metrics, + persistence: persistence, + memfile: utils.NewSetOnce[block.ReadonlyDevice](), + rootfs: utils.NewSetOnce[block.ReadonlyDevice](), + snapfile: utils.NewSetOnce[File](), + metafile: utils.NewSetOnce[File](), }, nil } @@ -210,6 +217,14 @@ func (t *storageTemplate) Fetch(ctx context.Context, buildStore *build.DiffStore return nil } + // Wire the durable header before publishing the device (SetValue), so no + // caller can observe the memfile with its durable header unset: a pause + // then always parents off the deduped header rather than the provisional + // one being served. + if t.durableMemfileHeader != nil { + memfileStorage.SetDurableHeader(t.durableMemfileHeader) + } + if err := t.memfile.SetValue(memfileStorage); err != nil { return fmt.Errorf("failed to set memfile value: %w", err) } @@ -304,7 +319,22 @@ func (t *storageTemplate) SchedulingMetadata(ctx context.Context) *orchestrator. // dropping the rootfs affinity data too. var mh *header.Header if memfile, memfileErr := t.memfile.WaitWithContext(ctx); memfileErr == nil { - mh = memfile.Header() + // Use the durable (deduped) header, not the live one: during the + // provisional window memfile.Header() maps dirty pages to a synthetic build + // id that is never uploaded or registered, which would put a nonexistent + // layer into scheduling metadata. Non-blocking: this runs on the + // create/resume response path, so while a provisional swap is still pending + // (deduped header not yet known) report rootfs-only metadata rather than + // block the response for the whole dedup. No swap pending → live header. + if dh, ok := memfile.(interface { + DurableHeaderNow() (*header.Header, bool) + }); ok { + if h, ready := dh.DurableHeaderNow(); ready { + mh = h + } + } else { + mh = memfile.Header() + } } return scheduling.FromHeaders(rh.Metadata.BuildId, mh, rh, 0) diff --git a/packages/orchestrator/pkg/server/sandboxes.go b/packages/orchestrator/pkg/server/sandboxes.go index 17f2ac24ee..d92d712038 100644 --- a/packages/orchestrator/pkg/server/sandboxes.go +++ b/packages/orchestrator/pkg/server/sandboxes.go @@ -936,6 +936,9 @@ func (s *Server) snapshotAndCacheSandbox( snapshot.Metafile, snapshot.MemorySnapshot.Diff, snapshot.RootfsDiff, + snapshot.MemorySnapshot.ProvisionalDiffHeader, + snapshot.MemorySnapshot.ProvisionalDiff, + snapshot.MemorySnapshot.ProvisionalSwapDone, ) if err != nil { return nil, fmt.Errorf("error adding snapshot to template cache: %w", err) diff --git a/packages/orchestrator/pkg/template/build/layer/layer_executor.go b/packages/orchestrator/pkg/template/build/layer/layer_executor.go index edc8c3f6cf..147ed7a70e 100644 --- a/packages/orchestrator/pkg/template/build/layer/layer_executor.go +++ b/packages/orchestrator/pkg/template/build/layer/layer_executor.go @@ -283,6 +283,9 @@ func (lb *LayerExecutor) PauseAndUpload( snapshot.Metafile, snapshot.MemorySnapshot.Diff, snapshot.RootfsDiff, + snapshot.MemorySnapshot.ProvisionalDiffHeader, + snapshot.MemorySnapshot.ProvisionalDiff, + snapshot.MemorySnapshot.ProvisionalSwapDone, ) if err != nil { err = errors.Join(err, snapshot.Close(context.WithoutCancel(ctx))) diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index b5b4201025..5a95eebc7c 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -162,8 +162,11 @@ var ( // MemfdDedupInflightServeFlag lets a resume that overlaps an in-flight // memfile dedup serve dirty pages straight from the still-mapped memfd - // instead of blocking on the dedup drain. Only affects the memfd-dedup - // path; off restores the prior wait-for-drain behavior. + // instead of blocking until dedup finishes. It gates both windows: serving + // via a provisional local header while dedup is still computing the deduped + // header, and serving during the dedup drain before the compacted diff is + // ready. Only affects the memfd-dedup path; off restores the prior + // wait-for-dedup behavior. MemfdDedupInflightServeFlag = NewBoolFlag("memfd-dedup-inflight-serve", false) // PeerToPeerChunkTransferFlag enables peer-to-peer chunk routing. From 629461088463588e673af992392b7bdbc7b928ef Mon Sep 17 00:00:00 2001 From: Nikita Kalyazin Date: Thu, 2 Jul 2026 14:21:20 +0100 Subject: [PATCH 9/9] test(orch): cover MemfdIdentitySource.Slice and provisional fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Co-Authored-By: Claude Opus 4.8 (1M context) --- .../pkg/sandbox/block/memfd_test.go | 8 +- .../pkg/sandbox/build/cache_test.go | 31 +++++ .../pkg/sandbox/build/durable_header_test.go | 110 ++++++++++++++++++ .../pkg/sandbox/provisional_test.go | 41 +++++++ .../template/peerserver/header_test.go | 43 +++++++ .../sandbox/template/storage_template_test.go | 82 +++++++++++++ 6 files changed, 314 insertions(+), 1 deletion(-) create mode 100644 packages/orchestrator/pkg/sandbox/build/durable_header_test.go create mode 100644 packages/orchestrator/pkg/sandbox/provisional_test.go create mode 100644 packages/orchestrator/pkg/sandbox/template/storage_template_test.go diff --git a/packages/orchestrator/pkg/sandbox/block/memfd_test.go b/packages/orchestrator/pkg/sandbox/block/memfd_test.go index bc8b59ce6a..911013527e 100644 --- a/packages/orchestrator/pkg/sandbox/block/memfd_test.go +++ b/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)) } diff --git a/packages/orchestrator/pkg/sandbox/build/cache_test.go b/packages/orchestrator/pkg/sandbox/build/cache_test.go index c7aee36a3f..d6a1cfd227 100644 --- a/packages/orchestrator/pkg/sandbox/build/cache_test.go +++ b/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() diff --git a/packages/orchestrator/pkg/sandbox/build/durable_header_test.go b/packages/orchestrator/pkg/sandbox/build/durable_header_test.go new file mode 100644 index 0000000000..b0e601acff --- /dev/null +++ b/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) +} diff --git a/packages/orchestrator/pkg/sandbox/provisional_test.go b/packages/orchestrator/pkg/sandbox/provisional_test.go new file mode 100644 index 0000000000..d105c4632a --- /dev/null +++ b/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) +} diff --git a/packages/orchestrator/pkg/sandbox/template/peerserver/header_test.go b/packages/orchestrator/pkg/sandbox/template/peerserver/header_test.go index 30fd502e52..d0728a275b 100644 --- a/packages/orchestrator/pkg/sandbox/template/peerserver/header_test.go +++ b/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() diff --git a/packages/orchestrator/pkg/sandbox/template/storage_template_test.go b/packages/orchestrator/pkg/sandbox/template/storage_template_test.go new file mode 100644 index 0000000000..10d99a1121 --- /dev/null +++ b/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()) +}