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 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.

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/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,26 @@ func (b *File) SwapHeader(h *header.Header) {
 	b.header.Store(h)
 }
 
+// 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
+}
+
 // 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

@@ -176,6 +176,13 @@ func (s *DiffStore) Has(d Diff) bool {
 	return s.cache.Has(d.CacheKey())
 }
 
+// Delete evicts a single entry by key, running the eviction callback (which
+// closes the Diff). Used to drop the short-lived provisional memfile entry once
+// the deduped header has been swapped in, so it doesn't linger for the TTL.
+func (s *DiffStore) Delete(key DiffStoreKey) {
+	s.cache.Delete(key)
+}
+
 // Lookup returns the cached Diff for the given key without initialising a new one.
 // Returns (nil, false) if the key is not present in the cache.
 func (s *DiffStore) Lookup(key DiffStoreKey) (Diff, bool) {

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,46 @@ 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 || !enabled || originalMemfile == nil || originalHeader == nil || diffMetadata == nil {
+		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))
+
+		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))
+
+		return nil, nil, nil
+	}
+
+	return provisionalHeader, provisionalDiff, dc.MarkSwapped
 }
 
 func pauseProcessRootfs(

packages/orchestrator/pkg/sandbox/template/cache.go

@@ -224,23 +224,56 @@ 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 {
+			// Assumption: this provisional entry stays resident in the build store
+			// until the SwapHeader below (i.e. for the provisional window — the
+			// dedup compare). 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 could not be rebuilt (createDiff has nothing to fetch),
+			// failing the read. This holds because the store TTL (hours) dwarfs the
+			// window (seconds) and eviction is LRU-by-TTL, so this freshly-added
+			// entry is evicted last. We rely on that rather than guarding the miss:
+			// the eviction can't realistically occur, so a fallback isn't warranted.
+			c.buildStore.Add(provisionalMemfileDiff)
+		}
+	}
 
 	switch rootfsDiff.(type) {
 	case *build.NoDiff:
 	default:
 		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.
+	localMemfileHeader := memfileHeader
+	if provisionalMemfileHeader != nil {
+		localMemfileHeader = resolvedHeader(provisionalMemfileHeader)
+	}
+
 	storageTemplate, err := newTemplateFromStorage(
 		c.config.BuilderConfig,
 		buildId,
-		memfileHeader,
+		localMemfileHeader,
 		rootfsHeader,
 		c.persistence,
 		c.blockMetrics,
@@ -253,6 +286,59 @@ 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.
+	if provisionalMemfileHeader != nil {
+		// Record the deduped header as the memfile's durable header (before this
+		// returns, i.e. before any resume can re-pause) so a pause that parents
+		// off this still-provisional template waits for the deduped header rather
+		// than persisting the provisional header's synthetic build-id mappings,
+		// which have no storage object and are unreadable on a cold/cross-node
+		// resume.
+		if mem, memErr := storageTemplate.Memfile(ctx); memErr != nil {
+			logger.L().Warn(ctx, "provisional memfile: set durable header: get memfile", zap.Error(memErr))
+		} else if dh, ok := mem.(interface {
+			SetDurableHeader(f *utils.SetOnce[*header.Header])
+		}); ok {
+			dh.SetDurableHeader(memfileHeader)
+		}
+
+		swapCtx := context.WithoutCancel(ctx)
+		go func() {
+			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
+			}
+			mem.SwapHeader(deduped)
+
+			// Reads now route to the deduped build id. Let the dedup goroutine
+			// release the memfd (it keeps it mapped until this fires or a grace
+			// elapses), and drop the provisional entry so it doesn't linger in the
+			// store for the TTL. On an earlier error return the provisional entry
+			// stays as the only memfile source until the grace releases the memfd.
+			if provisionalSwapDone != nil {
+				provisionalSwapDone()
+			}
+			if provisionalMemfileDiff != nil {
+				c.buildStore.Delete(provisionalMemfileDiff.CacheKey())
+			}
+		}()
+	}
+
 	return nil
 }
 

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,14 @@ func (d *Storage) SwapHeader(h *header.Header) {
 	d.source.SwapHeader(h)
 }
 
+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) Close() error {
 	return nil
 }

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)

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

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.