Replace bits-and-blooms/bitset with roaring bitmaps

packages/envd/internal/services/cgroups/cgroup2_test.go

@@ -39,7 +39,7 @@ func TestCgroupRoundTrip(t *testing.T) {
 		return
 	}
 
-	maxTimeout := time.Second * 5
+	maxTimeout := time.Second * 15
 
 	t.Run("process does not die without cgroups", func(t *testing.T) {
 		t.Parallel()

packages/orchestrator/go.mod

@@ -17,6 +17,7 @@ require (
 	cloud.google.com/go/storage v1.59.2
 	connectrpc.com/connect v1.18.1
 	github.com/Merovius/nbd v0.0.0-20240812113926-fd65a54c9949
+	github.com/RoaringBitmap/roaring/v2 v2.18.0
 	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
@@ -97,7 +98,6 @@ 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

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

@@ -13,7 +13,7 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/bits-and-blooms/bitset"
+	"github.com/RoaringBitmap/roaring/v2"
 	"github.com/edsrzf/mmap-go"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
@@ -117,7 +117,7 @@ func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMet
 	}
 
 	if c.mmap == nil {
-		return header.NewDiffMetadata(c.blockSize, bitset.New(0)), nil
+		return header.NewDiffMetadata(c.blockSize, roaring.New()), nil
 	}
 
 	f, err := os.Open(c.filePath)
@@ -136,7 +136,7 @@ func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMet
 		logger.L().Warn(ctx, "error syncing file", zap.Error(err))
 	}
 
-	diffMetadata := header.NewDiffMetadata(c.blockSize, c.dirty.BitSet())
+	diffMetadata := header.NewDiffMetadata(c.blockSize, c.dirty.Bitmap())
 
 	dst := int(out.Fd())
 	var writeOffset int64
@@ -194,7 +194,7 @@ func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMet
 	telemetry.SetAttributes(ctx,
 		attribute.Int64("copy_ms", time.Since(copyStart).Milliseconds()),
 		attribute.Int64("total_size_bytes", c.size),
-		attribute.Int64("dirty_size_bytes", int64(diffMetadata.Dirty.Count())*c.blockSize),
+		attribute.Int64("dirty_size_bytes", int64(diffMetadata.Dirty.GetCardinality())*c.blockSize),
 		attribute.Int64("total_ranges", totalRanges),
 	)
 

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

@@ -257,8 +257,8 @@ func TestCacheExportToDiff_ZeroDirtyBlockEmittedAsDirtyPayload(t *testing.T) {
 	diffMetadata, err := cache.ExportToDiff(t.Context(), out)
 	require.NoError(t, err)
 
-	require.EqualValues(t, 1, diffMetadata.Dirty.Count(), "zero-filled dirty block should be emitted as dirty payload")
-	require.EqualValues(t, 0, diffMetadata.Empty.Count(), "zero-filled dirty block should not be tracked in empty metadata")
+	require.EqualValues(t, 1, diffMetadata.Dirty.GetCardinality(), "zero-filled dirty block should be emitted as dirty payload")
+	require.EqualValues(t, 0, diffMetadata.Empty.GetCardinality(), "zero-filled dirty block should not be tracked in empty metadata")
 
 	stat, err := out.Stat()
 	require.NoError(t, err)
@@ -335,8 +335,8 @@ func TestCacheExportToDiff_MixedDirtyBlocksKeepsZeroBlockInDiff(t *testing.T) {
 	diffMetadata, err := cache.ExportToDiff(t.Context(), out)
 	require.NoError(t, err)
 
-	require.EqualValues(t, 2, diffMetadata.Dirty.Count())
-	require.EqualValues(t, 0, diffMetadata.Empty.Count(), "mixed export should still skip empty tracking for zero-filled dirty blocks")
+	require.EqualValues(t, 2, diffMetadata.Dirty.GetCardinality())
+	require.EqualValues(t, 0, diffMetadata.Empty.GetCardinality(), "mixed export should still skip empty tracking for zero-filled dirty blocks")
 
 	_, err = out.Seek(0, io.SeekStart)
 	require.NoError(t, err)
@@ -399,10 +399,10 @@ func TestCacheExportToDiff_NonContiguousDirtyBlocksPreserveRangeOrder(t *testing
 	diffMetadata, err := cache.ExportToDiff(t.Context(), out)
 	require.NoError(t, err)
 
-	require.EqualValues(t, 2, diffMetadata.Dirty.Count())
-	require.True(t, diffMetadata.Dirty.Test(0))
-	require.True(t, diffMetadata.Dirty.Test(3))
-	require.EqualValues(t, 0, diffMetadata.Empty.Count())
+	require.EqualValues(t, 2, diffMetadata.Dirty.GetCardinality())
+	require.True(t, diffMetadata.Dirty.Contains(0))
+	require.True(t, diffMetadata.Dirty.Contains(3))
+	require.EqualValues(t, 0, diffMetadata.Empty.GetCardinality())
 
 	_, err = out.Seek(0, io.SeekStart)
 	require.NoError(t, err)

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

@@ -3,7 +3,7 @@ package block
 import (
 	"iter"
 
-	"github.com/bits-and-blooms/bitset"
+	"github.com/RoaringBitmap/roaring/v2"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
@@ -36,24 +36,13 @@ func NewRangeFromBlocks(startIdx, numberOfBlocks, blockSize int64) Range {
 	}
 }
 
-// bitsetRanges returns a sequence of the ranges of the set bits of the bitset.
-func BitsetRanges(b *bitset.BitSet, blockSize int64) iter.Seq[Range] {
+// BitsetRanges returns a sequence of the ranges of the set bits of the bitmap.
+func BitsetRanges(b *roaring.Bitmap, blockSize int64) iter.Seq[Range] {
 	return func(yield func(Range) bool) {
-		start, found := b.NextSet(0)
-
-		for found {
-			end, endOk := b.NextClear(start)
-			if !endOk {
-				yield(NewRangeFromBlocks(int64(start), int64(b.Len()-start), blockSize))
-
-				return
-			}
-
-			if !yield(NewRangeFromBlocks(int64(start), int64(end-start), blockSize)) {
+		for start, endExcl := range b.Ranges() {
+			if !yield(NewRangeFromBlocks(int64(start), int64(endExcl)-int64(start), blockSize)) {
 				return
 			}
-
-			start, found = b.NextSet(end + 1)
 		}
 	}
 }

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

@@ -5,7 +5,7 @@ import (
 	"slices"
 	"testing"
 
-	"github.com/bits-and-blooms/bitset"
+	"github.com/RoaringBitmap/roaring/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -290,7 +290,7 @@ func TestRange_Offsets_Iteration(t *testing.T) {
 
 func TestBitsetRanges_Empty(t *testing.T) {
 	t.Parallel()
-	b := bitset.New(100)
+	b := roaring.New()
 	blockSize := int64(4096)
 
 	ranges := slices.Collect(BitsetRanges(b, blockSize))
@@ -299,8 +299,8 @@ func TestBitsetRanges_Empty(t *testing.T) {
 
 func TestBitsetRanges_SingleBit(t *testing.T) {
 	t.Parallel()
-	b := bitset.New(100)
-	b.Set(5)
+	b := roaring.New()
+	b.Add(5)
 	blockSize := int64(4096)
 
 	ranges := slices.Collect(BitsetRanges(b, blockSize))
@@ -313,12 +313,12 @@ func TestBitsetRanges_SingleBit(t *testing.T) {
 
 func TestBitsetRanges_Contiguous(t *testing.T) {
 	t.Parallel()
-	b := bitset.New(100)
+	b := roaring.New()
 	// Set bits 2, 3, 4, 5
-	b.Set(2)
-	b.Set(3)
-	b.Set(4)
-	b.Set(5)
+	b.Add(2)
+	b.Add(3)
+	b.Add(4)
+	b.Add(5)
 	blockSize := int64(4096)
 
 	ranges := slices.Collect(BitsetRanges(b, blockSize))
@@ -331,15 +331,15 @@ func TestBitsetRanges_Contiguous(t *testing.T) {
 
 func TestBitsetRanges_MultipleRanges(t *testing.T) {
 	t.Parallel()
-	b := bitset.New(100)
+	b := roaring.New()
 	// Set bits 1, 2, 3 (contiguous)
-	b.Set(1)
-	b.Set(2)
-	b.Set(3)
+	b.Add(1)
+	b.Add(2)
+	b.Add(3)
 	// Gap
 	// Set bits 7, 8 (contiguous)
-	b.Set(7)
-	b.Set(8)
+	b.Add(7)
+	b.Add(8)
 	blockSize := int64(4096)
 
 	ranges := slices.Collect(BitsetRanges(b, blockSize))
@@ -356,9 +356,9 @@ func TestBitsetRanges_MultipleRanges(t *testing.T) {
 
 func TestBitsetRanges_AllSet(t *testing.T) {
 	t.Parallel()
-	b := bitset.New(10)
-	for i := range uint(10) {
-		b.Set(i)
+	b := roaring.New()
+	for i := range uint32(10) {
+		b.Add(i)
 	}
 	blockSize := int64(4096)
 
@@ -372,10 +372,10 @@ func TestBitsetRanges_AllSet(t *testing.T) {
 
 func TestBitsetRanges_EndOfBitset(t *testing.T) {
 	t.Parallel()
-	b := bitset.New(20)
+	b := roaring.New()
 	// Set bits 15, 16, 17, 18, 19 (at the end)
-	for i := uint(15); i < 20; i++ {
-		b.Set(i)
+	for i := uint32(15); i < 20; i++ {
+		b.Add(i)
 	}
 	blockSize := int64(4096)
 
@@ -389,12 +389,12 @@ func TestBitsetRanges_EndOfBitset(t *testing.T) {
 
 func TestBitsetRanges_Sparse(t *testing.T) {
 	t.Parallel()
-	b := bitset.New(100)
+	b := roaring.New()
 	// Set individual bits with gaps
-	b.Set(0)
-	b.Set(10)
-	b.Set(20)
-	b.Set(30)
+	b.Add(0)
+	b.Add(10)
+	b.Add(20)
+	b.Add(30)
 	blockSize := int64(4096)
 
 	ranges := slices.Collect(BitsetRanges(b, blockSize))

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

@@ -4,23 +4,21 @@ import (
 	"iter"
 	"sync"
 
-	"github.com/bits-and-blooms/bitset"
+	"github.com/RoaringBitmap/roaring/v2"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
-	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
 type Tracker struct {
-	b  *bitset.BitSet
+	b  *roaring.Bitmap
 	mu sync.RWMutex
 
 	blockSize int64
 }
 
 func NewTracker(blockSize int64) *Tracker {
 	return &Tracker{
-		// The bitset resizes automatically based on the maximum set bit.
-		b:         bitset.New(0),
+		b:         roaring.New(),
 		blockSize: blockSize,
 	}
 }
@@ -29,27 +27,21 @@ func (t *Tracker) Has(off int64) bool {
 	t.mu.RLock()
 	defer t.mu.RUnlock()
 
-	return t.b.Test(uint(header.BlockIdx(off, t.blockSize)))
+	return t.b.Contains(uint32(header.BlockIdx(off, t.blockSize)))
 }
 
 func (t *Tracker) Add(off int64) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
-	t.b.Set(uint(header.BlockIdx(off, t.blockSize)))
+	t.b.Add(uint32(header.BlockIdx(off, t.blockSize)))
 }
 
 func (t *Tracker) Reset() {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
-	t.b.ClearAll()
-}
-
-// BitSet returns the bitset.
-// This is not safe to use concurrently.
-func (t *Tracker) BitSet() *bitset.BitSet {
-	return t.b
+	t.b.Clear()
 }
 
 func (t *Tracker) BlockSize() int64 {
@@ -70,11 +62,11 @@ func (t *Tracker) Offsets() iter.Seq[int64] {
 	t.mu.RLock()
 	defer t.mu.RUnlock()
 
-	return bitsetOffsets(t.b.Clone(), t.BlockSize())
-}
+	snapshot := t.b.Clone()
 
-func bitsetOffsets(b *bitset.BitSet, blockSize int64) iter.Seq[int64] {
-	return utils.TransformTo(b.EachSet(), func(idx uint) int64 {
-		return header.BlockOffset(int64(idx), blockSize)
-	})
+	return func(yield func(int64) bool) {
+		snapshot.Iterate(func(idx uint32) bool {
+			return yield(header.BlockOffset(int64(idx), t.blockSize))
+		})
+	}
 }

packages/orchestrator/pkg/sandbox/fc/client.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"runtime"
 
-	"github.com/bits-and-blooms/bitset"
+	"github.com/RoaringBitmap/roaring/v2"
 	"github.com/firecracker-microvm/firecracker-go-sdk"
 	"github.com/go-openapi/strfmt"
 
@@ -451,8 +451,8 @@ func (c *apiClient) memoryInfo(ctx context.Context, blockSize int64) (*header.Di
 	}
 
 	return &header.DiffMetadata{
-		Dirty:     bitset.From(res.Payload.Resident),
-		Empty:     bitset.From(res.Payload.Empty),
+		Dirty:     roaring.FromDense(res.Payload.Resident, false),
+		Empty:     roaring.FromDense(res.Payload.Empty, false),
 		BlockSize: blockSize,
 	}, nil
 }
@@ -468,8 +468,8 @@ func (c *apiClient) dirtyMemory(ctx context.Context, blockSize int64) (*header.D
 	}
 
 	return &header.DiffMetadata{
-		Dirty:     bitset.From(res.Payload.Bitmap),
-		Empty:     bitset.New(0),
+		Dirty:     roaring.FromDense(res.Payload.Bitmap, false),
+		Empty:     roaring.New(),
 		BlockSize: blockSize,
 	}, nil
 }

packages/orchestrator/pkg/sandbox/fc/memory.go

@@ -4,7 +4,7 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/bits-and-blooms/bitset"
+	"github.com/RoaringBitmap/roaring/v2"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
@@ -23,7 +23,7 @@ func (p *Process) DirtyMemory(ctx context.Context, blockSize int64) (*header.Dif
 
 func (p *Process) ExportMemory(
 	ctx context.Context,
-	include *bitset.BitSet,
+	include *roaring.Bitmap,
 	cachePath string,
 	blockSize int64,
 ) (*block.Cache, error) {

packages/orchestrator/pkg/sandbox/rootfs/direct.go

@@ -131,7 +131,7 @@ func (o *DirectProvider) exportToDiff(ctx context.Context, out io.Writer) (*head
 		return nil, fmt.Errorf("error flushing path: %w", err)
 	}
 
-	builder := header.NewDiffMetadataBuilder(int64(o.header.Metadata.Size), o.blockSize)
+	builder := header.NewDiffMetadataBuilder(o.blockSize)
 
 	f, err := os.Open(o.path)
 	if err != nil {