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

Co-authored-by: ValentaTomas <valenta.and.thomas@gmail.com>

Author: levb

packages/api/go.mod

@@ -101,7 +101,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/gopkg v0.1.4 // indirect
 	github.com/bytedance/sonic v1.15.0 // indirect

packages/api/go.sum

@@ -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.1 // 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
@@ -254,6 +255,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/go.mod

@@ -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"
@@ -47,26 +48,17 @@ func NewErrCacheClosed(filePath string) *CacheClosedError {
 }
 
 type Cache struct {
-	filePath         string
-	size             int64
-	blockSize        int64
-	mmap             *mmap.MMap
-	mu               sync.RWMutex
-	dirty            sync.Map
-	dirtyGranularity int64
-	dirtyFile        bool
-	closed           atomic.Bool
+	filePath  string
+	size      int64
+	blockSize int64
+	mmap      *mmap.MMap
+	mu        sync.RWMutex
+	dirty     *atomicbitset.Bitset
+	dirtyFile bool
+	closed    atomic.Bool
 }
 
-// NewCache creates a cache with dirty tracking at blockSize granularity.
-// 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 NewCacheWithDirtyGranularity(size, blockSize, blockSize, filePath, dirtyFile)
-}
-
-// NewCacheWithDirtyGranularity creates a cache with dirty tracking at the specified granularity.
-// For chunker caches, dirtyGranularity can be larger than blockSize to reduce dirty map overhead.
-func NewCacheWithDirtyGranularity(size, blockSize, dirtyGranularity int64, filePath string, dirtyFile bool) (*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)
@@ -76,11 +68,11 @@ func NewCacheWithDirtyGranularity(size, blockSize, dirtyGranularity int64, fileP
 
 	if size == 0 {
 		return &Cache{
-			filePath:         filePath,
-			size:             size,
-			blockSize:        blockSize,
-			dirtyGranularity: dirtyGranularity,
-			dirtyFile:        dirtyFile,
+			filePath:  filePath,
+			size:      size,
+			blockSize: blockSize,
+			dirtyFile: dirtyFile,
+			dirty:     atomicbitset.New(),
 		}, nil
 	}
 
@@ -100,12 +92,12 @@ func NewCacheWithDirtyGranularity(size, blockSize, dirtyGranularity int64, fileP
 	}
 
 	return &Cache{
-		mmap:             &mm,
-		filePath:         filePath,
-		size:             size,
-		blockSize:        blockSize,
-		dirtyGranularity: dirtyGranularity,
-		dirtyFile:        dirtyFile,
+		mmap:      &mm,
+		filePath:  filePath,
+		size:      size,
+		blockSize: blockSize,
+		dirtyFile: dirtyFile,
+		dirty:     atomicbitset.New(),
 	}, nil
 }
 
@@ -117,10 +109,6 @@ func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMet
 	ctx, childSpan := tracer.Start(ctx, "export-to-diff")
 	defer childSpan.End()
 
-	if c.dirtyGranularity != c.blockSize {
-		return nil, fmt.Errorf("ExportToDiff requires block-level dirty tracking (granularity %d != blockSize %d)", c.dirtyGranularity, c.blockSize)
-	}
-
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
@@ -129,11 +117,7 @@ func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMet
 	}
 
 	if c.mmap == nil {
-		return &header.DiffMetadata{
-			Dirty:     bitset.New(0),
-			Empty:     bitset.New(0),
-			BlockSize: c.blockSize,
-		}, nil
+		return header.NewDiffMetadata(c.blockSize, bitset.New(0)), nil
 	}
 
 	f, err := os.Open(c.filePath)
@@ -152,18 +136,7 @@ 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))
-
-		return true
-	})
-
-	diffMetadata := builder.Build()
-	telemetry.SetAttributes(ctx, attribute.Int64("build_metadata_ms", time.Since(buildStart).Milliseconds()))
+	diffMetadata := header.NewDiffMetadata(c.blockSize, c.dirty.BitSet())
 
 	dst := int(out.Fd())
 	var writeOffset int64
@@ -336,45 +309,17 @@ func (c *Cache) sliceDirect(off, length int64) ([]byte, error) {
 }
 
 func (c *Cache) isCached(off, length int64) bool {
-	// Zero-length is vacuously true (no-op)
-	if length <= 0 {
-		return true
-	}
-
-	// Make sure the offset is within the cache size
-	if off >= c.size {
-		return false
-	}
-
-	// 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)
+	start := uint64(header.BlockIdx(off, c.blockSize))
+	end := uint64(header.BlockCeilIdx(min(off+length, c.size), c.blockSize))
 
-	startKey := (off / c.dirtyGranularity) * c.dirtyGranularity
-	endKey := ((end - 1) / c.dirtyGranularity) * c.dirtyGranularity
-
-	for key := startKey; key <= endKey; key += c.dirtyGranularity {
-		if _, ok := c.dirty.Load(key); !ok {
-			return false
-		}
-	}
-
-	return true
+	return c.dirty.HasRange(start, end)
 }
 
 func (c *Cache) setIsCached(off, length int64) {
-	// Zero-length is a no-op
-	if length <= 0 {
-		return
-	}
-
-	end := off + length
+	start := uint64(header.BlockIdx(off, c.blockSize))
+	end := uint64(header.BlockCeilIdx(off+length, c.blockSize))
 
-	startKey := (off / c.dirtyGranularity) * c.dirtyGranularity
-	endKey := ((end - 1) / c.dirtyGranularity) * c.dirtyGranularity
-
-	for key := startKey; key <= endKey; key += c.dirtyGranularity {
-		c.dirty.Store(key, struct{}{})
-	}
+	c.dirty.SetRange(start, end)
 }
 
 // When using WriteAtWithoutLock you must ensure thread safety, ideally by only writing to the same block once and the exposing the slice.

packages/orchestrator/go.sum

@@ -22,11 +22,6 @@ import (
 )
 
 const (
-	// ChunkerDirtyGranularity is the dirty tracking resolution for chunker caches.
-	// Fetches happen in MemoryChunkSize (4MB) chunks, so tracking at blockSize (4KB)
-	// is wasteful. 2MB gives 2 entries per fetch instead of 1024.
-	ChunkerDirtyGranularity = 2 * 1024 * 1024 // 2 MiB
-
 	pullType       = "pull-type"
 	pullTypeLocal  = "local"
 	pullTypeRemote = "remote"
@@ -143,7 +138,7 @@ func NewFullFetchChunker(
 	cachePath string,
 	metrics metrics.Metrics,
 ) (*FullFetchChunker, error) {
-	cache, err := NewCacheWithDirtyGranularity(size, blockSize, ChunkerDirtyGranularity, cachePath, false)
+	cache, err := NewCache(size, blockSize, cachePath, false)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create file cache: %w", err)
 	}

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

@@ -172,7 +172,7 @@ func NewStreamingChunker(
 	minReadBatchSize int64,
 	ff *featureflags.Client,
 ) (*StreamingChunker, error) {
-	cache, err := NewCacheWithDirtyGranularity(size, blockSize, ChunkerDirtyGranularity, cachePath, false)
+	cache, err := NewCache(size, blockSize, cachePath, false)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create file cache: %w", err)
 	}
@@ -270,9 +270,7 @@ func (c *StreamingChunker) Slice(ctx context.Context, off, length int64) ([]byte
 		return nil, fmt.Errorf("failed to ensure data at %d-%d: %w", off, off+length, err)
 	}
 
-	// Use sliceDirect (no isCached check) since the waiter mechanism guarantees data is in the mmap.
-	// With coarse dirty granularity, isCached may return false during an active fetch even though
-	// the requested bytes have been written to the mmap.
+	// sliceDirect skips isCached — the waiter already confirmed the data is in the mmap.
 	b, cacheErr := c.cache.sliceDirect(off, length)
 	if cacheErr != nil {
 		timer.RecordRaw(ctx, length, chunkerAttrs.failLocalReadAgain)

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

@@ -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.1
 	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.12.0
@@ -270,6 +271,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/orchestrator/pkg/sandbox/block/streaming_chunk.go

@@ -0,0 +1,40 @@
+package atomicbitset
+
+import (
+	"sync"
+
+	roaring "github.com/RoaringBitmap/roaring/v2"
+	"github.com/bits-and-blooms/bitset"
+)
+
+type Bitset struct {
+	mu sync.RWMutex
+	bm *roaring.Bitmap
+}
+
+func New() *Bitset {
+	return &Bitset{
+		bm: roaring.New(),
+	}
+}
+
+func (b *Bitset) HasRange(start, end uint64) bool {
+	b.mu.RLock()
+	defer b.mu.RUnlock()
+
+	return b.bm.CardinalityInRange(start, end) == end-start
+}
+
+func (b *Bitset) SetRange(start, end uint64) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	b.bm.AddRange(start, end)
+}
+
+func (b *Bitset) BitSet() *bitset.BitSet {
+	b.mu.RLock()
+	defer b.mu.RUnlock()
+
+	return b.bm.ToBitSet()
+}