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.

Inline the 3-line pageState enum into userfaultfd.go and drop the
dedicated page_tracker.go now that pageTracker is gone.

Convert block.StateTracker's NewStateTracker / SetRange API from panics
to errors. Distinct-state validation and unsupported-state checks now
return fmt.Errorf descriptors; the userfaultfd-side init propagates the
constructor error through NewUserfaultfdFromFd, and the SetRange call
in the worker path logs and continues since these errors only fire on
programming bugs.

Author: ValentaTomas

packages/orchestrator/pkg/sandbox/block/state_tracker.go

@@ -17,11 +17,11 @@ type StateTracker[S comparable] struct {
 
 // NewStateTracker requires three distinct states. Duplicates are a
 // programming error — the switch in SetRange would silently favour the
-// first matching case and corrupt bitmap state — so we panic at
+// first matching case and corrupt bitmap state — so we reject them at
 // construction rather than defer the bug to a later SetRange call.
-func NewStateTracker[S comparable](defaultState, a, b S) *StateTracker[S] {
+func NewStateTracker[S comparable](defaultState, a, b S) (*StateTracker[S], error) {
 	if defaultState == a || defaultState == b || a == b {
-		panic(fmt.Sprintf("block.NewStateTracker: states must be distinct (default=%v a=%v b=%v)", defaultState, a, b))
+		return nil, fmt.Errorf("block.NewStateTracker: states must be distinct (default=%v a=%v b=%v)", defaultState, a, b)
 	}
 
 	return &StateTracker[S]{
@@ -30,15 +30,15 @@ func NewStateTracker[S comparable](defaultState, a, b S) *StateTracker[S] {
 		b:            b,
 		bmA:          roaring.New(),
 		bmB:          roaring.New(),
-	}
+	}, nil
 }
 
 // SetRange takes uint64 because roaring's range API allows end = 1<<32
 // (the half-open upper bound of a 32-bit bitmap); Get stays uint32 since
 // no 33-bit value can ever be a bitmap member.
-func (t *StateTracker[S]) SetRange(start, end uint64, state S) {
+func (t *StateTracker[S]) SetRange(start, end uint64, state S) error {
 	if end <= start {
-		return
+		return nil
 	}
 
 	t.mu.Lock()
@@ -58,9 +58,11 @@ func (t *StateTracker[S]) SetRange(start, end uint64, state S) {
 		// S is constrained only to comparable, so the compiler can't
 		// prove exhaustiveness. A silent no-op here would hide a
 		// programming error (caller added a state but forgot to wire
-		// it); panic makes it fail fast in tests.
-		panic(fmt.Sprintf("block.StateTracker.SetRange: unsupported state %v (only default=%v a=%v b=%v allowed)", state, t.defaultState, t.a, t.b))
+		// it); surfacing an error makes it fail fast in tests.
+		return fmt.Errorf("block.StateTracker.SetRange: unsupported state %v (only default=%v a=%v b=%v allowed)", state, t.defaultState, t.a, t.b)
 	}
+
+	return nil
 }
 
 func (t *StateTracker[S]) Export() (a, b *roaring.Bitmap) {

packages/orchestrator/pkg/sandbox/block/state_tracker_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 type ts uint8
@@ -23,44 +24,47 @@ func TestStateTracker(t *testing.T) {
 
 	t.Run("initial state is default", func(t *testing.T) {
 		t.Parallel()
-		s := NewStateTracker(tsDefault, tsA, tsB)
+		s, err := NewStateTracker(tsDefault, tsA, tsB)
+		require.NoError(t, err)
 		assert.Equal(t, tsDefault, s.Get(0))
 		assert.Equal(t, tsDefault, s.Get(123))
 	})
 
 	t.Run("transitions", func(t *testing.T) {
 		t.Parallel()
-		s := NewStateTracker(tsDefault, tsA, tsB)
+		s, err := NewStateTracker(tsDefault, tsA, tsB)
+		require.NoError(t, err)
 
-		s.SetRange(0, 1, tsA)
+		require.NoError(t, s.SetRange(0, 1, tsA))
 		assert.Equal(t, tsA, s.Get(0))
 
-		s.SetRange(0, 1, tsB)
+		require.NoError(t, s.SetRange(0, 1, tsB))
 		assert.Equal(t, tsB, s.Get(0), "a→b should flip the page")
 		bmA, bmB := s.Export()
 		assert.False(t, bmA.Contains(0), "a→b must clear bmA")
 		assert.True(t, bmB.Contains(0), "a→b must add to bmB")
 
-		s.SetRange(0, 1, tsA)
+		require.NoError(t, s.SetRange(0, 1, tsA))
 		assert.Equal(t, tsA, s.Get(0), "b→a should flip back")
 
-		s.SetRange(0, 1, tsDefault)
+		require.NoError(t, s.SetRange(0, 1, tsDefault))
 		assert.Equal(t, tsDefault, s.Get(0), "→default must clear")
 		bmA, bmB = s.Export()
 		assert.False(t, bmA.Contains(0))
 		assert.False(t, bmB.Contains(0))
 
-		s.SetRange(0, 1, tsA)
-		s.SetRange(0, 1, tsA)
+		require.NoError(t, s.SetRange(0, 1, tsA))
+		require.NoError(t, s.SetRange(0, 1, tsA))
 		assert.Equal(t, tsA, s.Get(0), "a→a is idempotent")
 	})
 
 	t.Run("partial overlap moves only the overlapping pages", func(t *testing.T) {
 		t.Parallel()
-		s := NewStateTracker(tsDefault, tsA, tsB)
+		s, err := NewStateTracker(tsDefault, tsA, tsB)
+		require.NoError(t, err)
 
-		s.SetRange(0, 10, tsA)
-		s.SetRange(3, 7, tsB)
+		require.NoError(t, s.SetRange(0, 10, tsA))
+		require.NoError(t, s.SetRange(3, 7, tsB))
 
 		for i := range uint32(3) {
 			assert.Equal(t, tsA, s.Get(i), "page %d outside overlap stays a", i)
@@ -77,38 +81,45 @@ func TestStateTracker(t *testing.T) {
 
 	t.Run("empty and inverted ranges are no-ops", func(t *testing.T) {
 		t.Parallel()
-		s := NewStateTracker(tsDefault, tsA, tsB)
+		s, err := NewStateTracker(tsDefault, tsA, tsB)
+		require.NoError(t, err)
 
-		s.SetRange(5, 5, tsA)
-		s.SetRange(7, 3, tsB)
+		require.NoError(t, s.SetRange(5, 5, tsA))
+		require.NoError(t, s.SetRange(7, 3, tsB))
 		bmA, bmB := s.Export()
 		assert.True(t, bmA.IsEmpty())
 		assert.True(t, bmB.IsEmpty())
 	})
 
 	t.Run("Export returns clones, not aliases", func(t *testing.T) {
 		t.Parallel()
-		s := NewStateTracker(tsDefault, tsA, tsB)
-		s.SetRange(0, 1, tsA)
+		s, err := NewStateTracker(tsDefault, tsA, tsB)
+		require.NoError(t, err)
+		require.NoError(t, s.SetRange(0, 1, tsA))
 
 		bmA, _ := s.Export()
 		bmA.Add(99)
 
 		assert.Equal(t, tsDefault, s.Get(99), "mutating a clone must not leak into the tracker")
 	})
 
-	t.Run("NewStateTracker rejects non-distinct states", func(t *testing.T) {
+	t.Run("NewStateTracker errors on non-distinct states", func(t *testing.T) {
 		t.Parallel()
-		assert.Panics(t, func() { NewStateTracker(tsA, tsA, tsB) }, "default == a must panic")
-		assert.Panics(t, func() { NewStateTracker(tsA, tsB, tsB) }, "a == b must panic")
-		assert.Panics(t, func() { NewStateTracker(tsA, tsB, tsA) }, "default == b must panic")
-		assert.Panics(t, func() { NewStateTracker(tsA, tsA, tsA) }, "all-equal must panic")
+		_, err := NewStateTracker(tsA, tsA, tsB)
+		require.Error(t, err, "default == a must error")
+		_, err = NewStateTracker(tsA, tsB, tsB)
+		require.Error(t, err, "a == b must error")
+		_, err = NewStateTracker(tsA, tsB, tsA)
+		require.Error(t, err, "default == b must error")
+		_, err = NewStateTracker(tsA, tsA, tsA)
+		require.Error(t, err, "all-equal must error")
 	})
 
-	t.Run("SetRange panics on unsupported state", func(t *testing.T) {
+	t.Run("SetRange errors on unsupported state", func(t *testing.T) {
 		t.Parallel()
-		s := NewStateTracker(tsDefault, tsA, tsB)
-		assert.Panics(t, func() { s.SetRange(0, 1, ts(99)) },
-			"unregistered state value must panic, not silently no-op")
+		s, err := NewStateTracker(tsDefault, tsA, tsB)
+		require.NoError(t, err)
+		require.Error(t, s.SetRange(0, 1, ts(99)),
+			"unregistered state value must error, not silently no-op")
 	})
 }

packages/orchestrator/pkg/sandbox/uffd/userfaultfd/page_tracker.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/rpc_services_test.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")
@@ -46,13 +47,23 @@ func hasEvent(revents, event int16) bool {
 	return revents&event != 0
 }
 
+// pageState tracks which UFFD page-management action has been applied
+// to each registered page. The default (zero) value is missing.
+type pageState uint8
+
+const (
+	missing pageState = iota
+	faulted
+	removed
+)
+
 type Userfaultfd struct {
 	fd Fd
 
 	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
@@ -88,11 +99,16 @@ func NewUserfaultfdFromFd(fd uintptr, src block.Slicer, m *memory.Mapping, logge
 		}
 	}
 
+	pageTracker, err := block.NewStateTracker(missing, faulted, removed)
+	if err != nil {
+		return nil, fmt.Errorf("failed to init page tracker: %w", err)
+	}
+
 	u := &Userfaultfd{
 		fd:              Fd(fd),
 		src:             src,
 		pageSize:        uintptr(blockSize),
-		pageTracker:     newPageTracker(uintptr(blockSize)),
+		pageTracker:     pageTracker,
 		prefetchTracker: block.NewPrefetchTracker(blockSize),
 		ma:              m,
 		logger:          logger,
@@ -418,7 +434,14 @@ 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)))
+	if err := u.pageTracker.SetRange(idx, idx+1, faulted); err != nil {
+		// Programming bug only — the serve loop is still healthy and
+		// the page is correctly installed in guest memory. Log and
+		// continue rather than abort.
+		u.logger.Error(ctx, "UFFD serve pageTracker SetRange error",
+			zap.Uint64("idx", idx), zap.Error(err))
+	}
 	u.prefetchTracker.Add(offset, accessType)
 
 	return nil