perf(block): sync.Map -> atomic bitmap for cache dirty tracking

packages/api/go.mod

@@ -102,7 +102,7 @@ require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bep/godartsass/v2 v2.3.2 // indirect
 	github.com/bep/golibsass v1.2.0 // indirect
-	github.com/bits-and-blooms/bitset v1.22.0 // indirect
+	github.com/bits-and-blooms/bitset v1.24.2 // indirect
 	github.com/bits-and-blooms/bloom/v3 v3.7.0 // indirect
 	github.com/bytedance/sonic v1.13.3 // indirect
 	github.com/bytedance/sonic/loader v0.2.4 // indirect

packages/orchestrator/Makefile

@@ -46,7 +46,7 @@ build-local: fetch-busybox
 
 .PHONY: build-debug
 build-debug: fetch-busybox
-	CGO_ENABLED=1 GOOS=linux GOARCH=$(BUILD_ARCH) go build -race -gcflags=all="-N -l" -o bin/orchestrator .
+	CGO_ENABLED=1 GOOS=linux GOARCH=$(BUILD_ARCH) go build -gcflags=all="-N -l" -o bin/orchestrator .
 
 .PHONY: run-debug
 run-debug:

packages/orchestrator/go.mod

@@ -20,7 +20,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/config v1.32.6
 	github.com/aws/aws-sdk-go-v2/credentials v1.19.6
 	github.com/aws/aws-sdk-go-v2/service/ecr v1.44.0
-	github.com/bits-and-blooms/bitset v1.22.0
+	github.com/bits-and-blooms/bitset v1.24.2
 	github.com/bmatcuk/doublestar/v4 v4.9.1
 	github.com/caarlos0/env/v11 v11.3.1
 	github.com/containernetworking/plugins v1.9.0
@@ -97,6 +97,7 @@ require (
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.54.0 // indirect
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.54.0 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
+	github.com/RoaringBitmap/roaring/v2 v2.16.0 // indirect
 	github.com/andybalholm/brotli v1.2.0 // indirect
 	github.com/armon/go-metrics v0.4.1 // indirect
 	github.com/aws/aws-sdk-go-v2 v1.41.0 // indirect
@@ -252,6 +253,7 @@ require (
 	github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect
 	github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect
 	github.com/morikuni/aec v1.0.0 // indirect
+	github.com/mschoch/smat v0.2.0 // indirect
 	github.com/oapi-codegen/oapi-codegen/v2 v2.5.1 // indirect
 	github.com/oasdiff/yaml v0.0.0-20250309154309-f31be36b4037 // indirect
 	github.com/oasdiff/yaml3 v0.0.0-20250309153720-d2182401db90 // indirect

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

@@ -20,6 +20,7 @@ import (
 	"go.uber.org/zap"
 	"golang.org/x/sys/unix"
 
+	"github.com/e2b-dev/infra/packages/shared/pkg/atomicbitset"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
@@ -52,13 +53,16 @@ type Cache struct {
 	blockSize int64
 	mmap      *mmap.MMap
 	mu        sync.RWMutex
-	dirty     sync.Map
+	dirty     atomicbitset.Bitset
 	dirtyFile bool
 	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)
@@ -96,6 +100,7 @@ func NewCache(size, blockSize int64, filePath string, dirtyFile bool) (*Cache, e
 		size:      size,
 		blockSize: blockSize,
 		dirtyFile: dirtyFile,
+		dirty:     atomicbitset.New(uint(header.TotalBlocks(size, blockSize)), bitsetImpl),
 	}, nil
 }
 
@@ -138,15 +143,10 @@ func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMet
 		logger.L().Warn(ctx, "error syncing file", zap.Error(err))
 	}
 
-	buildStart := time.Now()
-	builder := header.NewDiffMetadataBuilder(c.size, c.blockSize)
-
-	// We don't need to sort the keys as the bitset handles the ordering.
-	c.dirty.Range(func(key, _ any) bool {
-		builder.AddDirtyOffset(key.(int64))
+	dirty := c.dirty.BitSet()
 
-		return true
-	})
+	buildStart := time.Now()
+	builder := header.NewDiffMetadataBuilderFromDirtyBitSet(c.size, c.blockSize, dirty)
 
 	diffMetadata := builder.Build()
 	telemetry.SetAttributes(ctx, attribute.Int64("build_metadata_ms", time.Since(buildStart).Milliseconds()))
@@ -301,6 +301,14 @@ func (c *Cache) Slice(off, length int64) ([]byte, error) {
 	return nil, BytesNotAvailableError{}
 }
 
+func (c *Cache) startBlock(off int64) uint {
+	return uint(header.BlockIdx(off, c.blockSize))
+}
+
+func (c *Cache) endBlock(off int64) uint {
+	return uint((off + c.blockSize - 1) / c.blockSize)
+}
+
 func (c *Cache) isCached(off, length int64) bool {
 	// Make sure the offset is within the cache size
 	if off >= c.size {
@@ -309,23 +317,24 @@ func (c *Cache) isCached(off, length int64) bool {
 
 	// Cap if the length goes beyond the cache size, so we don't check for blocks that are out of bounds.
 	end := min(off+length, c.size)
-	// Recalculate the length based on the capped end, so we check for the correct blocks in case of capping.
-	length = end - off
 
-	for _, blockOff := range header.BlocksOffsets(length, c.blockSize) {
-		_, dirty := c.dirty.Load(off + blockOff)
-		if !dirty {
-			return false
-		}
+	lo := c.startBlock(off)
+	hi := c.endBlock(end)
+
+	// Fast path: single-block check (common case for NBD reads).
+	if hi-lo == 1 {
+		return c.dirty.Has(lo)
 	}
 
-	return true
+	return c.dirty.HasRange(lo, hi)
 }
 
 func (c *Cache) setIsCached(off, length int64) {
-	for _, blockOff := range header.BlocksOffsets(length, c.blockSize) {
-		c.dirty.Store(off+blockOff, struct{}{})
+	if length <= 0 {
+		return
 	}
+
+	c.dirty.SetRange(c.startBlock(off), c.endBlock(off+length))
 }
 
 // When using WriteAtWithoutLock you must ensure thread safety, ideally by only writing to the same block once and the exposing the slice.
@@ -533,9 +542,7 @@ func (c *Cache) copyProcessMemory(
 				return fmt.Errorf("failed to read memory: expected %d bytes, got %d", segmentSize, n)
 			}
 
-			for _, blockOff := range header.BlocksOffsets(segmentSize, c.blockSize) {
-				c.dirty.Store(offset+blockOff, struct{}{})
-			}
+			c.setIsCached(offset, segmentSize)
 
 			offset += segmentSize
 

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

@@ -98,17 +98,19 @@ func NewChunker(
 	cachePath string,
 	metrics metrics.Metrics,
 ) (Chunker, error) {
-	useStreaming, minReadBatchSizeKB := getChunkerConfig(ctx, featureFlags)
+	useStreaming, minReadBatchSizeKB, bitsetImpl := getChunkerConfig(ctx, featureFlags)
+
+	fmt.Printf("[DEBUG block.NewChunker] bitsetImpl=%q useStreaming=%v size=%d blockSize=%d cachePath=%s\n", bitsetImpl, useStreaming, size, blockSize, cachePath)
 
 	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 +121,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 +144,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

@@ -8,12 +8,13 @@ require (
 	cloud.google.com/go/artifactregistry v1.17.1
 	cloud.google.com/go/storage v1.59.2
 	connectrpc.com/connect v1.18.1
+	github.com/RoaringBitmap/roaring/v2 v2.16.0
 	github.com/aws/aws-sdk-go-v2 v1.41.0
 	github.com/aws/aws-sdk-go-v2/config v1.32.6
 	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
@@ -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