Use memfd wake path for UFFD installs - #2742
Conversation
Introduce a Cacher interface which abstracts the memory cache implementation, as seen by the diff/upload layer. Currently, it's only the Cache type that implements that. This is preparation for introducing a second cache type which is backed by the memfd used to map guest memory. Signed-off-by: Babis Chalios <babis.chalios@e2b.dev>
We are changing Firecracker to, optionally, back the guest memory using a memfd object. When enabled, Firecracker passes over the memfd file descriptor over the UFFD UDS, alongside the UFFD file descriptor, using SCM_RIGHTS. Change the UFFD serve logic to also parse the memfd file descriptor. When present, wrap the descriptor in a Memfd object. The object itself provides an interface that lets users access the guest memory from the memfd. UFFD logic exposes the Memfd object over a newly added method of the MemoryBackend interface, called Memfd(). The noop memory backend always returns nil for now, as Firecracker might only use memfd when resuming from a snapshot. Signed-off-by: Babis Chalios <babis.chalios@e2b.dev>
Change the ExportMemory() logic to export the memory via a MemfdCache when Firecracker has sent us a memfd file descriptor. Signed-off-by: Babis Chalios <babis.chalios@e2b.dev>
Add unit tests for memfd and MemfdCache functionality. Signed-off-by: Babis Chalios <babis.chalios@e2b.dev>
Add a feature flag that controls whether the orchestrator will instruct Firecracker to use memfd for backing the guest memory. Signed-off-by: Babis Chalios <babis.chalios@e2b.dev>
Trim verbose comments and drop trivial tests so the PR is easier to review. The behavioral changes are: - copyFromMemfd uses a fixed 2 MiB chunk (memfdCopyChunkSize) matching the source hugepage size, decoupled from cache.blockSize (which remains the dirty-tracking unit). - NewCacheFromMemfd no longer logs-and-swallows the memfd close error; it returns it. The size==0 fast path drops; the loop is a no-op when ranges sum to 0. - UseMemFdFlag comment matches the mmap-based implementation.
NewCacheFromMemfd already closes the memfd on every error path; calling memfd.Close() again here is a (harmless) double-close that muddies ownership semantics.
Trim the diff further: - MemfdCache embeds *Cache, dropping six delegate methods and the custom Close (Cache.Close handles the file; the memfd is already closed before the wrapper is returned). - copyFromMemfd performs one copy per range instead of a chunked inner loop; the memfdCopyChunkSize constant goes away. Cancellation checks move to range boundaries. - Memfd.Slice uses a simple nil-check instead of sync.Once for the lazy mmap; no concurrent callers in either the sync or the upcoming async path. - Drop TestMemfdCache_DirtyBitmap (covered by MultipleRanges + the BitsetRanges adapter is exercised elsewhere) and the chunk-boundary test (no chunked loop anymore). Consolidate SliceOutOfBounds.
NewFromFd now mmaps the fd eagerly via fstat and returns (*Memfd, error). The size field on Memfd goes away (use len(mmap)), the lazy-init nil-check in Slice goes away, and the explicit size computation in uffd.go is dropped — the kernel already knows the size. Also standardize on golang.org/x/sys/unix (drop the mixed syscall import).
Memfd has a single owner across all paths: NewCacheFromMemfd consumes it during construction, and the UFFD handshake transfers ownership via atomic Swap. There is no path that calls Close twice on the same Memfd, so the m.mmap=nil / m.fd=-1 / nil-check sentinels are dead weight. Drop them and document the single-use contract.
Instead of materializing a []Range in the caller just to iterate it once inside NewCacheFromMemfd, take the dirty *roaring.Bitmap and the block size directly. Total cache size comes from cardinality * blockSize, and copyFromMemfd iterates BitsetRanges in place. Drops the now-thin exportMemoryFromMemfd helper in fc/memory.go.
The gRPC handler already injects sandbox/team/template contexts via ctx, so team and template targeting for UseMemFdFlag already worked. Add a sandbox-type attribute (sandbox vs build) and pass the explicit sandboxLDContext to BoolFlag so flags can roll out to production sandboxes separately from template-builds.
Cacher was a vague -er name for an interface with seven methods. The type exists purely so the diff/upload layer can accept either *Cache or *MemfdCache; DiffSource names that role.
- copyFromMemfd uses ctx.Err() instead of the select/default form. - pauseProcessMemory runs ExportMemory before ToDiffHeader so the memfd is owned by ExportMemory throughout; the conditional close on ToDiffHeader failure goes away. - fc.ExportMemory returns NewCacheFromMemfd directly; the "create MemfdCache" wrap is redundant with the inner error context.
PR 2522 introduced MemfdCache (wrapper around *Cache) and the DiffSource interface purely so the async-copy follow-up could attach extra state without churning callers. PR 2522 itself never uses the indirection — NewCacheFromMemfd returns *Cache, fc.ExportMemory returns *block.Cache, localDiff takes *block.Cache. Drop the scaffolding from this PR; the async PR introduces its own wrapper when it needs the override behavior. Also: - Inline the copyFromMemfd loop into NewCacheFromMemfd (single use). - Trim tests to the two that earn their keep: non-adjacent blocks and the non-zero range-start regression. - Tighten the use_memfd/firecracker-fd comments. - Loop over fds for cleanup instead of two manual closes.
FC < 1.14 rejects the use_memfd field on snapshot load (deny_unknown_fields on MemoryBackend), so combining FCSupportsMemfd(version) with the flag avoids hard-failing resumes when the flag is flipped on across a heterogeneous fleet.
Before introducing deduplicated caches, all files including sandbox memory were using a block size of 2M, i.e. data were stored in chunks of 2M. With deduplicated caches, memory files produced by PAUSE and SNAPSHOT operations might use 4K block sizes. This means that using a simple `mmap` over a slice of the snapshot data won't work any more. UFFD gives us 2M aligned page faults to handle, so we need to do a bit more work to find the actual data. Change the logic to always fall through to ReadAt for getting the data we need to serve a page fault. Signed-off-by: Babis Chalios <babis.chalios@e2b.dev>
Within the sandbox guest sees 4K pages. On the host side we back sandbox memory with 2MB huge pages. Firecracker tracks dirty pages (with the help of UFFD WP async mechanism) at the huge page granularity. This means that even if the guest touches a single 4K page will mark the entire corresponding 2M memory range as dirty and include it in the memory snapshot diff we're creating during PAUSE and SNASPHOT operations. This increases the memory snapshots we take, which reflects on the amount of data we store to GCS. Now, we always start sandboxes from a snapshot, so we have knowledge of the intial memory contents. This allows us to iterate the previous state of the memory (when we started/resumed the sandbox) and see exactly which 4K ranges were touched by the guest and only store this in the produced Diff file that we keep in the Snapshot. Introduce dedup logic for the Cache and MemfdCache objects which does exactly that. In the case of Cache, the current implementation requires us to create the Cache file at 2MiB granularity and then deduplicate it creating a new Cache that includes only the actually changed 4K ranges (the blockSize of that second Cache is 4K). In the case of MemfdCache, we compare the memfd contents against the base snapshot and build the cache in one pass. Signed-off-by: Babis Chalios <babis.chalios@e2b.dev> wip: small fix
Wire the deduplication logic using a feature flag called `memfile-diff-dedup`. When the flag is true, pauseProcessMemory will receive a non nil pointer to the original memory file and use it to deduplicate the data of the produced diff. Signed-off-by: Babis Chalios <babis.chalios@e2b.dev> wip: small fix
Drop the unresolved <<<<<<< / ======= / >>>>>>> markers and the duplicate fullDirty helper that slipped into memfd_test.go. The HEAD side already had fullDirty defined at the top of the file and used it consistently, so we keep that shape. Unblocks orchestrator lint, unit tests and arm64 cross-compile.
Drop prose comments that explain context already obvious from the code, keeping only those documenting non-obvious intent or required invariants.
Set Metadata.BlockSize from DiffMetadata.BlockSize in ToDiffHeader so each generation's header advertises its actual mapping granularity. Reverts getMapping to align by Metadata.BlockSize and ValidateMappings to validate by Metadata.BlockSize — no longer need PageSize special-cases since the deduped header carries BlockSize=PageSize itself.
- Add dedupPages core; (*Cache).Dedup and NewCacheFromMemfdDeduped become thin wrappers selecting the page source. - Fix build.File.ReadAt to zero-fill uuid.Nil ranges; this fixes the stale-baseBuf read in the dedup path. Drop the now-redundant clear(b) in faultPage. - Drop duplicated TestNewCacheFromMemfdDeduped_* (covered by TestCacheDedup_*). - Inline UpsampleBitmap (single caller) and drop bitmap.go/bitmap_test.go.
The override broke dirty-tracking on resume from a deduped parent: u.memfile.BlockSize() returns Metadata.BlockSize, which feeds f.DirtyMemory, but FC tracks at its own page size. Keep Metadata.BlockSize == FC page size; ValidateMappings/getMapping align at PageSize so mixed-granularity mappings still pass.
Skip empty iovecs before pwritev so zero-length batches complete cleanly.
Remove the unused memfd test override now that integration dedup coverage runs through the default memory export path.
2672565 to
b34def1
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b34def151e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Classify zero source pages before comparing against the base so zero bytes always map through Empty instead of uploading dirty zero pages.
fd8dcb0 to
b6eb4a2
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b6eb4a220c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
b6eb4a2 to
3bedd2f
Compare
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Memfd write clobbers on EEXIST
- Changed faultPageViaMemfdWake to read into temporary buffer and only copy to memfd after UFFDIO_WAKE succeeds, preventing overwrite of guest-dirty bytes when concurrent path already mapped the page.
Or push these changes by commenting:
@cursor push 101bb9915e
Preview (101bb9915e)
diff --git a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/missing_wake.go b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/missing_wake.go
--- a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/missing_wake.go
+++ b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/missing_wake.go
@@ -68,7 +68,7 @@
return faultDiscarded, errors.Join(err, safeInvoke(onFailure))
}
- page := dst[offset : offset+pageSize]
+ tmpBuf := make([]byte, pageSize)
var dataErr error
var attempt int
@@ -76,7 +76,7 @@
retryLoop:
for attempt = range sliceMaxRetries + 1 {
var n int
- n, dataErr = source.ReadAt(ctx, page, offset)
+ n, dataErr = source.ReadAt(ctx, tmpBuf, offset)
if dataErr == nil && int64(n) != pageSize {
dataErr = fmt.Errorf("short read at %d: got %d, want %d", offset, n, pageSize)
}
@@ -115,9 +115,6 @@
}
if err := u.fd.wake(addr, u.pageSize); err != nil {
- // EEXIST: page already installed by a concurrent path. The memfd
- // write is idempotent (same source for on-demand and prefault),
- // matching UFFDIO_COPY's first-writer-wins semantics.
if errors.Is(err, unix.EEXIST) {
span.SetAttributes(attribute.Bool("uffd.already_mapped", true))
@@ -139,5 +136,7 @@
return faultDiscarded, fmt.Errorf("UFFDIO_WAKE: %w", joined)
}
+ copy(dst[offset:offset+pageSize], tmpBuf)
+
return faultInstalled, nil
}You can send follow-ups to the cloud agent here.
Reviewed by Cursor Bugbot for commit 3bedd2f. Configure here.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3bedd2fb28
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Resolve MISSING faults by writing into the FC-shared memfd and waking the faulter, instead of UFFDIO_COPY. For reads, arm WP before the populate so the kernel retry installs a write-protected PTE; for zero pages, a one-byte touch is enough to allocate (and zero-fill) the page in both the 4 KiB and hugetlb cases.
3bedd2f to
a9fc454
Compare
|
Dropped the EEXIST handling on the UFFDIO_WAKE branch in |
Resolve conflicts after main updated the integration dedup mode and related orchestrator paths.
Explain why SetMemfd(nil) precedes ExportPageStates: clearing the handler's memfd first stops new workers from picking the path up; the subsequent settleRequests.Lock acquire drains any worker that already loaded the pointer. Reversed order would race a fresh worker against ownership transfer to the caller.
This reverts commit 4c96e4d.


Install UFFD MISSING faults by writing source bytes into the FC-shared memfd and calling UFFDIO_WAKE, instead of UFFDIO_COPY. Read faults arm WP before populate so the kernel retry installs a write-protected PTE.
Behind
use-memfd-wake(sub-flag ofuse-memfd).resume-build -resume-benchcomparesdefault/memfd-copy/memfd-wakeon the same template.Stacked on #2590.