feat(memfd): async background copy via NewCacheFromMemfdAsync

ValentaTomas · ValentaTomas · commit 402c402ca3a7 · 2026-05-16T23:34:00.000-07:00
Introduce the MemfdCache wrapper (embedding *Cache) plus the
NewCacheFromMemfdAsync constructor: copy runs on a goroutine so gRPC
Pause can return as soon as the snapshot file and diff metadata are
written. The MemfdBackgroundCopyFlag gates the dispatch in
fc.ExportMemory; flag-off keeps the existing sync NewCacheFromMemfd
path untouched.

In-flight reads route through a memfdSource indexed by cache offset;
afterwards they delegate to the embedded Cache and the memfd is
closed. Slice returns BytesNotAvailableError while the copy is in
flight to prevent UAF on the asynchronous Munmap; callers fall back
to ReadAt or Wait first. localDiff takes block.DiffSource and uses a
Wait type-assertion in CachePath so existing FS-reading upload paths
see complete data.
diff --git a/packages/orchestrator/pkg/sandbox/block/device.go b/packages/orchestrator/pkg/sandbox/block/device.go
@@ -46,3 +46,16 @@ type Device interface {
 	io.WriterAt
 	WriteZeroesAt(off, length int64) (int, error)
 }
+
+// DiffSource is what the diff/upload layer reads from. *Cache satisfies it
+// directly; the async memfd path wraps *Cache in *MemfdCache to override
+// the read methods while the background copy is in flight.
+type DiffSource interface {
+	io.Closer
+	ReadAt(b []byte, off int64) (int, error)
+	Slice(off, length int64) ([]byte, error)
+	Size() (int64, error)
+	FileSize() (int64, error)
+	BlockSize() int64
+	Path() string
+}
diff --git a/packages/orchestrator/pkg/sandbox/block/memfd.go b/packages/orchestrator/pkg/sandbox/block/memfd.go
@@ -6,6 +6,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync"
+	"sync/atomic"
 
 	"github.com/RoaringBitmap/roaring/v2"
 	"golang.org/x/sys/unix"
@@ -73,26 +75,219 @@ func NewCacheFromMemfd(
 	if err != nil {
 		return nil, errors.Join(err, memfd.Close())
 	}
+	if err := copyFromMemfd(ctx, cache, memfd, dirty, blockSize); err != nil {
+		return nil, errors.Join(err, memfd.Close(), cache.Close())
+	}
+	if err := memfd.Close(); err != nil {
+		return nil, errors.Join(fmt.Errorf("close memfd: %w", err), cache.Close())
+	}
 
+	return cache, nil
+}
+
+func copyFromMemfd(ctx context.Context, cache *Cache, memfd *Memfd, dirty *roaring.Bitmap, blockSize int64) error {
 	var cacheOff int64
 	for r := range BitsetRanges(dirty, blockSize) {
 		if err := ctx.Err(); err != nil {
-			return nil, errors.Join(err, memfd.Close(), cache.Close())
+			return err
 		}
 
 		src, err := memfd.Slice(r.Start, r.Size)
 		if err != nil {
-			return nil, errors.Join(fmt.Errorf("memfd slice [%d,%d): %w", r.Start, r.Start+r.Size, err), memfd.Close(), cache.Close())
+			return fmt.Errorf("memfd slice [%d,%d): %w", r.Start, r.Start+r.Size, err)
 		}
 
 		copy((*cache.mmap)[cacheOff:cacheOff+r.Size], src)
 		cache.setIsCached(cacheOff, r.Size)
 		cacheOff += r.Size
 	}
 
-	if err := memfd.Close(); err != nil {
-		return nil, errors.Join(fmt.Errorf("close memfd: %w", err), cache.Close())
+	return nil
+}
+
+// MemfdCache wraps a Cache that is being populated from a memfd on a
+// background goroutine. While the copy is in flight, reads are served
+// directly from the memfd via memfdSource; afterwards they delegate to the
+// embedded Cache and the memfd is closed.
+type MemfdCache struct {
+	*Cache
+
+	mu     sync.RWMutex // guards src
+	src    *memfdSource // non-nil while the background copy is in flight
+	cancel context.CancelFunc
+	done   chan struct{}
+	err    atomic.Pointer[error]
+}
+
+// NewCacheFromMemfdAsync starts the memfd→cache copy on a goroutine so gRPC
+// Pause can return as soon as the snapshot file and diff metadata are
+// written; the FC stop and the memfd copy then run in parallel after the
+// response. The returned wrapper takes ownership of memfd; Close cancels
+// and joins the copy goroutine.
+func NewCacheFromMemfdAsync(
+	ctx context.Context,
+	blockSize int64,
+	filePath string,
+	memfd *Memfd,
+	dirty *roaring.Bitmap,
+) (*MemfdCache, error) {
+	cache, err := NewCache(int64(dirty.GetCardinality())*blockSize, blockSize, filePath, false)
+	if err != nil {
+		return nil, errors.Join(err, memfd.Close())
 	}
+	if dirty.IsEmpty() {
+		if closeErr := memfd.Close(); closeErr != nil {
+			return nil, errors.Join(fmt.Errorf("close memfd: %w", closeErr), cache.Close())
+		}
 
-	return cache, nil
+		return &MemfdCache{Cache: cache}, nil
+	}
+
+	// Detach from the request context so the copy can outlive Pause; Close
+	// drives cancellation.
+	copyCtx, cancel := context.WithCancel(context.WithoutCancel(ctx))
+	m := &MemfdCache{
+		Cache:  cache,
+		src:    newMemfdSource(memfd, dirty, blockSize),
+		cancel: cancel,
+		done:   make(chan struct{}),
+	}
+
+	go m.runCopy(copyCtx, dirty, blockSize)
+
+	return m, nil
+}
+
+func (m *MemfdCache) runCopy(ctx context.Context, dirty *roaring.Bitmap, blockSize int64) {
+	defer close(m.done)
+
+	err := copyFromMemfd(ctx, m.Cache, m.src.memfd, dirty, blockSize)
+	if err != nil {
+		m.err.Store(&err)
+	}
+
+	m.mu.Lock()
+	src := m.src
+	m.src = nil
+	m.mu.Unlock()
+
+	if closeErr := src.memfd.Close(); closeErr != nil {
+		joined := errors.Join(err, fmt.Errorf("close memfd: %w", closeErr))
+		m.err.Store(&joined)
+	}
+}
+
+// Wait blocks until the background copy completes (or ctx is cancelled).
+func (m *MemfdCache) Wait(ctx context.Context) error {
+	if m.done == nil {
+		return nil
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-m.done:
+	}
+	if errPtr := m.err.Load(); errPtr != nil {
+		return *errPtr
+	}
+
+	return nil
+}
+
+func (m *MemfdCache) ReadAt(b []byte, off int64) (int, error) {
+	m.mu.RLock()
+	if m.src != nil {
+		defer m.mu.RUnlock()
+
+		return m.src.readAt(b, off)
+	}
+	m.mu.RUnlock()
+
+	return m.Cache.ReadAt(b, off)
+}
+
+// Slice returns BytesNotAvailableError while the copy is in flight: the
+// memfd-backed slice would outlive the RLock and could be Munmap'd
+// asynchronously. Callers fall back to ReadAt or Wait first.
+func (m *MemfdCache) Slice(off, length int64) ([]byte, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	if m.src != nil {
+		return nil, BytesNotAvailableError{}
+	}
+
+	return m.Cache.Slice(off, length)
+}
+
+func (m *MemfdCache) Close() error {
+	if m.cancel != nil {
+		m.cancel()
+		<-m.done
+	}
+
+	return m.Cache.Close()
+}
+
+// memfdSource indexes the memfd-backed ranges by cache offset so reads can
+// be served from the memfd while the background copy is still in flight.
+type memfdSource struct {
+	memfd   *Memfd
+	entries []memfdRange
+}
+
+type memfdRange struct {
+	cacheStart int64
+	srcStart   int64
+	size       int64
+}
+
+func newMemfdSource(memfd *Memfd, dirty *roaring.Bitmap, blockSize int64) *memfdSource {
+	var entries []memfdRange
+	var cacheOff int64
+	for r := range BitsetRanges(dirty, blockSize) {
+		entries = append(entries, memfdRange{cacheStart: cacheOff, srcStart: r.Start, size: r.Size})
+		cacheOff += r.Size
+	}
+
+	return &memfdSource{memfd: memfd, entries: entries}
+}
+
+func (s *memfdSource) findEntry(cacheOff int64) int {
+	lo, hi := 0, len(s.entries)
+	for lo < hi {
+		mid := (lo + hi) / 2
+		if s.entries[mid].cacheStart > cacheOff {
+			hi = mid
+		} else {
+			lo = mid + 1
+		}
+	}
+	i := lo - 1
+	if i < 0 || cacheOff >= s.entries[i].cacheStart+s.entries[i].size {
+		return -1
+	}
+
+	return i
+}
+
+func (s *memfdSource) readAt(b []byte, cacheOff int64) (int, error) {
+	n := 0
+	for n < len(b) {
+		i := s.findEntry(cacheOff + int64(n))
+		if i < 0 {
+			return n, nil
+		}
+		e := s.entries[i]
+		offsetInEntry := cacheOff + int64(n) - e.cacheStart
+		toCopy := min(int64(len(b)-n), e.size-offsetInEntry)
+		src, err := s.memfd.Slice(e.srcStart+offsetInEntry, toCopy)
+		if err != nil {
+			return n, fmt.Errorf("memfd slice: %w", err)
+		}
+		copy(b[n:n+int(toCopy)], src)
+		n += int(toCopy)
+	}
+
+	return n, nil
 }
diff --git a/packages/orchestrator/pkg/sandbox/block/memfd_test.go b/packages/orchestrator/pkg/sandbox/block/memfd_test.go
@@ -3,7 +3,9 @@
 package block
 
 import (
+	"context"
 	"crypto/rand"
+	"os"
 	"testing"
 
 	"github.com/RoaringBitmap/roaring/v2"
@@ -76,3 +78,57 @@ func TestNewCacheFromMemfd_NonZeroRangeStart(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, expected[pageSize*3:pageSize*5], got)
 }
+
+// Cancelling the parent ctx after construction must not abort the in-flight
+// copy: NewCacheFromMemfdAsync detaches via context.WithoutCancel so the
+// copy outlives the Pause RPC. Cancellation goes through Close.
+func TestNewCacheFromMemfdAsync_ParentContextCancellationDoesNotAbort(t *testing.T) {
+	t.Parallel()
+
+	pageSize := int64(header.PageSize)
+	numPages := uint32(16)
+	memfd, expected := newTestMemfd(t, pageSize*int64(numPages))
+
+	dirty := roaring.New()
+	dirty.AddRange(0, uint64(numPages))
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cache, err := NewCacheFromMemfdAsync(ctx, pageSize, t.TempDir()+"/cache", memfd, dirty)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = cache.Close() })
+
+	cancel()
+	require.NoError(t, cache.Wait(t.Context()))
+
+	got := make([]byte, len(expected))
+	_, err = cache.ReadAt(got, 0)
+	require.NoError(t, err)
+	require.Equal(t, expected, got)
+}
+
+// After Wait, the cache file on disk has the full payload and Slice no
+// longer returns BytesNotAvailableError.
+func TestNewCacheFromMemfdAsync_WaitFlushesToFile(t *testing.T) {
+	t.Parallel()
+
+	pageSize := int64(header.PageSize)
+	numPages := uint32(12)
+	memfd, expected := newTestMemfd(t, pageSize*int64(numPages))
+
+	dirty := roaring.New()
+	dirty.AddRange(0, uint64(numPages))
+
+	cachePath := t.TempDir() + "/cache"
+	cache, err := NewCacheFromMemfdAsync(t.Context(), pageSize, cachePath, memfd, dirty)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = cache.Close() })
+
+	require.NoError(t, cache.Wait(t.Context()))
+
+	_, err = cache.Slice(0, pageSize)
+	require.NoError(t, err)
+
+	fromFile, err := os.ReadFile(cachePath)
+	require.NoError(t, err)
+	require.Equal(t, expected, fromFile)
+}
diff --git a/packages/orchestrator/pkg/sandbox/build/local_diff.go b/packages/orchestrator/pkg/sandbox/build/local_diff.go
@@ -80,14 +80,14 @@ func (f *LocalDiffFile) CloseToDiff(
 
 type localDiff struct {
 	cacheKey DiffStoreKey
-	cache    *block.Cache
+	cache    block.DiffSource
 }
 
 var _ Diff = (*localDiff)(nil)
 
 func NewLocalDiffFromCache(
 	cacheKey DiffStoreKey,
-	cache *block.Cache,
+	cache block.DiffSource,
 ) (Diff, error) {
 	return &localDiff{
 		cache:    cache,
@@ -110,6 +110,14 @@ func newLocalDiff(
 }
 
 func (b *localDiff) CachePath() (string, error) {
+	if w, ok := b.cache.(interface {
+		Wait(ctx context.Context) error
+	}); ok {
+		if err := w.Wait(context.Background()); err != nil {
+			return "", fmt.Errorf("memfd copy: %w", err)
+		}
+	}
+
 	return b.cache.Path(), nil
 }
 
diff --git a/packages/orchestrator/pkg/sandbox/fc/memory.go b/packages/orchestrator/pkg/sandbox/fc/memory.go
@@ -28,7 +28,7 @@ func (p *Process) exportMemoryFromFc(
 	include *roaring.Bitmap,
 	cachePath string,
 	blockSize int64,
-) (*block.Cache, error) {
+) (block.DiffSource, error) {
 	m, err := p.client.memoryMapping(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get memory mappings: %w", err)
@@ -64,10 +64,14 @@ func (p *Process) ExportMemory(
 	cachePath string,
 	blockSize int64,
 	memfd *block.Memfd,
-) (*block.Cache, error) {
+	bgCopy bool,
+) (block.DiffSource, error) {
 	if memfd == nil {
 		return p.exportMemoryFromFc(ctx, include, cachePath, blockSize)
 	}
+	if bgCopy {
+		return block.NewCacheFromMemfdAsync(ctx, blockSize, cachePath, memfd, include)
+	}
 
 	return block.NewCacheFromMemfd(ctx, blockSize, cachePath, memfd, include)
 }
diff --git a/packages/orchestrator/pkg/sandbox/sandbox.go b/packages/orchestrator/pkg/sandbox/sandbox.go
@@ -1138,6 +1138,7 @@ func (s *Sandbox) Pause(
 		s.config.DefaultCacheDir,
 		s.process,
 		s.memory.Memfd(ctx),
+		s.featureFlags.BoolFlag(ctx, featureflags.MemfdBackgroundCopyFlag, sandboxLDContext(s.Runtime, s.Config)),
 	)
 	if err != nil {
 		return nil, fmt.Errorf("error while post processing: %w", err)
@@ -1206,13 +1207,14 @@ func pauseProcessMemory(
 	cacheDir string,
 	fc *fc.Process,
 	memfd *block.Memfd,
+	bgCopy bool,
 ) (d build.Diff, h *header.Header, e error) {
 	ctx, span := tracer.Start(ctx, "process-memory")
 	defer span.End()
 
 	// ExportMemory owns memfd and closes it on all paths.
 	memfileDiffPath := build.GenerateDiffCachePath(cacheDir, buildID.String(), build.Memfile)
-	cache, err := fc.ExportMemory(ctx, diffMetadata.Dirty, memfileDiffPath, diffMetadata.BlockSize, memfd)
+	cache, err := fc.ExportMemory(ctx, diffMetadata.Dirty, memfileDiffPath, diffMetadata.BlockSize, memfd, bgCopy)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to export memory: %w", err)
 	}
diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go
