test(uffd): add cross-process scaffolding for gated and async ops (#2475)

* test(uffd): add cross-process scaffolding for gated and async ops

Pre-work for the upcoming REMOVE event handling PR. Lands as a runtime
no-op on main: no production code changes, no new UFFD features enabled.

Test-helper additions (all *_test.go):

- testConfig.gated and per-op `async bool`
- operationModeRemove / ServePause / ServeResume / Sleep
- testHandler.servePause / serveResume hooks driven by control pipes
  (fd 7/8) into the helper subprocess
- helper subprocess pause/resume goroutine that tears down and restarts
  the UFFD serve loop on demand
- executeAll() async path that fans out goroutines and joins on test
  context
- expectRemoved + checkDirtiness handling for read/write/remove ordering
- executeRemove via unix.MADV_DONTNEED
- unregister() helper around UFFDIO_UNREGISTER for cleanup
- pageStateEntries() now drains settleRequests then takes the
  pageTracker RLock directly so the snapshot stays correct for future
  writers (e.g. REMOVE) that don't go through settleRequests

operationModeRemove and the gated/async paths are exercised only by the
follow-up PR; here they are dead code at runtime, kept to minimise the
diff of that PR.

* test(uffd): restore early uffd close cleanup, keep unregister late

* test(uffd): make gated cleanup idempotent to avoid pause-then-exit hang

Codex review caught a deadlock: in GO_GATED mode, handling 'P' replaced
the stop closure with one that drained `exitUffd`, but the deferred
final cleanup still pointed at the original closure. If the helper exited
between 'P' and 'R' (e.g. parent test failure, ctx cancel, SIGUSR1), the
defer ran the original cleanup again and blocked forever on `<-exitUffd`
— the channel had already been drained and the original Serve goroutine
had already exited. The hung helper would in turn hang the parent's
cmd.Wait().

Replace the two-variable (cleanup / stopServe) layout with a single
mutex-guarded `stopFn` that is reset to a no-op after firing, and have
both the gated 'P' handler and the deferred final cleanup go through
the same `stopServe` wrapper. 'R' now just installs a fresh `stopFn`
that drains the new goroutine.

* test(uffd): drop REMOVE-specific bits, keep only gated/async scaffolding

Pull out everything that's not strictly needed for "gated pause/resume +
async ops" so PR T is the smallest possible scaffolding diff and the
follow-up REMOVE PR carries those pieces alongside the production code
that actually needs them:

- helpers_test.go: drop operationModeRemove, executeRemove, expectRemoved
  and the REMOVE-aware checkDirtiness branch — none of the PR T tests
  call them, and they only become meaningful once UFFD_FEATURE_EVENT_REMOVE
  is set in the follow-up.
- fd_helpers_test.go: revert; drop unregister() helper. Without
  UFFD_FEATURE_EVENT_REMOVE, MADV_DONTNEED never queues UFFD events, so
  munmap doesn't block on un-acked events and there's nothing to
  unregister around.
- cross_process_helpers_test.go:
  - drop unregister(uffdFd, ...) call in late cleanup (same reason as
    above).
  - revert pageStateEntries() to the simple settleRequests.Lock pattern;
    the rework was forward-looking for a REMOVE handler that doesn't go
    through settleRequests and is dead code on main.
  - keep cleanup → stopFn/stopServe restructure (mandatory for the
    gated tear-down/respawn path) but restore the original
    fdExit.SignalExit error handling that the previous version had
    silently dropped.
  - revert incidental cosmetic reformats (exitUffd placement, defer
    multi-line) so the diff is purely additive in the unchanged paths.

* test(uffd): address review comments on cross-process scaffolding

- Guard executeOperation against nil servePause/serveResume hooks when a
  non-gated test misuses operationModeServePause/Resume — return a clear
  error instead of panicking on a nil function call (cursor + codex).
- Make pageStateEntries() also hold pageTracker.mu.RLock() so the
  snapshot stays correct for any future writer that doesn't go through
  settleRequests (claude — matches the PR description).
- Make gated 'P' check the gateSyncFile.Write error and propagate via
  cancel() so a write failure doesn't leave the parent's servePause
  closure blocked indefinitely on Read (claude).
- Make the gated command goroutine reject 'P'-when-paused and
  'R'-when-running so a stray or duplicate resume can't leak an
  untracked Serve goroutine and break later pauses (codex).

* test(uffd): serialize executeRead under shared RWMutex with executeWrite

Previously executeWrite took h.mutex.Lock() but executeRead held no lock,
so an async write + async read on the same page would race on memoryArea
under go test -race even though the bytes happen to match (data.Slice is
deterministic per offset). Switch h.mutex to sync.RWMutex; reads take
RLock so concurrent reads still trigger UFFD faults in parallel, writes
take Lock so they exclude both other writers and any in-flight reads.

No active test in this PR mixes async read/write yet; this is preemptive
against the upcoming REMOVE PR's matrix tests.

Author: ValentaTomas

packages/orchestrator/pkg/sandbox/uffd/userfaultfd/cross_process_helpers_test.go

@@ -19,6 +19,7 @@ import (
 	"slices"
 	"strconv"
 	"strings"
+	"sync"
 	"syscall"
 	"testing"
 
@@ -110,6 +111,9 @@ func configureCrossProcessTest(t *testing.T, tt testConfig) (*testHandler, error
 	if tt.alwaysWP {
 		cmd.Env = append(cmd.Env, "GO_ALWAYS_WP=1")
 	}
+	if tt.gated {
+		cmd.Env = append(cmd.Env, "GO_GATED=1")
+	}
 
 	dup, err := syscall.Dup(int(uffdFd))
 	require.NoError(t, err)
@@ -153,12 +157,34 @@ func configureCrossProcessTest(t *testing.T, tt testConfig) (*testHandler, error
 		readySignal <- struct{}{}
 	}()
 
-	cmd.ExtraFiles = []*os.File{
+	extraFiles := []*os.File{
 		uffdFile,
 		contentReader,
 		offsetsWriter,
 		readyWriter,
 	}
+
+	var gateCmdWriter *os.File
+	var gateSyncReader *os.File
+	if tt.gated {
+		var gateCmdReader *os.File
+		gateCmdReader, gateCmdWriter, err = os.Pipe()
+		require.NoError(t, err)
+
+		var gateSyncWriter *os.File
+		gateSyncReader, gateSyncWriter, err = os.Pipe()
+		require.NoError(t, err)
+
+		t.Cleanup(func() {
+			gateCmdWriter.Close()
+			gateSyncReader.Close()
+		})
+
+		extraFiles = append(extraFiles, gateCmdReader)  // fd 7
+		extraFiles = append(extraFiles, gateSyncWriter) // fd 8
+	}
+
+	cmd.ExtraFiles = extraFiles
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 
@@ -169,6 +195,10 @@ func configureCrossProcessTest(t *testing.T, tt testConfig) (*testHandler, error
 	offsetsWriter.Close()
 	readyWriter.Close()
 	uffdFile.Close()
+	if tt.gated {
+		extraFiles[4].Close() // gateCmdReader
+		extraFiles[5].Close() // gateSyncWriter
+	}
 
 	t.Cleanup(func() {
 		signalErr := cmd.Process.Signal(syscall.SIGUSR1)
@@ -238,12 +268,31 @@ func configureCrossProcessTest(t *testing.T, tt testConfig) (*testHandler, error
 	case <-readySignal:
 	}
 
-	return &testHandler{
+	h := &testHandler{
 		memoryArea:     &memoryArea,
 		pagesize:       tt.pagesize,
 		data:           data,
 		pageStatesOnce: pageStatesOnce,
-	}, nil
+	}
+
+	if tt.gated {
+		h.servePause = func() error {
+			if _, err := gateCmdWriter.Write([]byte{'P'}); err != nil {
+				return err
+			}
+			var buf [1]byte
+			_, err := gateSyncReader.Read(buf[:])
+
+			return err
+		}
+		h.serveResume = func() error {
+			_, err := gateCmdWriter.Write([]byte{'R'})
+
+			return err
+		}
+	}
+
+	return h, nil
 }
 
 // Secondary process, orchestrator in our case
@@ -378,22 +427,111 @@ func crossProcessServe() error {
 		}
 	}()
 
-	cleanup := func() {
-		err := fdExit.SignalExit()
-		if err != nil {
-			msg := fmt.Errorf("error signaling exit: %w", err)
+	// stopFn drains whichever Serve goroutine is currently running. The
+	// running flag plus stopMu makes both pause-then-exit (no resume in
+	// between) and pause-resume-pause-exit safe, and rejects nonsensical
+	// command sequences from the gated channel: 'P' when already paused
+	// is a no-op, 'R' when already running is a no-op so a stray or
+	// duplicate resume can't leak an untracked Serve goroutine and break
+	// later pauses.
+	var (
+		stopMu  sync.Mutex
+		running = true
+		stopFn  = func() {
+			err := fdExit.SignalExit()
+			if err != nil {
+				msg := fmt.Errorf("error signaling exit: %w", err)
 
-			fmt.Fprint(os.Stderr, msg.Error())
+				fmt.Fprint(os.Stderr, msg.Error())
 
-			cancel(msg)
+				cancel(msg)
+
+				return
+			}
+
+			<-exitUffd
+		}
+	)
+
+	stopServe := func() {
+		stopMu.Lock()
+		if !running {
+			stopMu.Unlock()
 
 			return
 		}
+		fn := stopFn
+		stopFn = func() {}
+		running = false
+		stopMu.Unlock()
 
-		<-exitUffd
+		fn()
 	}
 
-	defer cleanup()
+	defer stopServe()
+
+	if os.Getenv("GO_GATED") == "1" {
+		gateCmdFile := os.NewFile(uintptr(7), "gate-cmd")
+		defer gateCmdFile.Close()
+
+		gateSyncFile := os.NewFile(uintptr(8), "gate-sync")
+		defer gateSyncFile.Close()
+
+		startServe := func() {
+			stopMu.Lock()
+			if running {
+				stopMu.Unlock()
+
+				return
+			}
+			stopMu.Unlock()
+
+			newExit, fdErr := fdexit.New()
+			if fdErr != nil {
+				cancel(fmt.Errorf("error creating fd exit: %w", fdErr))
+
+				return
+			}
+
+			done := make(chan struct{})
+			go func() {
+				defer close(done)
+				if err := uffd.Serve(ctx, newExit); err != nil {
+					cancel(fmt.Errorf("error serving: %w", err))
+				}
+			}()
+
+			stopMu.Lock()
+			stopFn = func() {
+				newExit.SignalExit()
+				<-done
+				newExit.Close()
+			}
+			running = true
+			stopMu.Unlock()
+		}
+
+		go func() {
+			var buf [1]byte
+			for {
+				if _, err := gateCmdFile.Read(buf[:]); err != nil {
+					return
+				}
+
+				switch buf[0] {
+				case 'P':
+					stopServe()
+					if _, err := gateSyncFile.Write([]byte{1}); err != nil {
+						cancel(fmt.Errorf("writing gate sync: %w", err))
+
+						return
+					}
+				case 'R':
+					startServe()
+				}
+			}
+		}()
+	}
 
 	exitSignal := make(chan os.Signal, 1)
 	signal.Notify(exitSignal, syscall.SIGUSR1)
@@ -423,12 +561,17 @@ type pageStateEntry struct {
 }
 
 // pageStateEntries returns a snapshot of every tracked page and its state.
-// It holds the settleRequests write lock so no in-flight faultPage worker
-// can mutate the pageTracker while we iterate.
+// It first drains in-flight faultPage workers via settleRequests.Lock(), then
+// holds the pageTracker RLock while iterating so any future writer that
+// doesn't go through settleRequests (e.g. the upcoming REMOVE handler)
+// still can't mutate the map under us.
 func (u *Userfaultfd) pageStateEntries() ([]pageStateEntry, error) {
 	u.settleRequests.Lock()
 	defer u.settleRequests.Unlock()
 
+	u.pageTracker.mu.RLock()
+	defer u.pageTracker.mu.RUnlock()
+
 	entries := make([]pageStateEntry, 0, len(u.pageTracker.m))
 	for addr, state := range u.pageTracker.m {
 		offset, err := u.ma.GetOffset(addr)

packages/orchestrator/pkg/sandbox/uffd/userfaultfd/helpers_test.go

@@ -3,10 +3,12 @@ package userfaultfd
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"slices"
 	"sync"
 	"testing"
+	"time"
 	"unsafe"
 
 	"github.com/RoaringBitmap/roaring/v2"
@@ -26,19 +28,28 @@ type testConfig struct {
 	operations []operation
 	// alwaysWP makes the handler copy with UFFDIO_COPY_MODE_WP for all faults.
 	alwaysWP bool
+	// gated enables pause/resume control over the handler's serve loop.
+	gated bool
 }
 
 type operationMode uint32
 
 const (
 	operationModeRead operationMode = 1 << iota
 	operationModeWrite
+	operationModeServePause
+	operationModeServeResume
+	// operationModeSleep pauses for a short duration to let async goroutines
+	// enter their blocking syscalls before proceeding.
+	operationModeSleep
 )
 
 type operation struct {
 	// Offset in bytes. Must be smaller than the (numberOfPages-1) * pagesize as it reads a page and it must be aligned to the pagesize from the testConfig.
 	offset int64
 	mode   operationMode
+	// async runs the operation in a background goroutine.
+	async bool
 }
 
 // handlerPageStates is a snapshot of the pageTracker grouped by state. It
@@ -71,16 +82,43 @@ type testHandler struct {
 	// pageStatesOnce returns a per-state snapshot of the handler's pageTracker.
 	// It can only be called once.
 	pageStatesOnce func() (handlerPageStates, error)
-	mutex          sync.Mutex
+	// servePause and serveResume gate the UFFD event loop in the child process.
+	// Tests use them to deterministically batch a sequence of UFFD events
+	// before more faults are processed.
+	servePause  func() error
+	serveResume func() error
+	mutex       sync.RWMutex
 }
 
 func (h *testHandler) executeAll(t *testing.T, operations []operation) {
 	t.Helper()
 
+	var asyncErrors []chan error
+
 	for i, op := range operations {
+		if op.async {
+			errCh := make(chan error, 1)
+			asyncErrors = append(asyncErrors, errCh)
+
+			go func() {
+				errCh <- h.executeOperation(t.Context(), op)
+			}()
+
+			continue
+		}
+
 		err := h.executeOperation(t.Context(), op)
 		require.NoError(t, err, "step %d: %v at offset %d", i, op.mode, op.offset)
 	}
+
+	for _, errCh := range asyncErrors {
+		select {
+		case err := <-errCh:
+			require.NoError(t, err, "async operation")
+		case <-t.Context().Done():
+			t.Fatal("timed out waiting for async operation")
+		}
+	}
 }
 
 type pageExpectation uint8
@@ -135,19 +173,43 @@ func (h *testHandler) executeOperation(ctx context.Context, op operation) error
 		return h.executeRead(ctx, op)
 	case operationModeWrite:
 		return h.executeWrite(ctx, op)
+	case operationModeServePause:
+		if h.servePause == nil {
+			return errors.New("operationModeServePause requires testConfig.gated = true")
+		}
+
+		return h.servePause()
+	case operationModeServeResume:
+		if h.serveResume == nil {
+			return errors.New("operationModeServeResume requires testConfig.gated = true")
+		}
+
+		return h.serveResume()
+	case operationModeSleep:
+		time.Sleep(50 * time.Millisecond)
+
+		return nil
 	default:
 		return fmt.Errorf("invalid operation mode: %d", op.mode)
 	}
 }
 
 func (h *testHandler) executeRead(ctx context.Context, op operation) error {
-	readBytes := (*h.memoryArea)[op.offset : op.offset+int64(h.pagesize)]
-
 	expectedBytes, err := h.data.Slice(ctx, op.offset, int64(h.pagesize))
 	if err != nil {
 		return err
 	}
 
+	// Hold the read side of the memoryArea mutex while we touch the page.
+	// Reads can run concurrently with each other (RLock), but a parallel
+	// async write to the same page (executeWrite, write lock) is excluded
+	// so go test -race stays clean when a test plan mixes async read and
+	// async write at overlapping offsets.
+	h.mutex.RLock()
+	defer h.mutex.RUnlock()
+
+	readBytes := (*h.memoryArea)[op.offset : op.offset+int64(h.pagesize)]
+
 	// The bytes.Equal is the first place in this flow that actually touches the uffd managed memory and triggers the pagefault, so any deadlocks will manifest here.
 	if !bytes.Equal(readBytes, expectedBytes) {
 		idx, want, got := testutils.FirstDifferentByte(readBytes, expectedBytes)
@@ -165,6 +227,7 @@ func (h *testHandler) executeWrite(ctx context.Context, op operation) error {
 	}
 
 	// An unprotected parallel write to map might result in an undefined behavior.
+	// Lock excludes both other writers and any concurrent executeRead.
 	h.mutex.Lock()
 	defer h.mutex.Unlock()