Skip to content

Commit a55db6c

Browse files
ValentaTomasValentaTomas
committed
refactor: replace bits-and-blooms/bitset with roaring bitmaps
Use roaring v2.17.0 with the new Ranges() iterator throughout. Remove bitset as a direct dependency.
1 parent 8245988 commit a55db6c

19 files changed

Lines changed: 212 additions & 209 deletions

File tree

‎packages/orchestrator/go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ require (
1717
cloud.google.com/go/storage v1.59.2
1818
connectrpc.com/connect v1.18.1
1919
github.com/Merovius/nbd v0.0.0-20240812113926-fd65a54c9949
20+
github.com/RoaringBitmap/roaring/v2 v2.17.0
2021
github.com/aws/aws-sdk-go-v2/config v1.32.6
2122
github.com/aws/aws-sdk-go-v2/credentials v1.19.6
2223
github.com/aws/aws-sdk-go-v2/service/ecr v1.44.0
@@ -97,7 +98,6 @@ require (
9798
github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.54.0 // indirect
9899
github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.54.0 // indirect
99100
github.com/Microsoft/go-winio v0.6.2 // indirect
100-
github.com/RoaringBitmap/roaring/v2 v2.16.1 // indirect
101101
github.com/andybalholm/brotli v1.2.0 // indirect
102102
github.com/armon/go-metrics v0.4.1 // indirect
103103
github.com/aws/aws-sdk-go-v2 v1.41.0 // indirect

‎packages/orchestrator/go.sum‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,14 @@ import (
1313
"syscall"
1414
"time"
1515

16-
"github.com/bits-and-blooms/bitset"
16+
"github.com/RoaringBitmap/roaring/v2"
1717
"github.com/edsrzf/mmap-go"
1818
"go.opentelemetry.io/otel"
1919
"go.opentelemetry.io/otel/attribute"
2020
"go.uber.org/zap"
2121
"golang.org/x/sys/unix"
2222

23-
"github.com/e2b-dev/infra/packages/shared/pkg/atomicbitset"
23+
"github.com/e2b-dev/infra/packages/shared/pkg/syncroaring"
2424
"github.com/e2b-dev/infra/packages/shared/pkg/logger"
2525
"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
2626
"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
@@ -53,7 +53,7 @@ type Cache struct {
5353
blockSize int64
5454
mmap *mmap.MMap
5555
mu sync.RWMutex
56-
dirty *atomicbitset.Bitset
56+
dirty *syncroaring.Bitset
5757
dirtyFile bool
5858
closed atomic.Bool
5959
}
@@ -72,7 +72,7 @@ func NewCache(size, blockSize int64, filePath string, dirtyFile bool) (*Cache, e
7272
size: size,
7373
blockSize: blockSize,
7474
dirtyFile: dirtyFile,
75-
dirty: atomicbitset.New(),
75+
dirty: syncroaring.New(),
7676
}, nil
7777
}
7878

@@ -97,7 +97,7 @@ func NewCache(size, blockSize int64, filePath string, dirtyFile bool) (*Cache, e
9797
size: size,
9898
blockSize: blockSize,
9999
dirtyFile: dirtyFile,
100-
dirty: atomicbitset.New(),
100+
dirty: syncroaring.New(),
101101
}, nil
102102
}
103103

@@ -117,7 +117,7 @@ func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMet
117117
}
118118

119119
if c.mmap == nil {
120-
return header.NewDiffMetadata(c.blockSize, bitset.New(0)), nil
120+
return header.NewDiffMetadata(c.blockSize, roaring.New()), nil
121121
}
122122

123123
f, err := os.Open(c.filePath)
@@ -136,7 +136,7 @@ func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMet
136136
logger.L().Warn(ctx, "error syncing file", zap.Error(err))
137137
}
138138

139-
diffMetadata := header.NewDiffMetadata(c.blockSize, c.dirty.BitSet())
139+
diffMetadata := header.NewDiffMetadata(c.blockSize, c.dirty.UnsafeBitmap())
140140

141141
dst := int(out.Fd())
142142
var writeOffset int64
@@ -194,7 +194,7 @@ func (c *Cache) ExportToDiff(ctx context.Context, out *os.File) (*header.DiffMet
194194
telemetry.SetAttributes(ctx,
195195
attribute.Int64("copy_ms", time.Since(copyStart).Milliseconds()),
196196
attribute.Int64("total_size_bytes", c.size),
197-
attribute.Int64("dirty_size_bytes", int64(diffMetadata.Dirty.Count())*c.blockSize),
197+
attribute.Int64("dirty_size_bytes", int64(diffMetadata.Dirty.GetCardinality())*c.blockSize),
198198
attribute.Int64("total_ranges", totalRanges),
199199
)
200200

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

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -257,8 +257,8 @@ func TestCacheExportToDiff_ZeroDirtyBlockEmittedAsDirtyPayload(t *testing.T) {
257257
diffMetadata, err := cache.ExportToDiff(t.Context(), out)
258258
require.NoError(t, err)
259259

260-
require.EqualValues(t, 1, diffMetadata.Dirty.Count(), "zero-filled dirty block should be emitted as dirty payload")
261-
require.EqualValues(t, 0, diffMetadata.Empty.Count(), "zero-filled dirty block should not be tracked in empty metadata")
260+
require.EqualValues(t, 1, diffMetadata.Dirty.GetCardinality(), "zero-filled dirty block should be emitted as dirty payload")
261+
require.EqualValues(t, 0, diffMetadata.Empty.GetCardinality(), "zero-filled dirty block should not be tracked in empty metadata")
262262

263263
stat, err := out.Stat()
264264
require.NoError(t, err)
@@ -335,8 +335,8 @@ func TestCacheExportToDiff_MixedDirtyBlocksKeepsZeroBlockInDiff(t *testing.T) {
335335
diffMetadata, err := cache.ExportToDiff(t.Context(), out)
336336
require.NoError(t, err)
337337

338-
require.EqualValues(t, 2, diffMetadata.Dirty.Count())
339-
require.EqualValues(t, 0, diffMetadata.Empty.Count(), "mixed export should still skip empty tracking for zero-filled dirty blocks")
338+
require.EqualValues(t, 2, diffMetadata.Dirty.GetCardinality())
339+
require.EqualValues(t, 0, diffMetadata.Empty.GetCardinality(), "mixed export should still skip empty tracking for zero-filled dirty blocks")
340340

341341
_, err = out.Seek(0, io.SeekStart)
342342
require.NoError(t, err)
@@ -399,10 +399,10 @@ func TestCacheExportToDiff_NonContiguousDirtyBlocksPreserveRangeOrder(t *testing
399399
diffMetadata, err := cache.ExportToDiff(t.Context(), out)
400400
require.NoError(t, err)
401401

402-
require.EqualValues(t, 2, diffMetadata.Dirty.Count())
403-
require.True(t, diffMetadata.Dirty.Test(0))
404-
require.True(t, diffMetadata.Dirty.Test(3))
405-
require.EqualValues(t, 0, diffMetadata.Empty.Count())
402+
require.EqualValues(t, 2, diffMetadata.Dirty.GetCardinality())
403+
require.True(t, diffMetadata.Dirty.Contains(0))
404+
require.True(t, diffMetadata.Dirty.Contains(3))
405+
require.EqualValues(t, 0, diffMetadata.Empty.GetCardinality())
406406

407407
_, err = out.Seek(0, io.SeekStart)
408408
require.NoError(t, err)

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

Lines changed: 5 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,8 @@ package block
33
import (
44
"iter"
55

6-
"github.com/bits-and-blooms/bitset"
7-
86
"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
7+
"github.com/e2b-dev/infra/packages/shared/pkg/syncroaring"
98
)
109

1110
type Range struct {
@@ -36,24 +35,13 @@ func NewRangeFromBlocks(startIdx, numberOfBlocks, blockSize int64) Range {
3635
}
3736
}
3837

39-
// bitsetRanges returns a sequence of the ranges of the set bits of the bitset.
40-
func BitsetRanges(b *bitset.BitSet, blockSize int64) iter.Seq[Range] {
38+
// BitsetRanges returns a sequence of the ranges of the set bits of the bitmap.
39+
func BitsetRanges(b *syncroaring.ReadOnly, blockSize int64) iter.Seq[Range] {
4140
return func(yield func(Range) bool) {
42-
start, found := b.NextSet(0)
43-
44-
for found {
45-
end, endOk := b.NextClear(start)
46-
if !endOk {
47-
yield(NewRangeFromBlocks(int64(start), int64(b.Len()-start), blockSize))
48-
41+
for start, endExcl := range b.Ranges() {
42+
if !yield(NewRangeFromBlocks(int64(start), int64(endExcl)-int64(start), blockSize)) {
4943
return
5044
}
51-
52-
if !yield(NewRangeFromBlocks(int64(start), int64(end-start), blockSize)) {
53-
return
54-
}
55-
56-
start, found = b.NextSet(end + 1)
5745
}
5846
}
5947
}

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

Lines changed: 35 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,11 @@ import (
55
"slices"
66
"testing"
77

8-
"github.com/bits-and-blooms/bitset"
8+
"github.com/RoaringBitmap/roaring/v2"
99
"github.com/stretchr/testify/assert"
1010
"github.com/stretchr/testify/require"
11+
12+
"github.com/e2b-dev/infra/packages/shared/pkg/syncroaring"
1113
)
1214

1315
// rangeOffsets returns the block offsets contained in the range.
@@ -290,20 +292,20 @@ func TestRange_Offsets_Iteration(t *testing.T) {
290292

291293
func TestBitsetRanges_Empty(t *testing.T) {
292294
t.Parallel()
293-
b := bitset.New(100)
295+
b := roaring.New()
294296
blockSize := int64(4096)
295297

296-
ranges := slices.Collect(BitsetRanges(b, blockSize))
298+
ranges := slices.Collect(BitsetRanges(syncroaring.NewReadOnly(b), blockSize))
297299
assert.Empty(t, ranges)
298300
}
299301

300302
func TestBitsetRanges_SingleBit(t *testing.T) {
301303
t.Parallel()
302-
b := bitset.New(100)
303-
b.Set(5)
304+
b := roaring.New()
305+
b.Add(5)
304306
blockSize := int64(4096)
305307

306-
ranges := slices.Collect(BitsetRanges(b, blockSize))
308+
ranges := slices.Collect(BitsetRanges(syncroaring.NewReadOnly(b), blockSize))
307309
require.Len(t, ranges, 1)
308310
assert.Equal(t, Range{
309311
Start: 20480, // 5 * 4096
@@ -313,15 +315,15 @@ func TestBitsetRanges_SingleBit(t *testing.T) {
313315

314316
func TestBitsetRanges_Contiguous(t *testing.T) {
315317
t.Parallel()
316-
b := bitset.New(100)
318+
b := roaring.New()
317319
// Set bits 2, 3, 4, 5
318-
b.Set(2)
319-
b.Set(3)
320-
b.Set(4)
321-
b.Set(5)
320+
b.Add(2)
321+
b.Add(3)
322+
b.Add(4)
323+
b.Add(5)
322324
blockSize := int64(4096)
323325

324-
ranges := slices.Collect(BitsetRanges(b, blockSize))
326+
ranges := slices.Collect(BitsetRanges(syncroaring.NewReadOnly(b), blockSize))
325327
require.Len(t, ranges, 1)
326328
assert.Equal(t, Range{
327329
Start: 8192, // 2 * 4096
@@ -331,18 +333,18 @@ func TestBitsetRanges_Contiguous(t *testing.T) {
331333

332334
func TestBitsetRanges_MultipleRanges(t *testing.T) {
333335
t.Parallel()
334-
b := bitset.New(100)
336+
b := roaring.New()
335337
// Set bits 1, 2, 3 (contiguous)
336-
b.Set(1)
337-
b.Set(2)
338-
b.Set(3)
338+
b.Add(1)
339+
b.Add(2)
340+
b.Add(3)
339341
// Gap
340342
// Set bits 7, 8 (contiguous)
341-
b.Set(7)
342-
b.Set(8)
343+
b.Add(7)
344+
b.Add(8)
343345
blockSize := int64(4096)
344346

345-
ranges := slices.Collect(BitsetRanges(b, blockSize))
347+
ranges := slices.Collect(BitsetRanges(syncroaring.NewReadOnly(b), blockSize))
346348
require.Len(t, ranges, 2)
347349
assert.Equal(t, Range{
348350
Start: 4096, // 1 * 4096
@@ -356,13 +358,13 @@ func TestBitsetRanges_MultipleRanges(t *testing.T) {
356358

357359
func TestBitsetRanges_AllSet(t *testing.T) {
358360
t.Parallel()
359-
b := bitset.New(10)
360-
for i := range uint(10) {
361-
b.Set(i)
361+
b := roaring.New()
362+
for i := range uint32(10) {
363+
b.Add(i)
362364
}
363365
blockSize := int64(4096)
364366

365-
ranges := slices.Collect(BitsetRanges(b, blockSize))
367+
ranges := slices.Collect(BitsetRanges(syncroaring.NewReadOnly(b), blockSize))
366368
require.Len(t, ranges, 1)
367369
assert.Equal(t, Range{
368370
Start: 0,
@@ -372,14 +374,14 @@ func TestBitsetRanges_AllSet(t *testing.T) {
372374

373375
func TestBitsetRanges_EndOfBitset(t *testing.T) {
374376
t.Parallel()
375-
b := bitset.New(20)
377+
b := roaring.New()
376378
// Set bits 15, 16, 17, 18, 19 (at the end)
377-
for i := uint(15); i < 20; i++ {
378-
b.Set(i)
379+
for i := uint32(15); i < 20; i++ {
380+
b.Add(i)
379381
}
380382
blockSize := int64(4096)
381383

382-
ranges := slices.Collect(BitsetRanges(b, blockSize))
384+
ranges := slices.Collect(BitsetRanges(syncroaring.NewReadOnly(b), blockSize))
383385
require.Len(t, ranges, 1)
384386
assert.Equal(t, Range{
385387
Start: 61440, // 15 * 4096
@@ -389,15 +391,15 @@ func TestBitsetRanges_EndOfBitset(t *testing.T) {
389391

390392
func TestBitsetRanges_Sparse(t *testing.T) {
391393
t.Parallel()
392-
b := bitset.New(100)
394+
b := roaring.New()
393395
// Set individual bits with gaps
394-
b.Set(0)
395-
b.Set(10)
396-
b.Set(20)
397-
b.Set(30)
396+
b.Add(0)
397+
b.Add(10)
398+
b.Add(20)
399+
b.Add(30)
398400
blockSize := int64(4096)
399401

400-
ranges := slices.Collect(BitsetRanges(b, blockSize))
402+
ranges := slices.Collect(BitsetRanges(syncroaring.NewReadOnly(b), blockSize))
401403
require.Len(t, ranges, 4)
402404
assert.Equal(t, Range{Start: 0, Size: 4096}, ranges[0])
403405
assert.Equal(t, Range{Start: 40960, Size: 4096}, ranges[1])

0 commit comments

Comments
 (0)