Skip to content

Commit 7424c06

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 7424c06

18 files changed

Lines changed: 124 additions & 164 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 & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ package block
33
import (
44
"iter"
55

6-
"github.com/bits-and-blooms/bitset"
6+
"github.com/RoaringBitmap/roaring/v2"
77

88
"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
99
)
@@ -36,24 +36,13 @@ func NewRangeFromBlocks(startIdx, numberOfBlocks, blockSize int64) Range {
3636
}
3737
}
3838

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] {
39+
// BitsetRanges returns a sequence of the ranges of the set bits of the bitmap.
40+
func BitsetRanges(b *roaring.Bitmap, blockSize int64) iter.Seq[Range] {
4141
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-
49-
return
50-
}
51-
52-
if !yield(NewRangeFromBlocks(int64(start), int64(end-start), blockSize)) {
42+
for start, endExcl := range b.Ranges() {
43+
if !yield(NewRangeFromBlocks(int64(start), int64(endExcl)-int64(start), blockSize)) {
5344
return
5445
}
55-
56-
start, found = b.NextSet(end + 1)
5746
}
5847
}
5948
}

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

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ 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"
1111
)
@@ -290,7 +290,7 @@ func TestRange_Offsets_Iteration(t *testing.T) {
290290

291291
func TestBitsetRanges_Empty(t *testing.T) {
292292
t.Parallel()
293-
b := bitset.New(100)
293+
b := roaring.New()
294294
blockSize := int64(4096)
295295

296296
ranges := slices.Collect(BitsetRanges(b, blockSize))
@@ -299,8 +299,8 @@ func TestBitsetRanges_Empty(t *testing.T) {
299299

300300
func TestBitsetRanges_SingleBit(t *testing.T) {
301301
t.Parallel()
302-
b := bitset.New(100)
303-
b.Set(5)
302+
b := roaring.New()
303+
b.Add(5)
304304
blockSize := int64(4096)
305305

306306
ranges := slices.Collect(BitsetRanges(b, blockSize))
@@ -313,12 +313,12 @@ func TestBitsetRanges_SingleBit(t *testing.T) {
313313

314314
func TestBitsetRanges_Contiguous(t *testing.T) {
315315
t.Parallel()
316-
b := bitset.New(100)
316+
b := roaring.New()
317317
// Set bits 2, 3, 4, 5
318-
b.Set(2)
319-
b.Set(3)
320-
b.Set(4)
321-
b.Set(5)
318+
b.Add(2)
319+
b.Add(3)
320+
b.Add(4)
321+
b.Add(5)
322322
blockSize := int64(4096)
323323

324324
ranges := slices.Collect(BitsetRanges(b, blockSize))
@@ -331,15 +331,15 @@ func TestBitsetRanges_Contiguous(t *testing.T) {
331331

332332
func TestBitsetRanges_MultipleRanges(t *testing.T) {
333333
t.Parallel()
334-
b := bitset.New(100)
334+
b := roaring.New()
335335
// Set bits 1, 2, 3 (contiguous)
336-
b.Set(1)
337-
b.Set(2)
338-
b.Set(3)
336+
b.Add(1)
337+
b.Add(2)
338+
b.Add(3)
339339
// Gap
340340
// Set bits 7, 8 (contiguous)
341-
b.Set(7)
342-
b.Set(8)
341+
b.Add(7)
342+
b.Add(8)
343343
blockSize := int64(4096)
344344

345345
ranges := slices.Collect(BitsetRanges(b, blockSize))
@@ -356,9 +356,9 @@ func TestBitsetRanges_MultipleRanges(t *testing.T) {
356356

357357
func TestBitsetRanges_AllSet(t *testing.T) {
358358
t.Parallel()
359-
b := bitset.New(10)
360-
for i := range uint(10) {
361-
b.Set(i)
359+
b := roaring.New()
360+
for i := range uint32(10) {
361+
b.Add(i)
362362
}
363363
blockSize := int64(4096)
364364

@@ -372,10 +372,10 @@ func TestBitsetRanges_AllSet(t *testing.T) {
372372

373373
func TestBitsetRanges_EndOfBitset(t *testing.T) {
374374
t.Parallel()
375-
b := bitset.New(20)
375+
b := roaring.New()
376376
// Set bits 15, 16, 17, 18, 19 (at the end)
377-
for i := uint(15); i < 20; i++ {
378-
b.Set(i)
377+
for i := uint32(15); i < 20; i++ {
378+
b.Add(i)
379379
}
380380
blockSize := int64(4096)
381381

@@ -389,12 +389,12 @@ func TestBitsetRanges_EndOfBitset(t *testing.T) {
389389

390390
func TestBitsetRanges_Sparse(t *testing.T) {
391391
t.Parallel()
392-
b := bitset.New(100)
392+
b := roaring.New()
393393
// Set individual bits with gaps
394-
b.Set(0)
395-
b.Set(10)
396-
b.Set(20)
397-
b.Set(30)
394+
b.Add(0)
395+
b.Add(10)
396+
b.Add(20)
397+
b.Add(30)
398398
blockSize := int64(4096)
399399

400400
ranges := slices.Collect(BitsetRanges(b, blockSize))

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

Lines changed: 14 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -4,23 +4,21 @@ import (
44
"iter"
55
"sync"
66

7-
"github.com/bits-and-blooms/bitset"
7+
"github.com/RoaringBitmap/roaring/v2"
88

99
"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
10-
"github.com/e2b-dev/infra/packages/shared/pkg/utils"
1110
)
1211

1312
type Tracker struct {
14-
b *bitset.BitSet
13+
b *roaring.Bitmap
1514
mu sync.RWMutex
1615

1716
blockSize int64
1817
}
1918

2019
func NewTracker(blockSize int64) *Tracker {
2120
return &Tracker{
22-
// The bitset resizes automatically based on the maximum set bit.
23-
b: bitset.New(0),
21+
b: roaring.New(),
2422
blockSize: blockSize,
2523
}
2624
}
@@ -29,27 +27,21 @@ func (t *Tracker) Has(off int64) bool {
2927
t.mu.RLock()
3028
defer t.mu.RUnlock()
3129

32-
return t.b.Test(uint(header.BlockIdx(off, t.blockSize)))
30+
return t.b.Contains(uint32(header.BlockIdx(off, t.blockSize)))
3331
}
3432

3533
func (t *Tracker) Add(off int64) {
3634
t.mu.Lock()
3735
defer t.mu.Unlock()
3836

39-
t.b.Set(uint(header.BlockIdx(off, t.blockSize)))
37+
t.b.Add(uint32(header.BlockIdx(off, t.blockSize)))
4038
}
4139

4240
func (t *Tracker) Reset() {
4341
t.mu.Lock()
4442
defer t.mu.Unlock()
4543

46-
t.b.ClearAll()
47-
}
48-
49-
// BitSet returns the bitset.
50-
// This is not safe to use concurrently.
51-
func (t *Tracker) BitSet() *bitset.BitSet {
52-
return t.b
44+
t.b.Clear()
5345
}
5446

5547
func (t *Tracker) BlockSize() int64 {
@@ -66,15 +58,17 @@ func (t *Tracker) Clone() *Tracker {
6658
}
6759
}
6860

61+
// Offsets returns a lazy iterator over tracked block offsets.
62+
// It clones the internal bitmap so the iterator is safe to use after the lock is released.
6963
func (t *Tracker) Offsets() iter.Seq[int64] {
7064
t.mu.RLock()
7165
defer t.mu.RUnlock()
7266

73-
return bitsetOffsets(t.b.Clone(), t.BlockSize())
74-
}
67+
snapshot := t.b.Clone()
7568

76-
func bitsetOffsets(b *bitset.BitSet, blockSize int64) iter.Seq[int64] {
77-
return utils.TransformTo(b.EachSet(), func(idx uint) int64 {
78-
return header.BlockOffset(int64(idx), blockSize)
79-
})
69+
return func(yield func(int64) bool) {
70+
snapshot.Iterate(func(idx uint32) bool {
71+
return yield(header.BlockOffset(int64(idx), t.blockSize))
72+
})
73+
}
8074
}

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import (
55
"fmt"
66
"runtime"
77

8-
"github.com/bits-and-blooms/bitset"
8+
"github.com/RoaringBitmap/roaring/v2"
99
"github.com/firecracker-microvm/firecracker-go-sdk"
1010
"github.com/go-openapi/strfmt"
1111

@@ -451,8 +451,8 @@ func (c *apiClient) memoryInfo(ctx context.Context, blockSize int64) (*header.Di
451451
}
452452

453453
return &header.DiffMetadata{
454-
Dirty: bitset.From(res.Payload.Resident),
455-
Empty: bitset.From(res.Payload.Empty),
454+
Dirty: roaring.FromDense(res.Payload.Resident, false),
455+
Empty: roaring.FromDense(res.Payload.Empty, false),
456456
BlockSize: blockSize,
457457
}, nil
458458
}
@@ -468,8 +468,8 @@ func (c *apiClient) dirtyMemory(ctx context.Context, blockSize int64) (*header.D
468468
}
469469

470470
return &header.DiffMetadata{
471-
Dirty: bitset.From(res.Payload.Bitmap),
472-
Empty: bitset.New(0),
471+
Dirty: roaring.FromDense(res.Payload.Bitmap, false),
472+
Empty: roaring.New(),
473473
BlockSize: blockSize,
474474
}, nil
475475
}

0 commit comments

Comments
 (0)