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.

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,35 @@ 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
+}
+
 // 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

@@ -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,35 +225,129 @@ 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.
+	// 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)
 	}
 
 	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 {
+		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
+			}
+			// 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. Let the dedup goroutine
+			// release the memfd (it keeps it mapped until this fires or a grace
+			// elapses). 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).
+			if provisionalSwapDone != nil {
+				provisionalSwapDone()
+			}
+		}()
+	}
+
 	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,18 @@ 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) Close() error {
 	return nil
 }