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.

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/sandbox.go

@@ -1435,6 +1435,13 @@ 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
 	// 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
@@ -1481,7 +1488,7 @@ func (s *Sandbox) processMemorySnapshot(ctx context.Context, buildID uuid.UUID)
 		}
 	}
 
-	memfileDiff, memfileDiffHeader, err := pauseProcessMemory(
+	memfileDiff, memfileDiffHeader, provMemfileHeader, provMemfileDiff, err := pauseProcessMemory(
 		ctx,
 		buildID,
 		memfileHeader,
@@ -1501,11 +1508,13 @@ 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,
+		BlockSize:             memfileHeader.Metadata.BlockSize,
+		header:                memfileHeader,
+		newBytes:              memfileDiffMetadata.Dirty.GetCardinality() * uint64(memfileDiffMetadata.BlockSize),
 	}, nil
 }
 
@@ -1539,7 +1548,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, e error) {
 	ctx, span := tracer.Start(ctx, "process-memory")
 	defer span.End()
 
@@ -1552,17 +1561,24 @@ func pauseProcessMemory(
 		dedupInflightServe,
 	)
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to export memory: %w", err)
+		return 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, 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 = 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 +1605,44 @@ func pauseProcessMemory(
 		setHeader(meta.ToDiffHeader(ctx, originalHeader, buildID))
 	}()
 
-	return diff, headerOut, nil
+	return diff, headerOut, provisionalHeader, provisionalDiff, nil
+}
+
+// buildProvisionalMemfile builds the provisional local header + its memfd-backed
+// diff source. Returns (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) {
+	dc, ok := cache.(*block.DedupedMemfdCache)
+	if !ok || !enabled || originalMemfile == nil || originalHeader == nil || diffMetadata == nil {
+		return 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
+	}
+
+	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
+	}
+
+	return provisionalHeader, provisionalDiff
 }
 
 func pauseProcessRootfs(

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

@@ -224,23 +224,52 @@ 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,
 ) 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 +282,33 @@ 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 {
+		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)
+		}()
+	}
+
 	return nil
 }
 

packages/orchestrator/pkg/server/sandboxes.go

@@ -936,6 +936,8 @@ func (s *Server) snapshotAndCacheSandbox(
 		snapshot.Metafile,
 		snapshot.MemorySnapshot.Diff,
 		snapshot.RootfsDiff,
+		snapshot.MemorySnapshot.ProvisionalDiffHeader,
+		snapshot.MemorySnapshot.ProvisionalDiff,
 	)
 	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,8 @@ func (lb *LayerExecutor) PauseAndUpload(
 		snapshot.Metafile,
 		snapshot.MemorySnapshot.Diff,
 		snapshot.RootfsDiff,
+		snapshot.MemorySnapshot.ProvisionalDiffHeader,
+		snapshot.MemorySnapshot.ProvisionalDiff,
 	)
 	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.