e2b-dev · ValentaTomas · May 2, 2026 · Apr 21, 2026 · Apr 21, 2026 · Apr 21, 2026
@@ -0,0 +1,155 @@
+package testharness
+
+import (
+	"context"
+	"fmt"
+	"sync"
+)
+
+// Point identifies WHICH worker hook a barrier should park on. Values
+// must match the parent package's faultPhase iota so the test hook can
+// pass them through with a numeric cast.
+type Point uint8
+
+const (
+	// BeforeRLock parks the worker BEFORE settleRequests.RLock(), so a
+	// parallel writer can take the write lock immediately.
+	BeforeRLock Point = iota
+	// BeforeFaultPage parks the worker AFTER settleRequests.RLock but
+	// BEFORE the UFFDIO_COPY syscall, so a parent operation must still
+	// return even though a worker holds RLock.
+	BeforeFaultPage
+)
+
+// Registry is the child-process side of the barrier mechanism. The
+// per-fault hook on *Userfaultfd consults it by (addr, point) to decide
+// whether to park.
+type Registry struct {
+	mu     sync.Mutex
+	next   uint64
+	tokens map[uint64]*slot
+	byKey  map[key]uint64
+}
+
+type key struct {
+	addr  uintptr
+	point Point
+}
+
+type slot struct {
+	addr        uintptr
+	point       Point
+	arrived     chan struct{}
+	release     chan struct{}
+	arrivedOnce sync.Once
+}
+
+func NewRegistry() *Registry {
+	return &Registry{
+		tokens: make(map[uint64]*slot),
+		byKey:  make(map[key]uint64),
+	}
+}
+
+// Install registers a barrier at (addr, point) and returns its token.
+func (r *Registry) Install(addr uintptr, point Point) uint64 {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	r.next++
+	token := r.next
+	s := &slot{
+		addr:    addr,
+		point:   point,
+		arrived: make(chan struct{}),
+		release: make(chan struct{}),
+	}
+	r.tokens[token] = s
+	r.byKey[key{addr, point}] = token
+
+	return token
+}
+
+func (r *Registry) lookupByAddr(addr uintptr, point Point) *slot {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	token, ok := r.byKey[key{addr, point}]
+	if !ok {
+		return nil
+	}
+
+	return r.tokens[token]
+}
+
+// WaitArrived blocks until the worker hook for token has reached the
+// barrier point, or until ctx is cancelled.
+func (r *Registry) WaitArrived(ctx context.Context, token uint64) error {
+	r.mu.Lock()
+	s, ok := r.tokens[token]
+	r.mu.Unlock()
+	if !ok {
+		return fmt.Errorf("unknown barrier token %d", token)
+	}
+
+	select {
+	case <-s.arrived:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+// Release frees the barrier identified by token, allowing any parked
+// worker to proceed. Releasing an unknown token is a no-op.
+func (r *Registry) Release(token uint64) {
+	r.mu.Lock()
+	s, ok := r.tokens[token]
+	delete(r.tokens, token)
+	if ok {
+		delete(r.byKey, key{s.addr, s.point})
+	}
+	r.mu.Unlock()
+
+	if !ok {
+		return
+	}
+
+	select {
+	case <-s.release:
+	default:
+		close(s.release)
+	}
+}
+
+// ReleaseAll releases every still-installed barrier so any parked
+// worker can finish before the child's serve goroutine is joined.
+func (r *Registry) ReleaseAll() {
+	r.mu.Lock()
+	tokens := make([]uint64, 0, len(r.tokens))
+	for t := range r.tokens {
+		tokens = append(tokens, t)
+	}
+	r.mu.Unlock()
+
+	for _, t := range tokens {
+		r.Release(t)
+	}
+}
+
+// Hook returns the per-fault hook tests install on *Userfaultfd. Faults
+// at (addr, point) pairs without an Install'd slot are no-ops.
+func (r *Registry) Hook() func(addr uintptr, point Point) {
+	return func(addr uintptr, point Point) {
+		s := r.lookupByAddr(addr, point)
+		if s == nil {
+			return
+		}
+
+		s.arrivedOnce.Do(func() {
+			close(s.arrived)
+		})
+
+		<-s.release
+	}
+}
@@ -0,0 +1,75 @@
+package testharness
+
+import (
+	"io"
+	"net/rpc"
+	"net/rpc/jsonrpc"
+)
+
+// Client is the typed parent-side wrapper around the JSON-RPC channel
+// to the child helper process. It hides method-name strings, the Empty
+// placeholder pointers, and the wire-vs-domain type translation.
+type Client struct {
+	rpc  *rpc.Client
+	conn io.Closer
+}
+
+// NewClient wraps an already-connected duplex stream (typically one
+// half of a socketpair handed off via cmd.ExtraFiles). Closing the
+// returned Client closes the underlying conn.
+func NewClient(conn io.ReadWriteCloser) *Client {
+	return &Client{
+		rpc:  jsonrpc.NewClient(conn),
+		conn: conn,
+	}
+}
+
+func (c *Client) Bootstrap(args BootstrapArgs) error {
+	return c.rpc.Call("Lifecycle.Bootstrap", &args, &BootstrapReply{})
+}
+
+func (c *Client) WaitReady() error {
+	return c.rpc.Call("Lifecycle.WaitReady", &Empty{}, &Empty{})
+}
+
+func (c *Client) Shutdown() error {
+	return c.rpc.Call("Lifecycle.Shutdown", &Empty{}, &Empty{})
+}
+
+func (c *Client) Pause() error {
+	return c.rpc.Call("Paging.Pause", &Empty{}, &Empty{})
+}
+
+func (c *Client) Resume() error {
+	return c.rpc.Call("Paging.Resume", &Empty{}, &Empty{})
+}
+
+func (c *Client) PageStates() ([]PageStateEntry, error) {
+	var reply PageStatesReply
+	if err := c.rpc.Call("Paging.States", &Empty{}, &reply); err != nil {
+		return nil, err
+	}
+
+	return reply.Entries, nil
+}
+
+func (c *Client) InstallBarrier(addr uintptr, point Point) (uint64, error) {
+	var reply FaultBarrierReply
+	if err := c.rpc.Call("Barriers.Install", &FaultBarrierArgs{Addr: uint64(addr), Point: uint8(point)}, &reply); err != nil {
+		return 0, err
+	}
+
+	return reply.Token, nil
+}
+
+func (c *Client) WaitFaultHeld(token uint64) error {
+	return c.rpc.Call("Barriers.WaitHeld", &TokenArgs{Token: token}, &Empty{})
+}
+
+func (c *Client) ReleaseFault(token uint64) error {
+	return c.rpc.Call("Barriers.Release", &TokenArgs{Token: token}, &Empty{})
+}
+
+func (c *Client) Close() error {
+	return c.conn.Close()
+}
@@ -0,0 +1,48 @@
+// Package testharness provides the wire types, typed RPC client, and
+// barrier registry shared between the parent and child halves of the
+// userfaultfd test harness. It deliberately knows nothing about
+// *Userfaultfd internals so it can sit alongside the other uffd test
+// utilities and never get pulled into a production import path.
+package testharness
+
+// Empty stands in for net/rpc methods that take or return nothing;
+// net/rpc still requires both args and reply to be exported pointers.
+type Empty struct{}
+
+type BootstrapArgs struct {
+	MmapStart uint64
+	Pagesize  int64
+	TotalSize int64
+	AlwaysWP  bool
+	// Barriers gates the test-only worker hooks. Off by default so
+	// the worker hot path stays a single nil-pointer load + branch
+	// in non-race tests.
+	Barriers bool
+	Content  []byte
+}
+
+type BootstrapReply struct{}
+
+// PageStateEntry is the wire form of the parent package's pageState
+// enum; the parent translates back at the boundary.
+type PageStateEntry struct {
+	State  uint8
+	Offset uint64
+}
+
+type PageStatesReply struct {
+	Entries []PageStateEntry
+}
+
+type FaultBarrierArgs struct {
+	Addr  uint64
+	Point uint8
+}
+
+type FaultBarrierReply struct {
+	Token uint64
+}
+
+type TokenArgs struct {
+	Token uint64
+}
@@ -289,7 +289,7 @@ func TestAsyncWriteProtection(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			h, err := configureCrossProcessTest(t, testConfig{
+			h, err := configureCrossProcessTest(t.Context(), t, testConfig{
 				pagesize:      tt.pagesize,
 				numberOfPages: tt.numberOfPages,
 				alwaysWP:      tt.alwaysWP,