Add method for preemptive UFFD copy to FC memory

.gitignore

@@ -13,3 +13,4 @@ tests/periodic-test/build-template/e2b.toml
 .air
 go.work.sum
 .infisical.json
+.vscode/mise-tools/

.vscode/settings.json

@@ -89,5 +89,11 @@
     "editor.codeActionsOnSave": {
       "source.organizeImports": "explicit"
     }
+  },
+  "go.goroot": "${workspaceFolder}/.vscode/mise-tools/goRoot",
+  "go.alternateTools": {
+    "go": "${workspaceFolder}/.vscode/mise-tools/go",
+    "dlv": "${workspaceFolder}/.vscode/mise-tools/dlv",
+    "gopls": "${workspaceFolder}/.vscode/mise-tools/gopls"
   }
 }

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

@@ -0,0 +1,56 @@
+package block
+
+import (
+	"iter"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+)
+
+type Range struct {
+	// Start is the start address of the range in bytes.
+	// Start is inclusive.
+	Start int64
+	// Size is the size of the range in bytes.
+	Size int64
+}
+
+// NewRange creates a new range from a start address and size in bytes.
+func NewRange(start, size int64) Range {
+	return Range{
+		Start: start,
+		Size:  size,
+	}
+}
+
+func NewRangeFromBuildMap(b *header.BuildMap) Range {
+	return Range{
+		Start: int64(b.Offset),
+		Size:  int64(b.Length),
+	}
+}
+
+// NewRangeFromBlocks creates a new range from a start index and number of blocks.
+func NewRangeFromBlocks(startIdx, numberOfBlocks, blockSize int64) Range {
+	return Range{
+		Start: header.BlockOffset(startIdx, blockSize),
+		Size:  header.BlockOffset(numberOfBlocks, blockSize),
+	}
+}
+
+// End returns the end address of the range in bytes.
+// The end address is exclusive.
+func (r *Range) End() int64 {
+	return r.Start + r.Size
+}
+
+// Offsets returns the block offsets contained in the range.
+// This assumes the Range.Start is a multiple of the blockSize.
+func (r *Range) Offsets(blockSize int64) iter.Seq[int64] {
+	return func(yield func(offset int64) bool) {
+		for i := r.Start; i < r.End(); i += blockSize {
+			if !yield(i) {
+				return
+			}
+		}
+	}
+}

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

@@ -0,0 +1,150 @@
+package block
+
+import (
+	"iter"
+	"slices"
+	"sync"
+
+	"github.com/bits-and-blooms/bitset"
+
+	"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
+	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),
+		blockSize: blockSize,
+	}
+}
+
+func NewTrackerFromBitSet(b *bitset.BitSet, blockSize int64) *Tracker {
+	return &Tracker{
+		b:         b.Clone(),
+		blockSize: blockSize,
+	}
+}
+
+func NewTrackerFromMapping(mapping []*header.BuildMap, blockSize int64) *Tracker {
+	return NewTrackerFromRanges(slices.Collect(mappingRanges(mapping)), blockSize)
+}
+
+func NewTrackerFromRanges(ranges []Range, blockSize int64) *Tracker {
+	b := bitset.New(0)
+
+	for _, r := range ranges {
+		b.FlipRange(uint(header.BlockIdx(r.Start, blockSize)), uint(header.BlockIdx(r.End(), blockSize)))
+	}
+
+	return &Tracker{
+		b:         b,
+		blockSize: blockSize,
+	}
+}
+
+func (t *Tracker) Offsets() []int64 {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return slices.Collect(bitsetOffsets(t.b, t.blockSize))
+}
+
+func (t *Tracker) Ranges() []Range {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return slices.Collect(bitsetRanges(t.b, t.blockSize))
+}
+
+func (t *Tracker) Has(off int64) bool {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return t.b.Test(uint(header.BlockIdx(off, t.blockSize)))
+}
+
+func (t *Tracker) Add(off int64) bool {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	if t.b.Test(uint(header.BlockIdx(off, t.blockSize))) {
+		return false
+	}
+
+	t.b.Set(uint(header.BlockIdx(off, t.blockSize)))
+
+	return true
+}
+
+func (t *Tracker) Reset() {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	t.b.ClearAll()
+}
+
+// BitSet returns a clone of the bitset and the block size.
+func (t *Tracker) BitSet() *bitset.BitSet {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return t.b.Clone()
+}
+
+func (t *Tracker) BlockSize() int64 {
+	return t.blockSize
+}
+
+func (t *Tracker) Clone() *Tracker {
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+
+	return &Tracker{
+		b:         t.b.Clone(),
+		blockSize: t.BlockSize(),
+	}
+}
+
+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)
+	})
+}
+
+// bitsetRanges returns a sequence of the ranges of the set bits of the bitset.
+func bitsetRanges(b *bitset.BitSet, blockSize int64) iter.Seq[Range] {
+	return func(yield func(Range) bool) {
+		for start, ok := b.NextSet(0); ok; {
+			end, ok := b.NextClear(start)
+			if !ok {
+				yield(NewRange(header.BlockOffset(int64(start), blockSize), header.BlockOffset(int64(b.Len()-start), blockSize)))
+
+				return
+			}
+
+			if !yield(NewRange(header.BlockOffset(int64(start), blockSize), header.BlockOffset(int64(end-start), blockSize))) {
+				return
+			}
+
+			start, ok = b.NextSet(end + 1)
+		}
+	}
+}
+
+func mappingRanges(mapping []*header.BuildMap) iter.Seq[Range] {
+	return func(yield func(Range) bool) {
+		for _, buildMap := range mapping {
+			if !yield(NewRangeFromBuildMap(buildMap)) {
+				return
+			}
+		}
+	}
+}

packages/orchestrator/internal/sandbox/block/tracker_test.go

@@ -0,0 +1,9 @@
+package block
+
+import (
+	"testing"
+)
+
+func TestTrackerBitSetRanges(t *testing.T) {
+	t.Fatal("not implemented")
+}

packages/orchestrator/internal/sandbox/build/masked_overlay.go

@@ -0,0 +1,99 @@
+package build
+
+import (
+	"context"
+	"fmt"
+	"slices"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+)
+
+type MaskedOverlay struct {
+	lower block.ReadonlyDevice
+	upper block.ReadonlyDevice
+
+	mask *block.Tracker
+}
+
+var _ block.ReadonlyDevice = (*MaskedOverlay)(nil)
+
+func NewMaskedOverlay(lower, upper block.ReadonlyDevice) (*MaskedOverlay, error) {
+	if lower.BlockSize() != upper.BlockSize() {
+		return nil, fmt.Errorf("lower and upper block sizes do not match")
+	}
+
+	lowerSize, err := lower.Size()
+	if err != nil {
+		return nil, err
+	}
+
+	upperSize, err := upper.Size()
+	if err != nil {
+		return nil, err
+	}
+
+	if lowerSize != upperSize {
+		return nil, fmt.Errorf("lower and upper sizes do not match")
+	}
+
+	upperMapping := slices.Collect(upper.Header().FilterMapping(upper.Header().Metadata.BuildId))
+	mask := block.NewTrackerFromMapping(upperMapping, upper.BlockSize())
+
+	return &MaskedOverlay{
+		lower: lower,
+		upper: upper,
+		mask:  mask,
+	}, nil
+}
+
+func (m *MaskedOverlay) BlockSize() int64 {
+	return m.lower.BlockSize()
+}
+
+func (m *MaskedOverlay) Close() error {
+	return nil
+}
+
+func (m *MaskedOverlay) Header() *header.Header {
+	return m.lower.Header()
+}
+
+// TODO: We should be able to read by bigger chunks than just the block size
+func (m *MaskedOverlay) ReadAt(ctx context.Context, p []byte, off int64) (n int, err error) {
+	for n < len(p) {
+		currentOff := off + int64(n)
+		if m.mask.Has(currentOff) {
+			readN, err := m.upper.ReadAt(ctx, p[n:], currentOff)
+
+			n += readN
+
+			if err != nil {
+				return n, err
+			}
+		} else {
+			readN, err := m.lower.ReadAt(ctx, p[n:], currentOff)
+
+			n += readN
+
+			if err != nil {
+				return n, err
+			}
+		}
+	}
+
+	return n, nil
+}
+
+func (m *MaskedOverlay) Size() (int64, error) {
+	return m.lower.Size()
+}
+
+// TODO: We should handle the slice even for more than one block size
+func (m *MaskedOverlay) Slice(ctx context.Context, off, length int64) ([]byte, error) {
+	if m.mask.Has(off) {
+		return m.upper.Slice(ctx, off, length)
+	}
+
+	return m.lower.Slice(ctx, off, length)
+}

packages/orchestrator/internal/sandbox/sandbox.go

@@ -671,11 +671,18 @@ func (s *Sandbox) Pause(
 		return nil, fmt.Errorf("failed to pause VM: %w", err)
 	}
 
-	dirtyPages, err := s.memory.Disable(ctx)
+	err = s.memory.Disable(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to disable uffd: %w", err)
 	}
 
+	dirty, err := s.memory.Dirty(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get dirty pages: %w", err)
+	}
+
+	dirtyPages := dirty.BitSet()
+
 	// Snapfile is not closed as it's returned and cached for later use (like resume)
 	snapfile := template.NewLocalFileLink(snapshotTemplateFiles.CacheSnapfilePath())
 	// Memfile is also closed on diff creation processing
@@ -930,7 +937,7 @@ func serveMemory(
 	ctx, span := tracer.Start(ctx, "serve-memory")
 	defer span.End()
 
-	fcUffd, err := uffd.New(memfile, socketPath, memfile.BlockSize())
+	fcUffd, err := uffd.New(memfile, socketPath)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create uffd: %w", err)
 	}