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.

Author: ValentaTomas

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
+	}
+}

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")
+	})
+}