Skip to content

Use memfd wake path for UFFD installs - #2742

Merged
ValentaTomas merged 93 commits into
mainfrom
draft/uffd-minor-continue-zero-copy
May 25, 2026
Merged

ValentaTomas merged 93 commits into
mainfrom
draft/uffd-minor-continue-zero-copy

Conversation

@ValentaTomas

@ValentaTomas ValentaTomas commented May 20, 2026 •

Copy link
Copy Markdown
Member

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 of use-memfd). resume-build -resume-bench compares default / memfd-copy / memfd-wake on the same template.

Stacked on #2590.

bchalios and others added 30 commits May 15, 2026 14:08
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.
@ValentaTomas
ValentaTomas force-pushed the draft/uffd-minor-continue-zero-copy branch 2 times, most recently from 2672565 to b34def1 Compare May 23, 2026 00:47

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread packages/orchestrator/pkg/sandbox/uffd/userfaultfd/missing_wake.go Outdated
Classify zero source pages before comparing against the base so zero bytes always map through Empty instead of uploading dirty zero pages.
@ValentaTomas
ValentaTomas force-pushed the draft/uffd-minor-continue-zero-copy branch 2 times, most recently from fd8dcb0 to b6eb4a2 Compare May 23, 2026 01:14

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread packages/orchestrator/pkg/sandbox/uffd/userfaultfd/missing_wake.go
@ValentaTomas
ValentaTomas force-pushed the draft/uffd-minor-continue-zero-copy branch from b6eb4a2 to 3bedd2f Compare May 23, 2026 05:09

@cursor cursor Bot left a comment •

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Fix All in Cursor

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.

Create PR

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.

Comment thread packages/orchestrator/pkg/sandbox/uffd/userfaultfd/missing_wake.go Outdated

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread packages/orchestrator/pkg/sandbox/uffd/userfaultfd/missing_wake.go Outdated
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.
@ValentaTomas
ValentaTomas force-pushed the draft/uffd-minor-continue-zero-copy branch from 3bedd2f to a9fc454 Compare May 23, 2026 05:44
@ValentaTomas

Copy link
Copy Markdown
Member Author

Dropped the EEXIST handling on the UFFDIO_WAKE branch in missing_wake.go — UFFDIO_WAKE doesn't actually return EEXIST (that's a UFFDIO_COPY/UFFDIO_CONTINUE errno; WAKE only signals threads, doesn't check PTE state). The dead branch was an artifact of copy-pasting UFFDIO_COPY error handling and was triggering false-positive Bugbot/Codex flags about the stale-write race. With it gone, a real concurrent install can't be silently treated as success.

Base automatically changed from feat/compact-diff to main May 24, 2026 07:07
Resolve conflicts after main updated the integration dedup mode and related orchestrator paths.
Comment thread packages/orchestrator/pkg/sandbox/uffd/uffd.go
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.
@ValentaTomas
ValentaTomas requested a review from jakubno May 25, 2026 06:08
@ValentaTomas
ValentaTomas merged commit 4c96e4d into main May 25, 2026
51 checks passed
@ValentaTomas
ValentaTomas deleted the draft/uffd-minor-continue-zero-copy branch May 25, 2026 06:14
ValentaTomas added a commit that referenced this pull request May 25, 2026
jakubno pushed a commit that referenced this pull request May 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants