More concurrency fixes

Signed-off-by: Babis Chalios <babis.chalios@e2b.dev>

Author: bchalios

packages/orchestrator/internal/sandbox/uffd/userfaultfd/prefault.go

@@ -23,7 +23,9 @@ func (u *Userfaultfd) Prefault(ctx context.Context, offset int64, data []byte) e
 		return fmt.Errorf("data length (%d) does not match pagesize (%d)", len(data), u.pageSize)
 	}
 
-	return u.faultPage(ctx, nil, addr, offset, u.pageSize, directDataSource{data, int64(u.pageSize)}, nil, block.Prefetch)
+	_, err = u.faultPage(ctx, addr, offset, u.pageSize, directDataSource{data, int64(u.pageSize)}, nil, block.Prefetch)
+
+	return err
 }
 
 // directDataSource wraps a byte slice to implement block.Slicer for prefaulting.

packages/orchestrator/internal/sandbox/uffd/userfaultfd/userfaultfd.go

@@ -275,7 +275,9 @@ func (u *Userfaultfd) Serve(
 			u.pageTracker.markRemovedRange(start, end)
 		}
 
-		for _, pfEvent := range pfEvents {
+		for i := range pfEvents {
+			pfEvent := pfEvents[i]
+
 			// We don't handle minor page faults
 			if pfEvent.flags&UFFD_PAGEFAULT_FLAG_MINOR != 0 {
 				return fmt.Errorf("unexpected MINOR pagefault event, closing UFFD")
@@ -307,10 +309,38 @@ func (u *Userfaultfd) Serve(
 			case faulted:
 				continue
 			case removed:
-				u.wg.Go(func() error { return u.zeroPage(&pfEvent, addr, u.pageSize) })
+				u.wg.Go(func() error {
+					handled, err := u.zeroPage(addr, u.pageSize)
+					if err != nil {
+						return err
+					}
+
+					if !handled {
+						// This happens when a remove event arrives in the UFFD file descriptor while
+						// we are trying to copy()/zero() a page. We need to retry.
+						u.deferredRWMutex.Lock()
+						defer u.deferredRWMutex.Unlock()
+
+						u.deferredPfEvents = append(u.deferredPfEvents, pfEvent)
+					}
+
+					return nil
+				})
 			case unfaulted:
 				u.wg.Go(func() error {
-					return u.faultPage(ctx, &pfEvent, addr, offset, u.pageSize, u.src, fdExit.SignalExit, accessType)
+					handled, err := u.faultPage(ctx, addr, offset, u.pageSize, u.src, fdExit.SignalExit, accessType)
+					if err != nil {
+						return err
+					}
+
+					if !handled {
+						u.deferredRWMutex.Lock()
+						defer u.deferredRWMutex.Unlock()
+
+						u.deferredPfEvents = append(u.deferredPfEvents, pfEvent)
+					}
+
+					return nil
 				})
 			default:
 				return fmt.Errorf("unexpected pageState: %#v", state)
@@ -340,10 +370,9 @@ func (u *Userfaultfd) PrefetchData() block.PrefetchData {
 }
 
 func (u *Userfaultfd) zeroPage(
-	pfEvent *UffdPagefault,
 	addr uintptr,
 	pagesize uintptr,
-) error {
+) (bool, error) {
 	var err error
 
 	if pagesize == 4096 {
@@ -365,39 +394,31 @@ func (u *Userfaultfd) zeroPage(
 	if errors.Is(err, unix.EEXIST) {
 		u.pageTracker.markFaulted(addr)
 
-		return nil
+		return true, nil
 	}
 
 	if errors.Is(err, unix.EAGAIN) {
-		// This happens when a remove event arrives in the UFFD file descriptor while
-		// we are trying to copy()/zero() a page. We need to retry.
-		u.deferredRWMutex.Lock()
-		defer u.deferredRWMutex.Lock()
-
-		if pfEvent != nil {
-			u.deferredPfEvents = append(u.deferredPfEvents, *pfEvent)
-		}
-
-		return nil
+		return false, nil
 	}
 
-	if err != nil {
+	if err == nil {
 		u.pageTracker.markFaulted(addr)
+
+		return true, nil
 	}
 
-	return err
+	return false, err
 }
 
 func (u *Userfaultfd) faultPage(
 	ctx context.Context,
-	pfEvent *UffdPagefault,
 	addr uintptr,
 	offset int64,
 	pagesize uintptr,
 	source block.Slicer,
 	onFailure func() error,
 	accessType block.AccessType,
-) error {
+) (bool, error) {
 	span := trace.SpanFromContext(ctx)
 
 	// The RLock must be called inside the goroutine to ensure RUnlock runs via defer,
@@ -425,7 +446,7 @@ func (u *Userfaultfd) faultPage(
 		span.RecordError(joinedErr)
 		u.logger.Error(ctx, "UFFD serve data fetch error", zap.Error(joinedErr))
 
-		return fmt.Errorf("failed to read from source: %w", joinedErr)
+		return false, fmt.Errorf("failed to read from source: %w", joinedErr)
 	}
 
 	var copyMode CULong
@@ -443,18 +464,13 @@ func (u *Userfaultfd) faultPage(
 		span.SetAttributes(attribute.Bool("uffd.already_mapped", true))
 		u.pageTracker.markFaulted(addr)
 
-		return nil
+		return true, nil
 	}
 
 	if errors.Is(copyErr, unix.EAGAIN) {
 		u.logger.Debug(ctx, "UFFD copy page EAGAIN. Ignoring", zap.Uintptr("addr", addr))
 
-		u.deferredRWMutex.Lock()
-		defer u.deferredRWMutex.Unlock()
-
-		u.deferredPfEvents = append(u.deferredPfEvents, *pfEvent)
-
-		return nil
+		return false, nil
 	}
 
 	if copyErr != nil {
@@ -468,13 +484,13 @@ func (u *Userfaultfd) faultPage(
 		span.RecordError(joinedErr)
 		u.logger.Error(ctx, "UFFD serve uffdio copy error", zap.Error(joinedErr))
 
-		return fmt.Errorf("failed uffdio copy %w", joinedErr)
+		return false, fmt.Errorf("failed uffdio copy %w", joinedErr)
 	}
 
 	// Add the offset to the missing requests tracker with metadata.
 	u.missingRequests.Add(offset)
 	u.prefetchTracker.Add(offset, accessType)
 	u.pageTracker.markFaulted(addr)
 
-	return nil
+	return true, nil
 }