Extract low level UFFD handler to a separate package

.github/workflows/pr-tests.yml

@@ -33,6 +33,16 @@ jobs:
         with:
           job_key: unit-tests-${{ matrix.package }}
 
+      - name: Enable unprivileged uffd mode
+        run: |
+          echo 1 | sudo tee /proc/sys/vm/unprivileged_userfaultfd
+
+      - name: Enable hugepages
+        run: |
+          sudo mkdir -p /mnt/hugepages
+          sudo mount -t hugetlbfs none /mnt/hugepages
+          echo 128 | sudo tee /proc/sys/vm/nr_hugepages
+
       - name: Run tests
         working-directory: ${{ matrix.package }}
         run: go test -v ${{ matrix.test_path }}

.gitignore

@@ -13,3 +13,4 @@ tests/periodic-test/build-template/e2b.toml
 .air
 go.work.sum
 .infisical.json
+.vscode/mise-tools/

.vscode/settings.json

@@ -89,5 +89,11 @@
     "editor.codeActionsOnSave": {
       "source.organizeImports": "explicit"
     }
+  },
+  "go.goroot": "${workspaceFolder}/.vscode/mise-tools/goRoot",
+  "go.alternateTools": {
+    "go": "${workspaceFolder}/.vscode/mise-tools/go",
+    "dlv": "${workspaceFolder}/.vscode/mise-tools/dlv",
+    "gopls": "${workspaceFolder}/.vscode/mise-tools/gopls"
   }
 }

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

@@ -1,66 +1,70 @@
 package block
 
 import (
-	"context"
-	"fmt"
 	"sync"
-	"sync/atomic"
 
 	"github.com/bits-and-blooms/bitset"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
-type TrackedSliceDevice struct {
-	data      ReadonlyDevice
-	blockSize int64
+type Tracker struct {
+	b  *bitset.BitSet
+	mu sync.RWMutex
 
-	nilTracking atomic.Bool
-	dirty       *bitset.BitSet
-	dirtyMu     sync.Mutex
-	empty       []byte
+	blockSize int64
 }
 
-func NewTrackedSliceDevice(blockSize int64, device ReadonlyDevice) (*TrackedSliceDevice, error) {
-	return &TrackedSliceDevice{
-		data:      device,
-		empty:     make([]byte, blockSize),
+func NewTracker(blockSize int64) *Tracker {
+	return &Tracker{
+		// The bitset resizes automatically based on the maximum set bit.
+		b:         bitset.New(0),
 		blockSize: blockSize,
-	}, nil
+	}
 }
 
-func (t *TrackedSliceDevice) Disable() error {
-	size, err := t.data.Size()
-	if err != nil {
-		return fmt.Errorf("failed to get device size: %w", err)
-	}
+func (t *Tracker) Has(off int64) bool {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
 
-	t.dirty = bitset.New(uint(header.TotalBlocks(size, t.blockSize)))
-	// We are starting with all being dirty.
-	t.dirty.FlipRange(0, t.dirty.Len())
+	return t.b.Test(uint(header.BlockIdx(off, t.blockSize)))
+}
+
+func (t *Tracker) Add(off int64) bool {
+	t.mu.Lock()
+	defer t.mu.Unlock()
 
-	t.nilTracking.Store(true)
+	if t.b.Test(uint(header.BlockIdx(off, t.blockSize))) {
+		return false
+	}
 
-	return nil
+	t.b.Set(uint(header.BlockIdx(off, t.blockSize)))
+
+	return true
 }
 
-func (t *TrackedSliceDevice) Slice(ctx context.Context, off int64, length int64) ([]byte, error) {
-	if t.nilTracking.Load() {
-		t.dirtyMu.Lock()
-		t.dirty.Clear(uint(header.BlockIdx(off, t.blockSize)))
-		t.dirtyMu.Unlock()
+func (t *Tracker) Reset() {
+	t.mu.Lock()
+	defer t.mu.Unlock()
 
-		return t.empty, nil
-	}
+	t.b.ClearAll()
+}
 
-	return t.data.Slice(ctx, off, length)
+// BitSet returns a clone of the bitset and the block size.
+func (t *Tracker) BitSet() *bitset.BitSet {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return t.b.Clone()
 }
 
-// Return which bytes were not read since Disable.
-// This effectively returns the bytes that have been requested after paused vm and are not dirty.
-func (t *TrackedSliceDevice) Dirty() *bitset.BitSet {
-	t.dirtyMu.Lock()
-	defer t.dirtyMu.Unlock()
+func (t *Tracker) BlockSize() int64 {
+	return t.blockSize
+}
 
-	return t.dirty.Clone()
+func (t *Tracker) Clone() *Tracker {
+	return &Tracker{
+		b:         t.BitSet(),
+		blockSize: t.BlockSize(),
+	}
 }

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

@@ -0,0 +1,109 @@
+package block
+
+import (
+	"testing"
+)
+
+func TestTracker_AddAndHas(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	offset := int64(pageSize * 4)
+
+	// Initially should not be marked
+	if tr.Has(offset) {
+		t.Errorf("Expected offset %d not to be marked initially", offset)
+	}
+
+	// After adding, should be marked
+	tr.Add(offset)
+	if !tr.Has(offset) {
+		t.Errorf("Expected offset %d to be marked after Add", offset)
+	}
+}
+
+func TestTracker_Reset(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	offset := int64(pageSize * 4)
+
+	// Add offset and verify it's marked
+	tr.Add(offset)
+	if !tr.Has(offset) {
+		t.Errorf("Expected offset %d to be marked after Add", offset)
+	}
+
+	// After reset, should not be marked
+	tr.Reset()
+	if tr.Has(offset) {
+		t.Errorf("Expected offset %d to be cleared after Reset", offset)
+	}
+}
+
+func TestTracker_MultipleOffsets(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	offsets := []int64{0, pageSize, 2 * pageSize, 10 * pageSize}
+
+	// Add multiple offsets
+	for _, o := range offsets {
+		tr.Add(o)
+	}
+
+	// Verify all offsets are marked
+	for _, o := range offsets {
+		if !tr.Has(o) {
+			t.Errorf("Expected offset %d to be marked", o)
+		}
+	}
+}
+
+func TestTracker_ResetClearsAll(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	offsets := []int64{0, pageSize, 2 * pageSize, 10 * pageSize}
+
+	// Add multiple offsets
+	for _, o := range offsets {
+		tr.Add(o)
+	}
+
+	// Reset should clear all
+	tr.Reset()
+
+	// Verify all offsets are cleared
+	for _, o := range offsets {
+		if tr.Has(o) {
+			t.Errorf("Expected offset %d to be cleared after Reset", o)
+		}
+	}
+}
+
+func TestTracker_MisalignedOffset(t *testing.T) {
+	const pageSize = 4096
+	tr := NewTracker(pageSize)
+
+	// Test with misaligned offset
+	misalignedOffset := int64(123)
+	tr.Add(misalignedOffset)
+
+	// Should be set for the block containing the offset—that is, block 0 (0..4095)
+	if !tr.Has(misalignedOffset) {
+		t.Errorf("Expected misaligned offset %d to be marked (should mark its containing block)", misalignedOffset)
+	}
+
+	// Now check that any offset in the same block is also considered marked
+	anotherOffsetInSameBlock := int64(1000)
+	if !tr.Has(anotherOffsetInSameBlock) {
+		t.Errorf("Expected offset %d to be marked as in same block as %d", anotherOffsetInSameBlock, misalignedOffset)
+	}
+
+	// But not for a different block
+	offsetInNextBlock := int64(pageSize) // convert to int64 to match Has signature
+	if tr.Has(offsetInNextBlock) {
+		t.Errorf("Did not expect offset %d to be marked", offsetInNextBlock)
+	}
+}

packages/orchestrator/internal/sandbox/sandbox.go

@@ -671,10 +671,18 @@ func (s *Sandbox) Pause(
 		return nil, fmt.Errorf("failed to pause VM: %w", err)
 	}
 
-	if err := s.memory.Disable(); err != nil {
+	err = s.memory.Disable(ctx)
+	if err != nil {
 		return nil, fmt.Errorf("failed to disable uffd: %w", err)
 	}
 
+	dirty, err := s.memory.Dirty(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get dirty pages: %w", err)
+	}
+
+	dirtyPages := dirty.BitSet()
+
 	// Snapfile is not closed as it's returned and cached for later use (like resume)
 	snapfile := template.NewLocalFileLink(snapshotTemplateFiles.CacheSnapfilePath())
 	// Memfile is also closed on diff creation processing
@@ -718,7 +726,7 @@ func (s *Sandbox) Pause(
 		originalMemfile.Header(),
 		&MemoryDiffCreator{
 			memfile:    memfile,
-			dirtyPages: s.memory.Dirty(),
+			dirtyPages: dirtyPages,
 			blockSize:  originalMemfile.BlockSize(),
 			doneHook: func(context.Context) error {
 				return memfile.Close()
@@ -929,7 +937,7 @@ func serveMemory(
 	ctx, span := tracer.Start(ctx, "serve-memory")
 	defer span.End()
 
-	fcUffd, err := uffd.New(memfile, socketPath, memfile.BlockSize())
+	fcUffd, err := uffd.New(memfile, socketPath)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create uffd: %w", err)
 	}

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

@@ -0,0 +1,45 @@
+package memory
+
+import (
+	"fmt"
+)
+
+type Mapping struct {
+	Regions []Region
+}
+
+func NewMapping(regions []Region) *Mapping {
+	return &Mapping{Regions: regions}
+}
+
+// GetOffset returns the relative offset and the page size of the mapped range for a given address.
+func (m *Mapping) GetOffset(hostVirtAddr uintptr) (int64, uint64, error) {
+	for _, r := range m.Regions {
+		if hostVirtAddr >= r.BaseHostVirtAddr && hostVirtAddr < r.endHostVirtAddr() {
+			return r.shiftedOffset(hostVirtAddr), uint64(r.PageSize), nil
+		}
+	}
+
+	return 0, 0, fmt.Errorf("address %d not found in any mapping", hostVirtAddr)
+}
+
+// GetHostVirtAddr returns the host virtual address for a given offset.
+func (m *Mapping) GetHostVirtAddr(off int64) (int64, uint64, error) {
+	r, err := m.getHostVirtRegion(off)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	return int64(r.shiftedHostVirtAddr(off)), uint64(r.PageSize), nil
+}
+
+// getHostVirtRegion returns the region that contains the given offset.
+func (m *Mapping) getHostVirtRegion(off int64) (*Region, error) {
+	for _, r := range m.Regions {
+		if off >= int64(r.Offset) && off < r.endOffset() {
+			return &r, nil
+		}
+	}
+
+	return nil, fmt.Errorf("offset %d not found in any mapping", off)
+}