feat(block): add flat/roaring/sharded bitset implementations behind FF

Introduce three atomicbitset implementations behind a shared Bitset
interface, selectable at runtime via the chunker-config feature flag
("bitset" key): roaring (default), and atomic (flat/sharded auto).

- Flat: []atomic.Uint64, lock-free, for small bitmaps (≤64KB)
- Sharded: two-level with lazy shard allocation, lock-free, for large sparse bitmaps
- Roaring: compressed bitmap with RWMutex, default via FF
- Iterator/UnsafeIterator split for safe vs caller-synchronized iteration
- Conformance tests across all implementations with race detector

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: levb

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

@@ -54,8 +54,11 @@ type Cache struct {
 	closed    atomic.Bool
 }
 
-// When we are passing filePath that is a file that has content we want to server want to use dirtyFile = true.
 func NewCache(size, blockSize int64, filePath string, dirtyFile bool) (*Cache, error) {
+	return newCache(size, blockSize, filePath, dirtyFile, "")
+}
+
+func newCache(size, blockSize int64, filePath string, dirtyFile bool, bitsetImpl string) (*Cache, error) {
 	f, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE, 0o644)
 	if err != nil {
 		return nil, fmt.Errorf("error opening file: %w", err)
@@ -95,7 +98,7 @@ func NewCache(size, blockSize int64, filePath string, dirtyFile bool) (*Cache, e
 		size:      size,
 		blockSize: blockSize,
 		dirtyFile: dirtyFile,
-		dirty:     atomicbitset.New(uint(numBlocks)),
+		dirty:     atomicbitset.New(uint(numBlocks), bitsetImpl),
 	}, nil
 }
 
@@ -295,6 +298,10 @@ func (c *Cache) WriteAtWithoutLock(b []byte, off int64) (int, error) {
 
 // dirtySortedKeys returns a sorted list of dirty block offsets.
 func (c *Cache) dirtySortedKeys() []int64 {
+	if c.dirty == nil {
+		return nil
+	}
+
 	var keys []int64
 
 	for i := range c.dirty.Iterator() {

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

@@ -98,17 +98,17 @@ func NewChunker(
 	cachePath string,
 	metrics metrics.Metrics,
 ) (Chunker, error) {
-	useStreaming, minReadBatchSizeKB := getChunkerConfig(ctx, featureFlags)
+	useStreaming, minReadBatchSizeKB, bitsetImpl := getChunkerConfig(ctx, featureFlags)
 
 	if useStreaming {
-		return NewStreamingChunker(size, blockSize, upstream, cachePath, metrics, int64(minReadBatchSizeKB)*1024, featureFlags)
+		return newStreamingChunker(size, blockSize, upstream, cachePath, metrics, int64(minReadBatchSizeKB)*1024, featureFlags, bitsetImpl)
 	}
 
-	return NewFullFetchChunker(size, blockSize, upstream, cachePath, metrics)
+	return newFullFetchChunker(size, blockSize, upstream, cachePath, metrics, bitsetImpl)
 }
 
 // getChunkerConfig fetches the chunker-config feature flag and returns the parsed values.
-func getChunkerConfig(ctx context.Context, ff *featureflags.Client) (useStreaming bool, minReadBatchSizeKB int) {
+func getChunkerConfig(ctx context.Context, ff *featureflags.Client) (useStreaming bool, minReadBatchSizeKB int, bitsetImpl string) {
 	value := ff.JSONFlag(ctx, featureflags.ChunkerConfigFlag)
 
 	if v := value.GetByKey("useStreaming"); v.IsDefined() {
@@ -119,7 +119,11 @@ func getChunkerConfig(ctx context.Context, ff *featureflags.Client) (useStreamin
 		minReadBatchSizeKB = v.IntValue()
 	}
 
-	return useStreaming, minReadBatchSizeKB
+	if v := value.GetByKey("bitset"); v.IsDefined() {
+		bitsetImpl = v.StringValue()
+	}
+
+	return useStreaming, minReadBatchSizeKB, bitsetImpl
 }
 
 type FullFetchChunker struct {
@@ -138,7 +142,17 @@ func NewFullFetchChunker(
 	cachePath string,
 	metrics metrics.Metrics,
 ) (*FullFetchChunker, error) {
-	cache, err := NewCache(size, blockSize, cachePath, false)
+	return newFullFetchChunker(size, blockSize, base, cachePath, metrics, "")
+}
+
+func newFullFetchChunker(
+	size, blockSize int64,
+	base storage.SeekableReader,
+	cachePath string,
+	metrics metrics.Metrics,
+	bitsetImpl string,
+) (*FullFetchChunker, error) {
+	cache, err := newCache(size, blockSize, cachePath, false, bitsetImpl)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create file cache: %w", err)
 	}

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

@@ -172,7 +172,19 @@ func NewStreamingChunker(
 	minReadBatchSize int64,
 	ff *featureflags.Client,
 ) (*StreamingChunker, error) {
-	cache, err := NewCache(size, blockSize, cachePath, false)
+	return newStreamingChunker(size, blockSize, upstream, cachePath, metrics, minReadBatchSize, ff, "")
+}
+
+func newStreamingChunker(
+	size, blockSize int64,
+	upstream storage.StreamingReader,
+	cachePath string,
+	metrics metrics.Metrics,
+	minReadBatchSize int64,
+	ff *featureflags.Client,
+	bitsetImpl string,
+) (*StreamingChunker, error) {
+	cache, err := newCache(size, blockSize, cachePath, false, bitsetImpl)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create file cache: %w", err)
 	}
@@ -441,7 +453,7 @@ func (c *StreamingChunker) progressiveRead(ctx context.Context, s *fetchSession,
 // it can be tuned without restarting the service.
 func (c *StreamingChunker) getMinReadBatchSize(ctx context.Context) int64 {
 	if c.featureFlags != nil {
-		_, minKB := getChunkerConfig(ctx, c.featureFlags)
+		_, minKB, _ := getChunkerConfig(ctx, c.featureFlags)
 		if minKB > 0 {
 			return int64(minKB) * 1024
 		}

packages/shared/go.mod

@@ -13,7 +13,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.17.74
 	github.com/aws/aws-sdk-go-v2/service/ecr v1.44.0
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3
-	github.com/bits-and-blooms/bitset v1.22.0
+	github.com/bits-and-blooms/bitset v1.24.2
 	github.com/bsm/redislock v0.9.4
 	github.com/dchest/uniuri v1.2.0
 	github.com/gin-gonic/gin v1.10.1
@@ -87,6 +87,7 @@ require (
 	github.com/Masterminds/semver/v3 v3.4.0 // indirect
 	github.com/Masterminds/sprig/v3 v3.3.0 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
+	github.com/RoaringBitmap/roaring/v2 v2.16.0 // indirect
 	github.com/Workiva/go-datastructures v1.1.6 // indirect
 	github.com/alecthomas/units v0.0.0-20240927000941-0f3dac36c52b // indirect
 	github.com/apapsch/go-jsonmerge/v2 v2.0.0 // indirect
@@ -267,6 +268,7 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect
 	github.com/morikuni/aec v1.0.0 // indirect
+	github.com/mschoch/smat v0.2.0 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f // indirect
 	github.com/oklog/ulid v1.3.1 // indirect

packages/shared/go.sum

@@ -1,115 +1,40 @@
-// Package atomicbitset provides a fixed-size bitset with atomic set operations.
+// Package atomicbitset provides fixed-size bitset implementations.
+// All implementations are safe for concurrent HasRange and SetRange.
 package atomicbitset
 
 import (
+	"fmt"
 	"iter"
-	"math"
-	"math/bits"
-	"sync/atomic"
 )
 
-// Bitset is a fixed-size bitset backed by atomic uint64 words.
-// SetRange uses atomic OR, so concurrent writers are safe without
-// external locking.
-//
-// A Bitset must not be copied after first use (copies share the
-// underlying array).
-type Bitset struct {
-	words []atomic.Uint64
-	n     uint
+type Bitset interface {
+	Has(i uint) bool
+	HasRange(lo, hi uint) bool
+	SetRange(lo, hi uint)
+	Iterator() iter.Seq[uint]
+	UnsafeIterator() iter.Seq[uint]
+	Len() uint
 }
 
-// New returns a Bitset with capacity for n bits, all initially zero.
-func New(n uint) Bitset {
-	return Bitset{
-		words: make([]atomic.Uint64, (n+63)/64),
-		n:     n,
-	}
-}
-
-// Len returns the capacity in bits.
-func (b *Bitset) Len() uint { return b.n }
-
-// Has reports whether bit i is set. Out-of-range returns false.
-func (b *Bitset) Has(i uint) bool {
-	if i >= b.n {
-		return false
-	}
-
-	return b.words[i/64].Load()&(1<<(i%64)) != 0
-}
+const (
+	autoThreshold uint = 524_288 // 64 KB flat bitmap
 
-// wordMask returns a bitmask covering bits [lo, hi) within a single uint64 word.
-func wordMask(lo, hi uint) uint64 {
-	if hi-lo == 64 {
-		return math.MaxUint64
-	}
-
-	return ((1 << (hi - lo)) - 1) << lo
-}
-
-// HasRange reports whether every bit in [lo, hi) is set.
-// An empty range returns true. hi is capped to Len().
-// Returns false if lo is out of range and the range is non-empty.
-func (b *Bitset) HasRange(lo, hi uint) bool {
-	if lo >= hi {
-		return true
-	}
-	if hi > b.n {
-		hi = b.n
-	}
-	if lo >= hi {
-		return false
-	}
-	for i := lo; i < hi; {
-		w := i / 64
-		bit := i % 64
-		top := min(hi-w*64, 64)
-		mask := wordMask(bit, top)
-
-		if b.words[w].Load()&mask != mask {
-			return false
-		}
-		i = (w + 1) * 64
-	}
-
-	return true
-}
-
-// SetRange sets every bit in [lo, hi) using atomic OR.
-// hi is capped to Len().
-func (b *Bitset) SetRange(lo, hi uint) {
-	if hi > b.n {
-		hi = b.n
-	}
-	if lo >= hi {
-		return
-	}
-	for i := lo; i < hi; {
-		w := i / 64
-		bit := i % 64
-		top := min(hi-w*64, 64)
-
-		b.words[w].Or(wordMask(bit, top))
-
-		i = (w + 1) * 64
-	}
-}
+	// Valid impl values for New.
+	BitsetDefault = ""
+	BitsetRoaring = "roaring"
+	BitsetAtomic  = "atomic"
+)
 
-// Iterator returns an iterator over the indices of set bits
-// in ascending order.
-func (b *Bitset) Iterator() iter.Seq[uint] {
-	return func(yield func(uint) bool) {
-		for wi := range b.words {
-			word := b.words[wi].Load()
-			base := uint(wi) * 64
-			for word != 0 {
-				tz := uint(bits.TrailingZeros64(word))
-				if !yield(base + tz) {
-					return
-				}
-				word &= word - 1
-			}
+func New(n uint, impl string) Bitset {
+	switch impl {
+	case BitsetDefault, BitsetRoaring:
+		return NewRoaring(n)
+	case BitsetAtomic:
+		if n <= autoThreshold {
+			return NewFlat(n)
 		}
+		return NewSharded(n, DefaultShardBits)
+	default:
+		panic(fmt.Sprintf("atomicbitset: unknown implementation %q", impl))
 	}
 }