From 525cc2fe9d6394b40942463c8ed8b48c3b1b0b62 Mon Sep 17 00:00:00 2001
From: ValentaTomas <valenta.and.thomas@gmail.com>
Date: Sat, 2 May 2026 23:30:02 -0700
Subject: [PATCH 1/5] feat(block): add generic StateTracker

A small two-state-plus-default tracker backed by roaring bitmaps. Used by
upcoming UFFD work to track page states (Missing/Faulted/Removed) and by
NBD to track zero pages, replacing ad-hoc map-based trackers with O(1)
range ops and cheap snapshot exports.
---
 .../pkg/sandbox/block/state_tracker.go        | 85 +++++++++++++++++
 .../pkg/sandbox/block/state_tracker_test.go   | 95 +++++++++++++++++++
 2 files changed, 180 insertions(+)
 create mode 100644 packages/orchestrator/pkg/sandbox/block/state_tracker.go
 create mode 100644 packages/orchestrator/pkg/sandbox/block/state_tracker_test.go

diff --git a/packages/orchestrator/pkg/sandbox/block/state_tracker.go b/packages/orchestrator/pkg/sandbox/block/state_tracker.go
new file mode 100644
index 0000000000..79c1ae8a5c
--- /dev/null
+++ b/packages/orchestrator/pkg/sandbox/block/state_tracker.go
@@ -0,0 +1,85 @@
+package block
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/RoaringBitmap/roaring/v2"
+)
+
+type StateTracker[S comparable] struct {
+	mu sync.RWMutex
+
+	defaultState S
+	a, b         S
+	bmA, bmB     *roaring.Bitmap
+}
+
+// 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
+// construction rather than defer the bug to a later SetRange call.
+func NewStateTracker[S comparable](defaultState, a, b S) *StateTracker[S] {
+	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 &StateTracker[S]{
+		defaultState: defaultState,
+		a:            a,
+		b:            b,
+		bmA:          roaring.New(),
+		bmB:          roaring.New(),
+	}
+}
+
+// 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) {
+	if end <= start {
+		return
+	}
+
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	switch state {
+	case t.a:
+		t.bmA.AddRange(start, end)
+		t.bmB.RemoveRange(start, end)
+	case t.b:
+		t.bmB.AddRange(start, end)
+		t.bmA.RemoveRange(start, end)
+	case t.defaultState:
+		t.bmA.RemoveRange(start, end)
+		t.bmB.RemoveRange(start, end)
+	default:
+		// 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))
+	}
+}
+
+func (t *StateTracker[S]) Export() (a, b *roaring.Bitmap) {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return t.bmA.Clone(), t.bmB.Clone()
+}
+
+func (t *StateTracker[S]) Get(idx uint32) S {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	switch {
+	case t.bmA.Contains(idx):
+		return t.a
+	case t.bmB.Contains(idx):
+		return t.b
+	default:
+		return t.defaultState
+	}
+}
diff --git a/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go b/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go
new file mode 100644
index 0000000000..1de81bfd17
--- /dev/null
+++ b/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go
@@ -0,0 +1,95 @@
+package block
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type ts uint8
+
+const (
+	tsDefault ts = iota
+	tsA
+	tsB
+)
+
+// TestStateTracker exercises every transition pair (default↔a, default↔b,
+// a↔b, idempotent same-state) and confirms the two non-default bitmaps
+// stay disjoint.
+func TestStateTracker(t *testing.T) {
+	t.Parallel()
+
+	t.Run("transitions", func(t *testing.T) {
+		t.Parallel()
+		s := NewStateTracker(tsDefault, tsA, tsB)
+
+		s.SetRange(0, 1, tsA)
+		assert.Equal(t, tsA, s.Get(0))
+
+		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)
+		assert.Equal(t, tsA, s.Get(0), "b→a should flip back")
+
+		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)
+		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.SetRange(0, 10, tsA)
+		s.SetRange(3, 7, tsB)
+
+		for i := range uint32(3) {
+			assert.Equal(t, tsA, s.Get(i), "page %d outside overlap stays a", i)
+		}
+		for i := range uint32(4) {
+			page := i + 3
+			assert.Equal(t, tsB, s.Get(page), "page %d in overlap moves to b", page)
+		}
+		for i := range uint32(3) {
+			page := i + 7
+			assert.Equal(t, tsA, s.Get(page), "page %d outside overlap stays a", page)
+		}
+	})
+
+	t.Run("empty and inverted ranges are no-ops", func(t *testing.T) {
+		t.Parallel()
+		s := NewStateTracker(tsDefault, tsA, tsB)
+
+		s.SetRange(5, 5, tsA)
+		s.SetRange(7, 3, tsB)
+		bmA, bmB := s.Export()
+		assert.True(t, bmA.IsEmpty())
+		assert.True(t, bmB.IsEmpty())
+	})
+
+	t.Run("NewStateTracker rejects 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")
+	})
+
+	t.Run("SetRange panics 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")
+	})
+}

From bbef5f2af6c194b97694b54137fd8f73c417a15b Mon Sep 17 00:00:00 2001
From: ValentaTomas <valenta.and.thomas@gmail.com>
Date: Sat, 2 May 2026 23:56:01 -0700
Subject: [PATCH 2/5] 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.
---
 .../pkg/sandbox/block/state_tracker.go        | 18 ++++---
 .../pkg/sandbox/block/state_tracker_test.go   | 53 +++++++++++--------
 .../sandbox/uffd/userfaultfd/page_tracker.go  | 33 ------------
 .../uffd/userfaultfd/rpc_services_test.go     | 26 ++++-----
 .../sandbox/uffd/userfaultfd/userfaultfd.go   | 29 ++++++++--
 5 files changed, 81 insertions(+), 78 deletions(-)
 delete mode 100644 packages/orchestrator/pkg/sandbox/uffd/userfaultfd/page_tracker.go

diff --git a/packages/orchestrator/pkg/sandbox/block/state_tracker.go b/packages/orchestrator/pkg/sandbox/block/state_tracker.go
index 79c1ae8a5c..33fb473270 100644
--- a/packages/orchestrator/pkg/sandbox/block/state_tracker.go
+++ b/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) {
diff --git a/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go b/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go
index 1de81bfd17..90bee00e5b 100644
--- a/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go
+++ b/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
@@ -22,37 +23,39 @@ func TestStateTracker(t *testing.T) {
 
 	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)
@@ -69,27 +72,33 @@ 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("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")
 	})
 }
diff --git a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/page_tracker.go b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/page_tracker.go
deleted file mode 100644
index da76d310a8..0000000000
--- a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/page_tracker.go
+++ /dev/null
@@ -1,33 +0,0 @@
-package userfaultfd
-
-import "sync"
-
-type pageState uint8
-
-const (
-	missing pageState = iota
-	faulted
-)
-
-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
-	}
-}
diff --git a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/rpc_services_test.go b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/rpc_services_test.go
index 380e9c1d2f..b2a3562acf 100644
--- a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/rpc_services_test.go
+++ b/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
 }
diff --git a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
index 43b9fee492..1b9bb69695 100644
--- a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
+++ b/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")
@@ -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

From 1949d9bc215eeb484b745dd56a9163439910c932 Mon Sep 17 00:00:00 2001
From: ValentaTomas <valenta.and.thomas@gmail.com>
Date: Sun, 3 May 2026 15:58:06 -0700
Subject: [PATCH 3/5] refactor(state-tracker): trim verbose doc comments

---
 .../pkg/sandbox/block/state_tracker.go            | 15 ++++-----------
 .../pkg/sandbox/block/state_tracker_test.go       |  3 ---
 2 files changed, 4 insertions(+), 14 deletions(-)

diff --git a/packages/orchestrator/pkg/sandbox/block/state_tracker.go b/packages/orchestrator/pkg/sandbox/block/state_tracker.go
index 33fb473270..3633f39cee 100644
--- a/packages/orchestrator/pkg/sandbox/block/state_tracker.go
+++ b/packages/orchestrator/pkg/sandbox/block/state_tracker.go
@@ -15,10 +15,8 @@ type StateTracker[S comparable] struct {
 	bmA, bmB     *roaring.Bitmap
 }
 
-// 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 reject them at
-// construction rather than defer the bug to a later SetRange call.
+// NewStateTracker requires three distinct states; duplicates would alias in
+// SetRange's switch and silently corrupt the bitmaps.
 func NewStateTracker[S comparable](defaultState, a, b S) (*StateTracker[S], error) {
 	if defaultState == a || defaultState == b || a == b {
 		return nil, fmt.Errorf("block.NewStateTracker: states must be distinct (default=%v a=%v b=%v)", defaultState, a, b)
@@ -33,9 +31,7 @@ func NewStateTracker[S comparable](defaultState, a, b S) (*StateTracker[S], erro
 	}, 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.
+// SetRange takes uint64 to allow end = 1<<32 (roaring's half-open upper bound).
 func (t *StateTracker[S]) SetRange(start, end uint64, state S) error {
 	if end <= start {
 		return nil
@@ -55,10 +51,7 @@ func (t *StateTracker[S]) SetRange(start, end uint64, state S) error {
 		t.bmA.RemoveRange(start, end)
 		t.bmB.RemoveRange(start, end)
 	default:
-		// 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); surfacing an error makes it fail fast in tests.
+		// S is only `comparable`, so the compiler can't prove exhaustiveness.
 		return fmt.Errorf("block.StateTracker.SetRange: unsupported state %v (only default=%v a=%v b=%v allowed)", state, t.defaultState, t.a, t.b)
 	}
 
diff --git a/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go b/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go
index 90bee00e5b..fe01cb9134 100644
--- a/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go
+++ b/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go
@@ -15,9 +15,6 @@ const (
 	tsB
 )
 
-// TestStateTracker exercises every transition pair (default↔a, default↔b,
-// a↔b, idempotent same-state) and confirms the two non-default bitmaps
-// stay disjoint.
 func TestStateTracker(t *testing.T) {
 	t.Parallel()
 

From 84ac30ba6e00534bf5646acb75bd9bd36f0ca73c Mon Sep 17 00:00:00 2001
From: ValentaTomas <valenta.and.thomas@gmail.com>
Date: Mon, 4 May 2026 11:38:31 -0700
Subject: [PATCH 4/5] refactor(block): drop generics, unify 3 universal states

Replace the generic block.StateTracker[S] with a non-generic block.Tracker
over a fixed block.State enum (NotPresent default, Dirty, Zero). The
same three states cover both UFFD page management (NotPresent=missing,
Dirty=faulted, Zero=removed via DONTNEED/balloon) and upcoming NBD
overlay tracking (NotPresent=fall-through, Dirty=overlay-owned,
Zero=explicit zero/discard).
---
 .../pkg/sandbox/block/state_tracker.go        |  80 --------------
 .../pkg/sandbox/block/state_tracker_test.go   | 101 ------------------
 .../orchestrator/pkg/sandbox/block/tracker.go |  76 +++++++++++++
 .../pkg/sandbox/block/tracker_test.go         |  93 ++++++++++++++++
 .../uffd/testutils/testharness/wire.go        |   2 +-
 .../sandbox/uffd/userfaultfd/helpers_test.go  |   3 +-
 .../uffd/userfaultfd/rpc_services_test.go     |  11 +-
 .../sandbox/uffd/userfaultfd/userfaultfd.go   |  27 +----
 8 files changed, 181 insertions(+), 212 deletions(-)
 delete mode 100644 packages/orchestrator/pkg/sandbox/block/state_tracker.go
 delete mode 100644 packages/orchestrator/pkg/sandbox/block/state_tracker_test.go
 create mode 100644 packages/orchestrator/pkg/sandbox/block/tracker.go
 create mode 100644 packages/orchestrator/pkg/sandbox/block/tracker_test.go

diff --git a/packages/orchestrator/pkg/sandbox/block/state_tracker.go b/packages/orchestrator/pkg/sandbox/block/state_tracker.go
deleted file mode 100644
index 3633f39cee..0000000000
--- a/packages/orchestrator/pkg/sandbox/block/state_tracker.go
+++ /dev/null
@@ -1,80 +0,0 @@
-package block
-
-import (
-	"fmt"
-	"sync"
-
-	"github.com/RoaringBitmap/roaring/v2"
-)
-
-type StateTracker[S comparable] struct {
-	mu sync.RWMutex
-
-	defaultState S
-	a, b         S
-	bmA, bmB     *roaring.Bitmap
-}
-
-// NewStateTracker requires three distinct states; duplicates would alias in
-// SetRange's switch and silently corrupt the bitmaps.
-func NewStateTracker[S comparable](defaultState, a, b S) (*StateTracker[S], error) {
-	if defaultState == a || defaultState == b || a == b {
-		return nil, fmt.Errorf("block.NewStateTracker: states must be distinct (default=%v a=%v b=%v)", defaultState, a, b)
-	}
-
-	return &StateTracker[S]{
-		defaultState: defaultState,
-		a:            a,
-		b:            b,
-		bmA:          roaring.New(),
-		bmB:          roaring.New(),
-	}, nil
-}
-
-// SetRange takes uint64 to allow end = 1<<32 (roaring's half-open upper bound).
-func (t *StateTracker[S]) SetRange(start, end uint64, state S) error {
-	if end <= start {
-		return nil
-	}
-
-	t.mu.Lock()
-	defer t.mu.Unlock()
-
-	switch state {
-	case t.a:
-		t.bmA.AddRange(start, end)
-		t.bmB.RemoveRange(start, end)
-	case t.b:
-		t.bmB.AddRange(start, end)
-		t.bmA.RemoveRange(start, end)
-	case t.defaultState:
-		t.bmA.RemoveRange(start, end)
-		t.bmB.RemoveRange(start, end)
-	default:
-		// S is only `comparable`, so the compiler can't prove exhaustiveness.
-		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) {
-	t.mu.RLock()
-	defer t.mu.RUnlock()
-
-	return t.bmA.Clone(), t.bmB.Clone()
-}
-
-func (t *StateTracker[S]) Get(idx uint32) S {
-	t.mu.RLock()
-	defer t.mu.RUnlock()
-
-	switch {
-	case t.bmA.Contains(idx):
-		return t.a
-	case t.bmB.Contains(idx):
-		return t.b
-	default:
-		return t.defaultState
-	}
-}
diff --git a/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go b/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go
deleted file mode 100644
index fe01cb9134..0000000000
--- a/packages/orchestrator/pkg/sandbox/block/state_tracker_test.go
+++ /dev/null
@@ -1,101 +0,0 @@
-package block
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-)
-
-type ts uint8
-
-const (
-	tsDefault ts = iota
-	tsA
-	tsB
-)
-
-func TestStateTracker(t *testing.T) {
-	t.Parallel()
-
-	t.Run("transitions", func(t *testing.T) {
-		t.Parallel()
-		s, err := NewStateTracker(tsDefault, tsA, tsB)
-		require.NoError(t, err)
-
-		require.NoError(t, s.SetRange(0, 1, tsA))
-		assert.Equal(t, tsA, s.Get(0))
-
-		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")
-
-		require.NoError(t, s.SetRange(0, 1, tsA))
-		assert.Equal(t, tsA, s.Get(0), "b→a should flip back")
-
-		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))
-
-		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, err := NewStateTracker(tsDefault, tsA, tsB)
-		require.NoError(t, err)
-
-		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)
-		}
-		for i := range uint32(4) {
-			page := i + 3
-			assert.Equal(t, tsB, s.Get(page), "page %d in overlap moves to b", page)
-		}
-		for i := range uint32(3) {
-			page := i + 7
-			assert.Equal(t, tsA, s.Get(page), "page %d outside overlap stays a", page)
-		}
-	})
-
-	t.Run("empty and inverted ranges are no-ops", func(t *testing.T) {
-		t.Parallel()
-		s, err := NewStateTracker(tsDefault, tsA, tsB)
-		require.NoError(t, err)
-
-		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("NewStateTracker errors on non-distinct states", func(t *testing.T) {
-		t.Parallel()
-		_, 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 errors on unsupported state", func(t *testing.T) {
-		t.Parallel()
-		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")
-	})
-}
diff --git a/packages/orchestrator/pkg/sandbox/block/tracker.go b/packages/orchestrator/pkg/sandbox/block/tracker.go
new file mode 100644
index 0000000000..92f3307c16
--- /dev/null
+++ b/packages/orchestrator/pkg/sandbox/block/tracker.go
@@ -0,0 +1,76 @@
+package block
+
+import (
+	"sync"
+
+	"github.com/RoaringBitmap/roaring/v2"
+)
+
+type State uint8
+
+const (
+	// NotPresent: fall back to the previous layer.
+	NotPresent State = iota
+	// Dirty: this layer holds materialized data.
+	Dirty
+	// Zero: known-zero; no need to consult the previous layer.
+	Zero
+)
+
+type Tracker struct {
+	mu          sync.RWMutex
+	dirty, zero *roaring.Bitmap
+}
+
+func NewTracker() *Tracker {
+	return &Tracker{
+		dirty: roaring.New(),
+		zero:  roaring.New(),
+	}
+}
+
+// SetRange takes uint64 to allow end = 1<<32 (roaring's half-open upper bound).
+// Out-of-range values (end > 1<<32) are silently ignored; roaring is a 32-bit
+// bitmap and AddRange panics otherwise.
+func (t *Tracker) SetRange(start, end uint64, state State) {
+	if end <= start || end > 1<<32 {
+		return
+	}
+
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	switch state {
+	case Dirty:
+		t.dirty.AddRange(start, end)
+		t.zero.RemoveRange(start, end)
+	case Zero:
+		t.zero.AddRange(start, end)
+		t.dirty.RemoveRange(start, end)
+	case NotPresent:
+		t.dirty.RemoveRange(start, end)
+		t.zero.RemoveRange(start, end)
+	}
+}
+
+func (t *Tracker) Get(idx uint32) State {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	switch {
+	case t.dirty.Contains(idx):
+		return Dirty
+	case t.zero.Contains(idx):
+		return Zero
+	default:
+		return NotPresent
+	}
+}
+
+// Export returns clones of the dirty and zero bitmaps.
+func (t *Tracker) Export() (dirty, zero *roaring.Bitmap) {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return t.dirty.Clone(), t.zero.Clone()
+}
diff --git a/packages/orchestrator/pkg/sandbox/block/tracker_test.go b/packages/orchestrator/pkg/sandbox/block/tracker_test.go
new file mode 100644
index 0000000000..5ac64fade5
--- /dev/null
+++ b/packages/orchestrator/pkg/sandbox/block/tracker_test.go
@@ -0,0 +1,93 @@
+package block
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTracker(t *testing.T) {
+	t.Parallel()
+
+	t.Run("transitions", func(t *testing.T) {
+		t.Parallel()
+		s := NewTracker()
+
+		s.SetRange(0, 1, Dirty)
+		assert.Equal(t, Dirty, s.Get(0))
+
+		s.SetRange(0, 1, Zero)
+		assert.Equal(t, Zero, s.Get(0), "dirty→zero should flip the page")
+		bmDirty, bmZero := s.Export()
+		assert.False(t, bmDirty.Contains(0), "dirty→zero must clear dirty bitmap")
+		assert.True(t, bmZero.Contains(0), "dirty→zero must add to zero bitmap")
+
+		s.SetRange(0, 1, Dirty)
+		assert.Equal(t, Dirty, s.Get(0), "zero→dirty should flip back")
+
+		s.SetRange(0, 1, NotPresent)
+		assert.Equal(t, NotPresent, s.Get(0), "→not-present must clear")
+		bmDirty, bmZero = s.Export()
+		assert.False(t, bmDirty.Contains(0))
+		assert.False(t, bmZero.Contains(0))
+
+		s.SetRange(0, 1, Dirty)
+		s.SetRange(0, 1, Dirty)
+		assert.Equal(t, Dirty, s.Get(0), "dirty→dirty is idempotent")
+	})
+
+	t.Run("partial overlap moves only the overlapping pages", func(t *testing.T) {
+		t.Parallel()
+		s := NewTracker()
+
+		s.SetRange(0, 10, Dirty)
+		s.SetRange(3, 7, Zero)
+
+		for i := range uint32(3) {
+			assert.Equal(t, Dirty, s.Get(i), "page %d outside overlap stays dirty", i)
+		}
+		for i := range uint32(4) {
+			page := i + 3
+			assert.Equal(t, Zero, s.Get(page), "page %d in overlap moves to zero", page)
+		}
+		for i := range uint32(3) {
+			page := i + 7
+			assert.Equal(t, Dirty, s.Get(page), "page %d outside overlap stays dirty", page)
+		}
+	})
+
+	t.Run("empty and inverted ranges are no-ops", func(t *testing.T) {
+		t.Parallel()
+		s := NewTracker()
+
+		s.SetRange(5, 5, Dirty)
+		s.SetRange(7, 3, Zero)
+		bmDirty, bmZero := s.Export()
+		assert.True(t, bmDirty.IsEmpty())
+		assert.True(t, bmZero.IsEmpty())
+	})
+
+	t.Run("out-of-range bounds are no-ops", func(t *testing.T) {
+		t.Parallel()
+		s := NewTracker()
+
+		s.SetRange(0, 1<<33, Dirty)
+		s.SetRange(1<<32+1, 1<<32+5, Zero)
+		bmDirty, bmZero := s.Export()
+		assert.True(t, bmDirty.IsEmpty(), "end > 1<<32 must be ignored")
+		assert.True(t, bmZero.IsEmpty(), "start >= 1<<32 must be ignored")
+
+		s.SetRange(1<<32-1, 1<<32, Dirty)
+		assert.Equal(t, Dirty, s.Get(1<<32-1), "end == 1<<32 still works")
+	})
+
+	t.Run("Export returns clones", func(t *testing.T) {
+		t.Parallel()
+		s := NewTracker()
+
+		s.SetRange(0, 1, Dirty)
+		bmDirty, _ := s.Export()
+		bmDirty.Add(42)
+		assert.Equal(t, NotPresent, s.Get(42), "mutating export must not leak into tracker")
+	})
+}
diff --git a/packages/orchestrator/pkg/sandbox/uffd/testutils/testharness/wire.go b/packages/orchestrator/pkg/sandbox/uffd/testutils/testharness/wire.go
index 7a475b775d..7cc7111955 100644
--- a/packages/orchestrator/pkg/sandbox/uffd/testutils/testharness/wire.go
+++ b/packages/orchestrator/pkg/sandbox/uffd/testutils/testharness/wire.go
@@ -19,7 +19,7 @@ type BootstrapArgs struct {
 
 type BootstrapReply struct{}
 
-// PageStateEntry is the wire form of the parent package's pageState enum.
+// PageStateEntry is the wire form of a block.State for a single page offset.
 type PageStateEntry struct {
 	State  uint8
 	Offset uint64
diff --git a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/helpers_test.go b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/helpers_test.go
index 443e5cef83..045c68f586 100644
--- a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/helpers_test.go
+++ b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/helpers_test.go
@@ -14,6 +14,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/testutils"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/uffd/testutils/testharness"
 )
@@ -76,7 +77,7 @@ func (h *testHandler) pageStates() (handlerPageStates, error) {
 
 	var states handlerPageStates
 	for _, e := range entries {
-		if pageState(e.State) == faulted {
+		if block.State(e.State) == block.Dirty {
 			states.faulted = append(states.faulted, uint(e.Offset))
 		}
 	}
diff --git a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/rpc_services_test.go b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/rpc_services_test.go
index b2a3562acf..d6a01b8f2e 100644
--- a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/rpc_services_test.go
+++ b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/rpc_services_test.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/RoaringBitmap/roaring/v2"
 
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block"
 	"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"
@@ -194,9 +195,9 @@ func (u *Userfaultfd) pageStateEntries() ([]testharness.PageStateEntry, error) {
 	u.settleRequests.Lock()
 	defer u.settleRequests.Unlock()
 
-	bmFaulted, bmRemoved := u.pageTracker.Export()
-	entries := make([]testharness.PageStateEntry, 0, bmFaulted.GetCardinality()+bmRemoved.GetCardinality())
-	emit := func(bm *roaring.Bitmap, state pageState) {
+	bmDirty, bmZero := u.pageTracker.Export()
+	entries := make([]testharness.PageStateEntry, 0, bmDirty.GetCardinality()+bmZero.GetCardinality())
+	emit := func(bm *roaring.Bitmap, state block.State) {
 		for _, idx := range bm.ToArray() {
 			entries = append(entries, testharness.PageStateEntry{
 				State:  uint8(state),
@@ -204,8 +205,8 @@ func (u *Userfaultfd) pageStateEntries() ([]testharness.PageStateEntry, error) {
 			})
 		}
 	}
-	emit(bmFaulted, faulted)
-	emit(bmRemoved, removed)
+	emit(bmDirty, block.Dirty)
+	emit(bmZero, block.Zero)
 
 	return entries, nil
 }
diff --git a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
index 1b9bb69695..2250c3bf6d 100644
--- a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
+++ b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
@@ -47,23 +47,13 @@ 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 *block.StateTracker[pageState]
+	pageTracker *block.Tracker
 
 	// 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
@@ -99,16 +89,11 @@ 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:     pageTracker,
+		pageTracker:     block.NewTracker(),
 		prefetchTracker: block.NewPrefetchTracker(blockSize),
 		ma:              m,
 		logger:          logger,
@@ -435,13 +420,7 @@ retryLoop:
 	}
 
 	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.pageTracker.SetRange(idx, idx+1, block.Dirty)
 	u.prefetchTracker.Add(offset, accessType)
 
 	return nil

From dcb7de44e1145a74f1e5fc663a33691a5658ff9b Mon Sep 17 00:00:00 2001
From: ValentaTomas <valenta.and.thomas@gmail.com>
Date: Mon, 4 May 2026 11:43:16 -0700
Subject: [PATCH 5/5] refactor(block): SetRange takes uint32 bounds

---
 .../orchestrator/pkg/sandbox/block/tracker.go | 22 +++++++++----------
 .../pkg/sandbox/block/tracker_test.go         | 14 ------------
 .../sandbox/uffd/userfaultfd/userfaultfd.go   |  2 +-
 3 files changed, 12 insertions(+), 26 deletions(-)

diff --git a/packages/orchestrator/pkg/sandbox/block/tracker.go b/packages/orchestrator/pkg/sandbox/block/tracker.go
index 92f3307c16..18b750263b 100644
--- a/packages/orchestrator/pkg/sandbox/block/tracker.go
+++ b/packages/orchestrator/pkg/sandbox/block/tracker.go
@@ -29,27 +29,27 @@ func NewTracker() *Tracker {
 	}
 }
 
-// SetRange takes uint64 to allow end = 1<<32 (roaring's half-open upper bound).
-// Out-of-range values (end > 1<<32) are silently ignored; roaring is a 32-bit
-// bitmap and AddRange panics otherwise.
-func (t *Tracker) SetRange(start, end uint64, state State) {
-	if end <= start || end > 1<<32 {
+// SetRange sets state for indices in [start, end). The index math.MaxUint32
+// is unaddressable: end is the half-open upper bound and capped at MaxUint32.
+func (t *Tracker) SetRange(start, end uint32, state State) {
+	if end <= start {
 		return
 	}
 
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
+	s, e := uint64(start), uint64(end)
 	switch state {
 	case Dirty:
-		t.dirty.AddRange(start, end)
-		t.zero.RemoveRange(start, end)
+		t.dirty.AddRange(s, e)
+		t.zero.RemoveRange(s, e)
 	case Zero:
-		t.zero.AddRange(start, end)
-		t.dirty.RemoveRange(start, end)
+		t.zero.AddRange(s, e)
+		t.dirty.RemoveRange(s, e)
 	case NotPresent:
-		t.dirty.RemoveRange(start, end)
-		t.zero.RemoveRange(start, end)
+		t.dirty.RemoveRange(s, e)
+		t.zero.RemoveRange(s, e)
 	}
 }
 
diff --git a/packages/orchestrator/pkg/sandbox/block/tracker_test.go b/packages/orchestrator/pkg/sandbox/block/tracker_test.go
index 5ac64fade5..a681b7e4e8 100644
--- a/packages/orchestrator/pkg/sandbox/block/tracker_test.go
+++ b/packages/orchestrator/pkg/sandbox/block/tracker_test.go
@@ -67,20 +67,6 @@ func TestTracker(t *testing.T) {
 		assert.True(t, bmZero.IsEmpty())
 	})
 
-	t.Run("out-of-range bounds are no-ops", func(t *testing.T) {
-		t.Parallel()
-		s := NewTracker()
-
-		s.SetRange(0, 1<<33, Dirty)
-		s.SetRange(1<<32+1, 1<<32+5, Zero)
-		bmDirty, bmZero := s.Export()
-		assert.True(t, bmDirty.IsEmpty(), "end > 1<<32 must be ignored")
-		assert.True(t, bmZero.IsEmpty(), "start >= 1<<32 must be ignored")
-
-		s.SetRange(1<<32-1, 1<<32, Dirty)
-		assert.Equal(t, Dirty, s.Get(1<<32-1), "end == 1<<32 still works")
-	})
-
 	t.Run("Export returns clones", func(t *testing.T) {
 		t.Parallel()
 		s := NewTracker()
diff --git a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
index 2250c3bf6d..ec432d51e5 100644
--- a/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
+++ b/packages/orchestrator/pkg/sandbox/uffd/userfaultfd/userfaultfd.go
@@ -419,7 +419,7 @@ retryLoop:
 		return fmt.Errorf("failed uffdio copy: %w", joinedErr)
 	}
 
-	idx := uint64(header.BlockIdx(offset, int64(u.pageSize)))
+	idx := uint32(header.BlockIdx(offset, int64(u.pageSize)))
 	u.pageTracker.SetRange(idx, idx+1, block.Dirty)
 	u.prefetchTracker.Add(offset, accessType)