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

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

@@ -8,7 +8,6 @@ import (
 	"math"
 	"math/rand"
 	"os"
-	"slices"
 	"sync"
 	"sync/atomic"
 	"syscall"
@@ -19,6 +18,7 @@ import (
 	"go.opentelemetry.io/otel"
 	"golang.org/x/sys/unix"
 
+	"github.com/e2b-dev/infra/packages/shared/pkg/atomicbitset"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
@@ -49,7 +49,7 @@ type Cache struct {
 	blockSize int64
 	mmap      *mmap.MMap
 	mu        sync.RWMutex
-	dirty     sync.Map
+	dirty     atomicbitset.Bitset
 	dirtyFile bool
 	closed    atomic.Bool
 }
@@ -87,12 +87,15 @@ func NewCache(size, blockSize int64, filePath string, dirtyFile bool) (*Cache, e
 		return nil, fmt.Errorf("error mapping file: %w", err)
 	}
 
+	numBlocks := (size + blockSize - 1) / blockSize
+
 	return &Cache{
 		mmap:      &mm,
 		filePath:  filePath,
 		size:      size,
 		blockSize: blockSize,
 		dirtyFile: dirtyFile,
+		dirty:     atomicbitset.New(uint(numBlocks)),
 	}, nil
 }
 
@@ -245,31 +248,30 @@ func (c *Cache) Slice(off, length int64) ([]byte, error) {
 	return nil, BytesNotAvailableError{}
 }
 
+func (c *Cache) startBlock(off int64) uint {
+	return uint(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 {
 		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)
-	// 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
-		}
-	}
-
-	return true
+	return c.dirty.HasRange(c.startBlock(off), c.endBlock(end))
 }
 
 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.
@@ -291,16 +293,13 @@ func (c *Cache) WriteAtWithoutLock(b []byte, off int64) (int, error) {
 	return n, nil
 }
 
-// dirtySortedKeys returns a sorted list of dirty keys.
-// Key represents a block offset.
+// dirtySortedKeys returns a sorted list of dirty block offsets.
 func (c *Cache) dirtySortedKeys() []int64 {
 	var keys []int64
-	c.dirty.Range(func(key, _ any) bool {
-		keys = append(keys, key.(int64))
 
-		return true
-	})
-	slices.Sort(keys)
+	for i := range c.dirty.Iterator() {
+		keys = append(keys, int64(i)*c.blockSize)
+	}
 
 	return keys
 }
@@ -491,9 +490,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/shared/pkg/atomicbitset/bitset.go

@@ -0,0 +1,115 @@
+// Package atomicbitset provides a fixed-size bitset with atomic set operations.
+package atomicbitset
+
+import (
+	"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
+}
+
+// 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
+}
+
+// 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
+	}
+}
+
+// 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
+			}
+		}
+	}
+}