diff --git a/packages/orchestrator/pkg/sandbox/block/memfd.go b/packages/orchestrator/pkg/sandbox/block/memfd.go index 7bc59793f0..30fb71d98e 100644 --- a/packages/orchestrator/pkg/sandbox/block/memfd.go +++ b/packages/orchestrator/pkg/sandbox/block/memfd.go @@ -7,10 +7,13 @@ import ( "errors" "fmt" "os" + "sync" "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" @@ -232,12 +235,47 @@ 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 + // 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 + +// 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, @@ -250,21 +288,80 @@ 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](), + 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) 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, @@ -286,7 +383,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 } @@ -305,17 +402,55 @@ 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()))) + + // 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) + d.memfd = memfd + 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) - if closeErr := memfd.Close(); 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. 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") + } + } + 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 @@ -331,7 +466,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 +513,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 @@ -349,6 +528,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() @@ -360,6 +573,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 { diff --git a/packages/orchestrator/pkg/sandbox/block/memfd_test.go b/packages/orchestrator/pkg/sandbox/block/memfd_test.go index d7e4a66278..911013527e 100644 --- a/packages/orchestrator/pkg/sandbox/block/memfd_test.go +++ b/packages/orchestrator/pkg/sandbox/block/memfd_test.go @@ -6,9 +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" @@ -227,7 +232,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() }) @@ -248,3 +253,229 @@ 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) +} + +// 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. +// 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() + + 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: + } +} + +// 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) + } + // 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()) + + 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/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/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/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/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/sandbox.go b/packages/orchestrator/pkg/sandbox/sandbox.go index b7c30b2fbb..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, @@ -1494,17 +1519,21 @@ 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) } 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 } @@ -1537,7 +1566,8 @@ func pauseProcessMemory( dedupBestEffort bool, dedupDirectIO bool, dedupBudget block.DedupBudget, -) (d build.Diff, h *DiffHeader, e error) { + dedupInflightServe bool, +) (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() @@ -1547,9 +1577,10 @@ 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) + return nil, nil, nil, nil, nil, fmt.Errorf("failed to export memory: %w", err) } diff, err := build.NewLocalDiffFromCache( @@ -1557,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. @@ -1586,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/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.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/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()) +} 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 3a68f14f83..5a95eebc7c 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -160,6 +160,15 @@ 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 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. PeerToPeerChunkTransferFlag = NewBoolFlag("peer-to-peer-chunk-transfer", false) // PeerToPeerAsyncCheckpointFlag makes Checkpoint upload fire-and-forget instead 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, 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) + } +}