From 8a88dbf9cde6bf4feabe0f98ff3df7793adf1ef3 Mon Sep 17 00:00:00 2001
From: ValentaTomas <valenta.and.thomas@gmail.com>
Date: Mon, 27 Apr 2026 16:23:46 -0700
Subject: [PATCH] fix(uffd): re-read page state inside worker goroutine under
 settleRequests.RLock
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The Serve() loop previously read pageTracker state and captured
`source = u.src` in the parent loop, then dispatched a worker
goroutine. A REMOVE event for the same page that arrived between the
state read and the worker actually acquiring settleRequests.RLock()
would silently leave the worker with a stale `source = u.src`
snapshot. The worker would then UFFDIO_COPY src bytes into a page the
kernel had just MADV_DONTNEED'd, leaving pageTracker == removed and
the kernel page mapped with stale src data — and observably
deadlocking parent madvise() in the orchestrator unit-test suite.

Move the state lookup and source capture inside the goroutine, after
RLock(). The read+act+commit sequence is now atomic with respect to
the REMOVE batch (which takes settleRequests.Lock()).

# Conflicts:
#	packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
---
 .../sandbox/uffd/userfaultfd/userfaultfd.go   | 51 +++++++++++--------
 1 file changed, 30 insertions(+), 21 deletions(-)

diff --git a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
index 1593e9ef3a..da2200163b 100644
--- a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
+++ b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
@@ -334,24 +334,6 @@ func (u *Userfaultfd) Serve(
 				return fmt.Errorf("failed to map: %w", err)
 			}
 
-			var source block.Slicer
-
-			switch state := u.pageTracker.get(addr); state {
-			case faulted:
-				// Skip faulting the page. This has already been faulted, either during pre-faulting
-				// or because we handled another page fault on the same address in the current
-				// iteration. It can only transition out of `faulted` via a UFFD_EVENT_REMOVE, which
-				// will mark the page as `removed`.
-				// For this to work correctly, the used pages cannot be swappable.
-				continue
-			case removed:
-				// Fault the page as empty.
-			case missing:
-				source = u.src
-			default:
-				return fmt.Errorf("unexpected pageState: %#v", state)
-			}
-
 			u.wg.Go(func() error {
 				// Test-only barrier: park the worker BEFORE it takes
 				// RLock. While parked, the parent loop is free to take
@@ -363,12 +345,39 @@ func (u *Userfaultfd) Serve(
 					hook(addr)
 				}
 
-				// The RLock must be called inside the goroutine to ensure RUnlock runs via defer,
-				// even if the errgroup is cancelled or the goroutine returns early.
-				// This check protects us against race condition between marking the request for prefetching and accessing the prefetchTracker.
+				// The RLock must be acquired inside the goroutine — and it must be acquired
+				// BEFORE we read the pageTracker / u.src state — so that the read+act+commit
+				// sequence (state lookup → faultPage → setState) is atomic with respect to
+				// any concurrent REMOVE batch (which takes settleRequests.Lock()). If the
+				// state read happened in the parent loop, a REMOVE could land between the
+				// read and the goroutine acquiring the RLock, and the goroutine would still
+				// commit `faulted` afterwards, overwriting `removed`.
+				//
+				// This also protects the read of u.src: in the future src could be swapped
+				// out under settleRequests; reading it under the RLock keeps that safe.
 				u.settleRequests.RLock()
 				defer u.settleRequests.RUnlock()
 
+				var source block.Slicer
+
+				switch state := u.pageTracker.get(addr); state {
+				case faulted:
+					// Skip faulting the page. This has already been faulted, either during pre-faulting
+					// or because we handled another page fault on the same address in the current
+					// iteration. It can only transition out of `faulted` via a UFFD_EVENT_REMOVE, which
+					// will mark the page as `removed`.
+					// For this to work correctly, the used pages cannot be swappable.
+					return nil
+				case removed:
+					// Fault the page as empty (no source). The page was MADV_DONTNEED'd; the
+					// kernel still expects an UFFDIO_COPY/ZEROPAGE ack for the original
+					// MISSING fault, otherwise the faulting thread stays blocked.
+				case missing:
+					source = u.src
+				default:
+					return fmt.Errorf("unexpected pageState: %#v", state)
+				}
+
 				var accessType block.AccessType
 
 				if pf.flags&UFFD_PAGEFAULT_FLAG_WRITE == 0 {