Add UFFD dirty tracking

.github/actions/start-services/action.yml

@@ -101,7 +101,6 @@ runs:
         STORAGE_PROVIDER: "Local"
         ENVIRONMENT: "local"
         OTEL_COLLECTOR_GRPC_ENDPOINT: "localhost:4317"
-        MAX_PARALLEL_MEMFILE_SNAPSHOTTING: "2"
         SHARED_CHUNK_CACHE_PATH: "./.e2b-chunk-cache"
         EDGE_TOKEN: "abdcdefghijklmnop"
       run: |

packages/orchestrator/Makefile

@@ -49,7 +49,6 @@ run-debug:
 	GCP_DOCKER_REPOSITORY_NAME=$(GCP_DOCKER_REPOSITORY_NAME) \
 	GOOGLE_SERVICE_ACCOUNT_BASE64=$(GOOGLE_SERVICE_ACCOUNT_BASE64) \
 	OTEL_COLLECTOR_GRPC_ENDPOINT=$(OTEL_COLLECTOR_GRPC_ENDPOINT) \
-	MAX_PARALLEL_MEMFILE_SNAPSHOTTING=$(MAX_PARALLEL_MEMFILE_SNAPSHOTTING) \
 	./bin/orchestrator
 
 define setup_local_env

packages/orchestrator/benchmark_test.go

@@ -304,9 +304,8 @@ func BenchmarkBaseImageLaunch(b *testing.B) {
 type testCycle string
 
 const (
-	onlyStart        testCycle = "only-start"
-	startAndPause    testCycle = "start-and-pause"
-	startPauseResume testCycle = "start-pause-resume"
+	onlyStart     testCycle = "only-start"
+	startAndPause testCycle = "start-and-pause"
 )
 
 type testContainer struct {

packages/orchestrator/internal/sandbox/block/cache_test.go

@@ -1,9 +1,7 @@
 package block
 
 import (
-	"bytes"
 	"crypto/rand"
-	"fmt"
 	"math"
 	"os"
 	"syscall"
@@ -60,7 +58,7 @@ func TestCopyFromProcess_FullRange(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, int(size), n)
 
-	require.NoError(t, compareData(data[:n], mem[:n]))
+	require.NoError(t, testutils.ErrorFromByteSlicesDifference(mem[:n], data[:n]))
 }
 
 func TestCopyFromProcess_LargeRanges(t *testing.T) {
@@ -94,19 +92,19 @@ func TestCopyFromProcess_LargeRanges(t *testing.T) {
 	n, err := cache.ReadAt(data1, 0)
 	require.NoError(t, err)
 	require.Equal(t, int(pageSize), n)
-	require.NoError(t, compareData(data1[:n], mem[0:pageSize]))
+	require.NoError(t, testutils.ErrorFromByteSlicesDifference(mem[0:pageSize], data1[:n]))
 
 	data2 := make([]byte, pageSize)
 	n, err = cache.ReadAt(data2, int64(pageSize))
 	require.NoError(t, err)
 	require.Equal(t, int(pageSize), n)
-	require.NoError(t, compareData(data2[:n], mem[pageSize*3:pageSize*4]))
+	require.NoError(t, testutils.ErrorFromByteSlicesDifference(mem[pageSize*3:pageSize*4], data2[:n]))
 
 	data3 := make([]byte, pageSize)
 	n, err = cache.ReadAt(data3, int64(pageSize*2))
 	require.NoError(t, err)
 	require.Equal(t, int(pageSize), n)
-	require.NoError(t, compareData(data3[:n], mem[pageSize:pageSize*2]))
+	require.NoError(t, testutils.ErrorFromByteSlicesDifference(mem[pageSize:pageSize*2], data3[:n]))
 }
 
 func TestCopyFromProcess_MultipleRanges(t *testing.T) {
@@ -152,7 +150,7 @@ func TestCopyFromProcess_MultipleRanges(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, int(pageSize), n)
 
-		require.NoError(t, compareData(data[:n], mem[alignedOffset:alignedOffset+int64(pageSize)]))
+		require.NoError(t, testutils.ErrorFromByteSlicesDifference(mem[alignedOffset:alignedOffset+int64(pageSize)], data[:n]))
 	}
 }
 
@@ -192,13 +190,13 @@ func TestCopyFromProcess_HugepageToRegularPage(t *testing.T) {
 	n, err = cache.ReadAt(data, 0)
 	require.NoError(t, err)
 	require.Equal(t, int(pageSize*2), n)
-	require.NoError(t, compareData(data[:n], mem[0:pageSize*2]))
+	require.NoError(t, testutils.ErrorFromByteSlicesDifference(mem[0:pageSize*2], data[:n]))
 
 	data = make([]byte, pageSize*4)
 	n, err = cache.ReadAt(data, pageSize*2)
 	require.NoError(t, err)
 	require.Equal(t, int(pageSize*4), n)
-	require.NoError(t, compareData(data[:n], mem[pageSize*4:pageSize*8]))
+	require.NoError(t, testutils.ErrorFromByteSlicesDifference(mem[pageSize*4:pageSize*8], data[:n]))
 }
 
 func TestEmptyRanges(t *testing.T) {
@@ -218,17 +216,6 @@ func TestEmptyRanges(t *testing.T) {
 	})
 }
 
-func compareData(readBytes []byte, expectedBytes []byte) error {
-	// 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)
-
-		return fmt.Errorf("content mismatch: want '%x, got %x at index %d", want, got, idx)
-	}
-
-	return nil
-}
-
 func TestSplitOversizedRanges(t *testing.T) {
 	t.Parallel()
 
@@ -395,7 +382,7 @@ func TestCopyFromProcess_Exceed_MAX_RW_COUNT(t *testing.T) {
 	n, err = cache.ReadAt(data, 0)
 	require.NoError(t, err)
 	require.Equal(t, int(size), n)
-	require.NoError(t, compareData(data[:n], mem[0:size]))
+	require.NoError(t, testutils.ErrorFromByteSlicesDifference(mem[0:size], data[:n]))
 }
 
 func BenchmarkCopyFromHugepagesFile(b *testing.B) {

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

@@ -7,6 +7,8 @@ import (
 	"sync"
 )
 
+var ErrFdExit = errors.New("fd exit signal")
+
 // FdExit is a wrapper around a pipe that allows to signal the exit of the uffd.
 type FdExit struct {
 	r    *os.File

packages/orchestrator/internal/sandbox/uffd/testutils/diff_byte.go

@@ -1,20 +1,31 @@
 package testutils
 
+import (
+	"errors"
+	"fmt"
+)
+
 // FirstDifferentByte returns the first byte index where a and b differ.
 // It also returns the differing byte values (want, got).
 // If slices are identical, it returns idx -1.
-func FirstDifferentByte(a, b []byte) (idx int, want, got byte) {
-	smallerSize := min(len(a), len(b))
+func ErrorFromByteSlicesDifference(expected, actual []byte) error {
+	var errs []error
 
-	for i := range smallerSize {
-		if a[i] != b[i] {
-			return i, b[i], a[i]
-		}
+	if len(expected) > len(actual) {
+		errs = append(errs, fmt.Errorf("expected slice (%d bytes) is longer than actual slice (%d bytes)", len(expected), len(actual)))
+	} else if len(expected) < len(actual) {
+		errs = append(errs, fmt.Errorf("actual slice (%d bytes) is longer than expected slice (%d bytes)", len(actual), len(expected)))
 	}
 
-	if len(a) != len(b) {
-		return smallerSize, 0, 0
+	smallerSize := min(len(expected), len(actual))
+
+	for i := range smallerSize {
+		if expected[i] != actual[i] {
+			errs = append(errs, fmt.Errorf("first different byte: want '%x', got '%x' at index %d", expected[i], actual[i], i))
+
+			break
+		}
 	}
 
-	return -1, 0, 0
+	return errors.Join(errs...)
 }

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

@@ -140,7 +140,7 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error {
 	m := memory.NewMapping(regions)
 
 	uffd, err := userfaultfd.NewUserfaultfdFromFd(
-		uintptr(fds[0]),
+		userfaultfd.Fd(fds[0]),
 		u.memfile,
 		m,
 		logger.L().With(logger.WithSandboxID(sandboxId)),
@@ -164,6 +164,9 @@ func (u *Uffd) handle(ctx context.Context, sandboxId string) error {
 		ctx,
 		u.fdExit,
 	)
+	if errors.Is(err, fdexit.ErrFdExit) {
+		return nil
+	}
 	if err != nil {
 		return fmt.Errorf("failed handling uffd: %w", err)
 	}