Export rootfs in the background during pause()/snapshot()

docs/ARCHITECTURE.md

@@ -312,6 +312,14 @@ sequenceDiagram
   orchestrator pauses the VM, snapshots it, diffs memory (dirty-page tracking) and rootfs (COW
   cache) against the template, caches the snapshot locally, and uploads asynchronously to object
   storage (with a retry budget). The sandbox leaves the Redis catalog.
+  - **Deferred rootfs export** (gated by the `deferred-rootfs-export` flag in
+    `packages/shared/pkg/featureflags`): instead of diffing the rootfs on the pause critical
+    path, the orchestrator ejects the writable COW cache during pause and returns, then seals it
+    into the rootfs diff (reflink) in the background. This moves the rootfs-diff latency off the
+    pause, but the local snapshot's rootfs body isn't materialized until the seal finishes, so the
+    async upload — and any origin-node resume/prefetch that reads the rootfs diff — waits on the
+    seal. A seal failure is permanent (it never re-runs), so the upload fails fast rather than
+    retrying.
 - **Resume**: same path as creation, but placement prefers the **origin node** — if the snapshot
   is still in its local cache, resume avoids any object-storage reads. `Checkpoint` is a
   pause+resume in place used to persist state while keeping the sandbox running.

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

@@ -106,18 +106,62 @@ func (c *Cache) isClosed() bool {
 	return c.closed.Load()
 }
 
+// DiffMetadata returns the dirty/empty diff metadata from the tracker without
+// copying any block data. It lets a deferred/background seal build the diff
+// header (and scheduling metadata) synchronously while the actual reflink copy
+// (ExportToDiff) runs off the critical path. The result matches what
+// ExportToDiff computes internally, provided the cache is frozen (no writes) in
+// between — which is guaranteed once the sandbox has been stopped and the cache
+// ejected.
+func (c *Cache) DiffMetadata() (*header.DiffMetadata, error) {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	if c.isClosed() {
+		return nil, NewErrCacheClosed(c.filePath)
+	}
+
+	if c.mmap == nil {
+		return header.NewDiffMetadata(c.blockSize, nil, nil), nil
+	}
+
+	dirty, empty := c.tracker.Export()
+
+	return header.NewDiffMetadata(c.blockSize, dirty, empty), nil
+}
+
 func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMetadata, error) {
-	ctx, childSpan := tracer.Start(ctx, "export-to-diff")
-	defer childSpan.End()
+	c.mu.Lock()
+	defer c.mu.Unlock()
 
+	return c.exportToDiffLocked(ctx, out, nil)
+}
+
+// ExportToDiffWithMetadata copies the dirty ranges described by meta to out,
+// using meta's bitmap instead of re-reading the tracker. Callers that captured a
+// DiffMetadata earlier (e.g. the deferred rootfs seal, which reads it at setup
+// and exports later in the background) use this so the copied ranges are
+// guaranteed to match a header built from the same bitmap read.
+func (c *Cache) ExportToDiffWithMetadata(ctx context.Context, out *os.File, meta *header.DiffMetadata) (*header.DiffMetadata, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	return c.exportToDiffLocked(ctx, out, meta)
+}
+
+func (c *Cache) exportToDiffLocked(ctx context.Context, out *os.File, meta *header.DiffMetadata) (*header.DiffMetadata, error) {
+	ctx, childSpan := tracer.Start(ctx, "export-to-diff")
+	defer childSpan.End()
+
 	if c.isClosed() {
 		return nil, NewErrCacheClosed(c.filePath)
 	}
 
 	if c.mmap == nil {
+		if meta != nil {
+			return meta, nil
+		}
+
 		return header.NewDiffMetadata(c.blockSize, nil, nil), nil
 	}
 
@@ -137,8 +181,11 @@ func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMet
 		logger.L().Warn(ctx, "error syncing file", zap.Error(err))
 	}
 
-	dirty, empty := c.tracker.Export()
-	diffMetadata := header.NewDiffMetadata(c.blockSize, dirty, empty)
+	diffMetadata := meta
+	if diffMetadata == nil {
+		dirty, empty := c.tracker.Export()
+		diffMetadata = header.NewDiffMetadata(c.blockSize, dirty, empty)
+	}
 
 	dst := int(out.Fd())
 	var writeOffset int64

packages/orchestrator/pkg/sandbox/block/cache_diffmetadata_test.go

@@ -0,0 +1,55 @@
+//go:build linux
+
+package block
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+)
+
+// TestCacheDiffMetadata_MatchesExportToDiff verifies DiffMetadata (used to build
+// the diff header synchronously) produces the same dirty/empty bitmaps as the
+// metadata ExportToDiff computes while copying, so a background seal's header is
+// exact.
+func TestCacheDiffMetadata_MatchesExportToDiff(t *testing.T) {
+	t.Parallel()
+
+	blockSize := int64(header.PageSize)
+	numBlocks := int64(4)
+
+	c, err := NewCache(blockSize*numBlocks, blockSize, t.TempDir()+"/cache", false)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = c.Close() })
+
+	// Dirty block 0, zero block 2, leave 1 and 3 untouched.
+	dirtyBlock := make([]byte, blockSize)
+	for i := range dirtyBlock {
+		dirtyBlock[i] = 0xAB
+	}
+	_, err = c.WriteAt(dirtyBlock, 0)
+	require.NoError(t, err)
+	_, err = c.WriteZeroesAt(2*blockSize, blockSize)
+	require.NoError(t, err)
+
+	// Metadata read without copying.
+	meta, err := c.DiffMetadata()
+	require.NoError(t, err)
+
+	// Metadata computed by the actual export.
+	out, err := os.CreateTemp(t.TempDir(), "diff-*")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = out.Close() })
+
+	exportMeta, err := c.ExportToDiff(t.Context(), out)
+	require.NoError(t, err)
+
+	require.True(t, meta.Dirty.Equals(exportMeta.Dirty), "dirty bitmaps must match")
+	require.True(t, meta.Empty.Equals(exportMeta.Empty), "empty bitmaps must match")
+	require.Equal(t, exportMeta.BlockSize, meta.BlockSize)
+	require.EqualValues(t, 1, meta.Dirty.GetCardinality())
+	require.EqualValues(t, 1, meta.Empty.GetCardinality())
+}

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

@@ -98,6 +98,22 @@ func NewDiffStore(
 		// buildData will be deleted by calling buildData.Close()
 		defer ds.resetDelete(item.Key())
 
+		// A deferred diff whose background seal hasn't resolved would block Close()
+		// on the reflink; close it off the eviction goroutine so this callback
+		// isn't stalled. Bounded (the seal always resolves) and only reachable if
+		// the TTL is shortened below the seal time — the disk-pressure eviction
+		// path already skips unsealed diffs (see deferredDiff.sealed()).
+		if dd, ok := buildData.(*deferredDiff); ok && !dd.sealed() {
+			logCtx := context.WithoutCancel(ctx)
+			go func() {
+				if closeErr := dd.Close(); closeErr != nil {
+					logger.L().Warn(logCtx, "failed to close unsealed deferred diff", zap.Any("item_key", item.Key()), zap.Error(closeErr))
+				}
+			}()
+
+			return
+		}
+
 		if closeErr := buildData.Close(); closeErr != nil {
 			logger.L().Warn(ctx, "failed to cleanup build data cache for item", zap.Any("item_key", item.Key()), zap.Error(closeErr))
 		}
@@ -311,6 +327,14 @@ func (s *DiffStore) deleteOldestFromCache(ctx context.Context) (suc bool, e erro
 			return true
 		}
 
+		// Skip a deferred diff whose background rootfs seal hasn't resolved yet:
+		// FileSize below would block on the seal (stalling the sole eviction
+		// goroutine), and a fresh, still-sealing snapshot is exactly what a
+		// just-resumed peer needs. It becomes evictable once the seal resolves.
+		if dd, ok := item.Value().(*deferredDiff); ok && !dd.sealed() {
+			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))

packages/orchestrator/pkg/sandbox/build/cache_test.go

@@ -31,6 +31,7 @@ import (
 	blockmetrics "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block/metrics"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
 const (
@@ -716,3 +717,43 @@ func mustParseCfg(t *testing.T) cfg.Config {
 
 	return c
 }
+
+// A deferred rootfs diff whose background seal hasn't resolved must be skipped
+// by disk-pressure eviction: its data methods (FileSize) block on the seal, so
+// evicting it would stall the sole eviction goroutine, and a fresh, still-sealing
+// snapshot is exactly what a just-resumed peer needs. Once the seal resolves the
+// entry becomes evictable like any other.
+func TestDiffStoreEvictionSkipsUnsealedDeferredDiff(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 entry: a deferred rootfs diff whose background seal hasn't resolved.
+	promise := utils.NewSetOnce[Diff]()
+	deferred := NewDeferredDiff(GetDiffStoreKey("deferred-unsealed", Rootfs), blockSize, promise)
+	store.Add(deferred)
+
+	// A newer, ordinary diff behind it.
+	newer := newRootFSDiff(t, cachePath, "seal-newer")
+	store.Add(newer)
+
+	// Eviction must skip the unsealed deferred diff and fall through to the
+	// next-oldest instead of blocking on the seal.
+	ok, err := store.deleteOldestFromCache(t.Context())
+	require.NoError(t, err)
+	assert.True(t, ok)
+	assert.False(t, store.isBeingDeleted(deferred.CacheKey()), "unsealed deferred diff must be skipped")
+	assert.True(t, store.isBeingDeleted(newer.CacheKey()), "eviction must fall through to the next-oldest")
+
+	// Once the seal resolves, the deferred diff is evictable like any other entry.
+	require.NoError(t, promise.SetValue(newRootFSDiff(t, cachePath, "seal-resolved")))
+	ok, err = store.deleteOldestFromCache(t.Context())
+	require.NoError(t, err)
+	assert.True(t, ok)
+	assert.True(t, store.isBeingDeleted(deferred.CacheKey()), "sealed deferred diff must be evictable")
+}

packages/orchestrator/pkg/sandbox/build/deferred_diff.go

@@ -0,0 +1,123 @@
+//go:build linux
+
+package build
+
+import (
+	"context"
+	"errors"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
+)
+
+// ErrDeferredSealFailed marks a deferred rootfs seal that failed. The seal runs
+// exactly once, so the failure is permanent: the promise stays settled with this
+// error and every subsequent data-bearing method on the diff returns it. Callers
+// such as the pause-upload retry loop match on it to stop retrying a diff that
+// can never materialize (rather than burning the whole retry budget on it).
+var ErrDeferredSealFailed = errors.New("deferred rootfs seal failed")
+
+// deferredDiff is a Diff whose backing data is produced asynchronously. It is
+// returned synchronously from a pause that seals the rootfs in the background:
+// the cache key and block size are known up front (so DiffStore.Add and the
+// upload's compress-config validation work immediately), while every data-
+// bearing method blocks on the inner promise until the background seal resolves
+// the real Diff.
+//
+// The producer MUST always resolve the promise — with the sealed Diff on success
+// or an error on failure — otherwise the data methods (and Close) block forever.
+type deferredDiff struct {
+	cacheKey  DiffStoreKey
+	blockSize int64
+	inner     *utils.SetOnce[Diff]
+}
+
+var _ Diff = (*deferredDiff)(nil)
+
+// NewDeferredDiff wraps a promise of a Diff. cacheKey and blockSize are the
+// synchronously-known identity of the diff; inner is resolved by the background
+// sealer with the materialized Diff (or an error).
+func NewDeferredDiff(cacheKey DiffStoreKey, blockSize int64, inner *utils.SetOnce[Diff]) Diff {
+	return &deferredDiff{
+		cacheKey:  cacheKey,
+		blockSize: blockSize,
+		inner:     inner,
+	}
+}
+
+func (d *deferredDiff) CacheKey() DiffStoreKey {
+	return d.cacheKey
+}
+
+func (d *deferredDiff) BlockSize() int64 {
+	return d.blockSize
+}
+
+// sealed reports, without blocking, whether the background seal has resolved the
+// promise (with the materialized diff or an error). Cache eviction uses it to
+// skip a not-yet-sealed diff rather than block on the data methods (which wait
+// on the seal) or evict a fresh, still-in-flight snapshot.
+func (d *deferredDiff) sealed() bool {
+	select {
+	case <-d.inner.Done:
+		return true
+	default:
+		return false
+	}
+}
+
+func (d *deferredDiff) CachePath(ctx context.Context) (string, error) {
+	inner, err := d.inner.WaitWithContext(ctx)
+	if err != nil {
+		return "", err
+	}
+
+	return inner.CachePath(ctx)
+}
+
+func (d *deferredDiff) ReadAt(ctx context.Context, p []byte, off int64, ft *storage.FrameTable) (int, error) {
+	inner, err := d.inner.WaitWithContext(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	return inner.ReadAt(ctx, p, off, ft)
+}
+
+func (d *deferredDiff) Slice(ctx context.Context, off, length int64, ft *storage.FrameTable) ([]byte, error) {
+	inner, err := d.inner.WaitWithContext(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return inner.Slice(ctx, off, length, ft)
+}
+
+func (d *deferredDiff) Size(ctx context.Context) (int64, error) {
+	inner, err := d.inner.WaitWithContext(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	return inner.Size(ctx)
+}
+
+func (d *deferredDiff) FileSize(ctx context.Context) (int64, error) {
+	inner, err := d.inner.WaitWithContext(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	return inner.FileSize(ctx)
+}
+
+// Close waits for the seal to resolve and closes the materialized diff. If the
+// seal failed there is nothing to close (the producer cleans up the partial file
+// on error), so only close when the diff actually materialized.
+func (d *deferredDiff) Close() error {
+	if inner, err := d.inner.Wait(); err == nil {
+		return inner.Close()
+	}
+
+	return nil
+}