refactor(uffd): swap pageTracker for block.StateTracker, add removed state

Replace the map-based pageTracker with block.StateTracker[pageState], a
roaring-bitmap-backed tracker with O(1) range ops. pageState gains a
third value, removed, which is wired at the type level but not yet
written anywhere -- #2520 adds the REMOVE-event handler that produces
it. Page indices are computed at the call site via header.BlockIdx.
pageStateEntries is updated to iterate the exported bitmaps so the
cross-process test harness keeps working.

Author: ValentaTomas

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

@@ -1,33 +1,9 @@
 package userfaultfd
 
-import "sync"
-
 type pageState uint8
 
 const (
 	missing pageState = iota
 	faulted
+	removed
 )
-
-type pageTracker struct {
-	pageSize uintptr
-
-	m  map[uintptr]pageState
-	mu sync.RWMutex
-}
-
-func newPageTracker(pageSize uintptr) *pageTracker {
-	return &pageTracker{
-		pageSize: pageSize,
-		m:        make(map[uintptr]pageState),
-	}
-}
-
-func (pt *pageTracker) setState(start, end uintptr, state pageState) {
-	pt.mu.Lock()
-	defer pt.mu.Unlock()
-
-	for addr := start; addr < end; addr += pt.pageSize {
-		pt.m[addr] = state
-	}
-}

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

@@ -10,6 +10,8 @@ import (
 	"os"
 	"sync"
 
+	"github.com/RoaringBitmap/roaring/v2"
+
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/fdexit"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/memory"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/testutils/testharness"
@@ -186,24 +188,24 @@ func (p *Paging) Resume(_ *testharness.Empty, _ *testharness.Empty) error {
 }
 
 // pageStateEntries returns a wire-format snapshot of pageTracker.
-// settleRequests.Lock drains fault workers (mirrors PrefetchData);
-// pageTracker.mu.RLock is defensive against a future REMOVE writer
-// that mutates pageTracker.m outside settleRequests.
+// settleRequests.Lock drains fault workers (mirrors PrefetchData) so
+// the snapshot is consistent w.r.t. concurrent installs.
 func (u *Userfaultfd) pageStateEntries() ([]testharness.PageStateEntry, error) {
 	u.settleRequests.Lock()
 	defer u.settleRequests.Unlock()
 
-	u.pageTracker.mu.RLock()
-	defer u.pageTracker.mu.RUnlock()
-
-	entries := make([]testharness.PageStateEntry, 0, len(u.pageTracker.m))
-	for addr, state := range u.pageTracker.m {
-		offset, err := u.ma.GetOffset(addr)
-		if err != nil {
-			return nil, fmt.Errorf("address %#x not in mapping: %w", addr, err)
+	bmFaulted, bmRemoved := u.pageTracker.Export()
+	entries := make([]testharness.PageStateEntry, 0, bmFaulted.GetCardinality()+bmRemoved.GetCardinality())
+	emit := func(bm *roaring.Bitmap, state pageState) {
+		for _, idx := range bm.ToArray() {
+			entries = append(entries, testharness.PageStateEntry{
+				State:  uint8(state),
+				Offset: uint64(idx) * uint64(u.pageSize),
+			})
 		}
-		entries = append(entries, testharness.PageStateEntry{State: uint8(state), Offset: uint64(offset)})
 	}
+	emit(bmFaulted, faulted)
+	emit(bmRemoved, removed)
 
 	return entries, nil
 }

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

@@ -22,6 +22,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/fdexit"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/memory"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
 var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/userfaultfd")
@@ -52,7 +53,7 @@ type Userfaultfd struct {
 	src         block.Slicer
 	ma          *memory.Mapping
 	pageSize    uintptr
-	pageTracker *pageTracker
+	pageTracker *block.StateTracker[pageState]
 
 	// We use the settleRequests to guard the pageTracker so we can access a consistent state of the pageTracker after the requests are finished.
 	settleRequests sync.RWMutex
@@ -92,7 +93,7 @@ func NewUserfaultfdFromFd(fd uintptr, src block.Slicer, m *memory.Mapping, logge
 		fd:              Fd(fd),
 		src:             src,
 		pageSize:        uintptr(blockSize),
-		pageTracker:     newPageTracker(uintptr(blockSize)),
+		pageTracker:     block.NewStateTracker(missing, faulted, removed),
 		prefetchTracker: block.NewPrefetchTracker(blockSize),
 		ma:              m,
 		logger:          logger,
@@ -418,7 +419,8 @@ retryLoop:
 		return fmt.Errorf("failed uffdio copy: %w", joinedErr)
 	}
 
-	u.pageTracker.setState(addr, addr+u.pageSize, faulted)
+	idx := uint64(header.BlockIdx(offset, int64(u.pageSize)))
+	u.pageTracker.SetRange(idx, idx+1, faulted)
 	u.prefetchTracker.Add(offset, accessType)
 
 	return nil