feat(orch): serve resume from a provisional header during dedup

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 <nikita.kalyazin@e2b.dev>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: kalyazin

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")
 		}
 	}

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

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

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,17 +1580,24 @@ 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(
 		build.GetDiffStoreKey(buildID.String(), build.Memfile),
 		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(