Track dirty pages via WP

.github/actions/start-services/action.yml

@@ -101,7 +101,6 @@ runs:
         STORAGE_PROVIDER: "Local"
         ENVIRONMENT: "local"
         OTEL_COLLECTOR_GRPC_ENDPOINT: "localhost:4317"
-        MAX_PARALLEL_MEMFILE_SNAPSHOTTING: "2"
         SHARED_CHUNK_CACHE_PATH: "./.e2b-chunk-cache"
         EDGE_TOKEN: "abdcdefghijklmnop"
       run: |

packages/api/internal/cfg/model.go

@@ -6,7 +6,7 @@ const (
 	DefaultKernelVersion = "vmlinux-6.1.158"
 	// The Firecracker version the last tag + the short SHA (so we can build our dev previews)
 	// TODO: The short tag here has only 7 characters — the one from our build pipeline will likely have exactly 8 so this will break.
-	DefaultFirecrackerVersion = "v1.12.1_d990331"
+	DefaultFirecrackerVersion = "v1.12.1_g8769de7cf"
 )
 
 type Config struct {

packages/orchestrator/Makefile

@@ -47,7 +47,6 @@ run-debug:
 	GCP_DOCKER_REPOSITORY_NAME=$(GCP_DOCKER_REPOSITORY_NAME) \
 	GOOGLE_SERVICE_ACCOUNT_BASE64=$(GOOGLE_SERVICE_ACCOUNT_BASE64) \
 	OTEL_COLLECTOR_GRPC_ENDPOINT=$(OTEL_COLLECTOR_GRPC_ENDPOINT) \
-	MAX_PARALLEL_MEMFILE_SNAPSHOTTING=$(MAX_PARALLEL_MEMFILE_SNAPSHOTTING) \
 	./bin/orchestrator
 
 define setup_local_env

packages/orchestrator/internal/sandbox/diffcreator.go

@@ -3,14 +3,11 @@ package sandbox
 import (
 	"context"
 	"errors"
-	"fmt"
 	"io"
-	"os"
 
 	"github.com/bits-and-blooms/bitset"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/rootfs"
-	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 )
 
@@ -28,29 +25,25 @@ func (r *RootfsDiffCreator) process(ctx context.Context, out io.Writer) (*header
 }
 
 type MemoryDiffCreator struct {
-	memfile    *storage.TemporaryMemfile
+	memory     io.ReaderAt
 	dirtyPages *bitset.BitSet
 	blockSize  int64
 	doneHook   func(context.Context) error
 }
 
 func (r *MemoryDiffCreator) process(ctx context.Context, out io.Writer) (h *header.DiffMetadata, e error) {
 	defer func() {
-		err := r.doneHook(ctx)
-		if err != nil {
-			e = errors.Join(e, err)
+		if r.doneHook != nil {
+			err := r.doneHook(ctx)
+			if err != nil {
+				e = errors.Join(e, err)
+			}
 		}
 	}()
 
-	memfileSource, err := os.Open(r.memfile.Path())
-	if err != nil {
-		return nil, fmt.Errorf("failed to open memfile: %w", err)
-	}
-	defer memfileSource.Close()
-
 	return header.WriteDiffWithTrace(
 		ctx,
-		memfileSource,
+		r.memory,
 		r.blockSize,
 		r.dirtyPages,
 		out,

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

@@ -126,13 +126,11 @@ func (c *apiClient) pauseVM(ctx context.Context) error {
 func (c *apiClient) createSnapshot(
 	ctx context.Context,
 	snapfilePath string,
-	memfilePath string,
 ) error {
 	snapshotConfig := operations.CreateSnapshotParams{
 		Context: ctx,
 		Body: &models.SnapshotCreateParams{
 			SnapshotType: models.SnapshotCreateParamsSnapshotTypeFull,
-			MemFilePath:  &memfilePath,
 			SnapshotPath: &snapfilePath,
 		},
 	}
@@ -301,3 +299,17 @@ func (c *apiClient) startVM(ctx context.Context) error {
 
 	return nil
 }
+
+// vmInfo retrieves general information about an instance from the Firecracker API.
+func (c *apiClient) instanceInfo(ctx context.Context) (*models.InstanceInfo, error) {
+	req := operations.DescribeInstanceParams{
+		Context: ctx,
+	}
+
+	resp, err := c.client.Operations.DescribeInstance(&req)
+	if err != nil {
+		return nil, fmt.Errorf("error retrieving vm info: %w", err)
+	}
+
+	return resp.Payload, nil
+}

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

@@ -0,0 +1,32 @@
+package fc
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/memory"
+)
+
+func (p *Process) Memory(ctx context.Context) (*memory.View, error) {
+	pid, err := p.Pid()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get process pid: %w", err)
+	}
+
+	info, err := p.client.instanceInfo(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get instance info: %w", err)
+	}
+
+	mapping, err := memory.NewMappingFromFCInfo(info.MemoryRegions)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create memory mapping: %w", err)
+	}
+
+	view, err := memory.NewView(pid, mapping)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create memory view: %w", err)
+	}
+
+	return view, nil
+}

packages/orchestrator/internal/sandbox/fc/process.go

@@ -487,9 +487,9 @@ func (p *Process) Pause(ctx context.Context) error {
 }
 
 // CreateSnapshot VM needs to be paused before creating a snapshot.
-func (p *Process) CreateSnapshot(ctx context.Context, snapfilePath string, memfilePath string) error {
+func (p *Process) CreateSnapshot(ctx context.Context, snapfilePath string) error {
 	ctx, childSpan := tracer.Start(ctx, "create-snapshot-fc")
 	defer childSpan.End()
 
-	return p.client.createSnapshot(ctx, snapfilePath, memfilePath)
+	return p.client.createSnapshot(ctx, snapfilePath)
 }

packages/orchestrator/internal/sandbox/sandbox.go

@@ -673,10 +673,6 @@ func (s *Sandbox) Shutdown(ctx context.Context) error {
 		return fmt.Errorf("failed to pause VM: %w", err)
 	}
 
-	if _, err := s.memory.Disable(ctx); err != nil {
-		return fmt.Errorf("failed to disable uffd: %w", err)
-	}
-
 	// This is required because the FC API doesn't support passing /dev/null
 	tf, err := storage.TemplateFiles{
 		BuildID:            uuid.New().String(),
@@ -688,32 +684,17 @@ func (s *Sandbox) Shutdown(ctx context.Context) error {
 	}
 	defer tf.Close()
 
-	// The snapfile is required only because the FC API doesn't support passing /dev/null
 	snapfile := template.NewLocalFileLink(tf.CacheSnapfilePath())
 	defer snapfile.Close()
 
-	// The memfile is required only because the FC API doesn't support passing /dev/null
-	memfile, err := storage.AcquireTmpMemfile(ctx, s.config, tf.BuildID)
-	if err != nil {
-		return fmt.Errorf("failed to acquire memfile snapshot: %w", err)
-	}
-	defer memfile.Close()
-
 	err = s.process.CreateSnapshot(
 		ctx,
 		snapfile.Path(),
-		memfile.Path(),
 	)
 	if err != nil {
 		return fmt.Errorf("error creating snapshot: %w", err)
 	}
 
-	// Close the memfile right after the snapshot to release the lock.
-	err = memfile.Close()
-	if err != nil {
-		return fmt.Errorf("error closing memfile: %w", err)
-	}
-
 	// This should properly flush rootfs to the underlying device.
 	err = s.Close(ctx)
 	if err != nil {
@@ -753,41 +734,16 @@ func (s *Sandbox) Pause(
 	// Stop the health check before pausing the VM
 	s.Checks.Stop()
 
-	if err := s.process.Pause(ctx); err != nil {
-		return nil, fmt.Errorf("failed to pause VM: %w", err)
-	}
-
-	// This disables the uffd and returns the dirty pages.
-	// With FC async io engine, there can be some further writes to the memory during the actual create snapshot process,
-	// but as we are still including even read pages as dirty so this should not introduce more bugs right now.
-	dirty, err := s.memory.Disable(ctx)
+	err = s.process.Pause(ctx)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get dirty pages: %w", err)
+		return nil, fmt.Errorf("failed to pause VM: %w", err)
 	}
 
 	// Snapfile is not closed as it's returned and cached for later use (like resume)
 	snapfile := template.NewLocalFileLink(snapshotTemplateFiles.CacheSnapfilePath())
 	cleanup.AddNoContext(ctx, snapfile.Close)
-	// Memfile is also closed on diff creation processing
-	/* The process of snapshotting memory is as follows:
-	1. Pause FC via API
-	2. Snapshot FC via API—memory dump to “file on disk” that is actually tmpfs, because it is too slow
-	3. Create the diff - copy the diff pages from tmpfs to normal disk file
-	4. Delete tmpfs file
-	5. Unlock so another snapshot can use tmpfs space
-	*/
-	memfile, err := storage.AcquireTmpMemfile(ctx, s.config, buildID.String())
-	if err != nil {
-		return nil, fmt.Errorf("failed to acquire memfile snapshot: %w", err)
-	}
-	// Close the file even if an error occurs
-	defer memfile.Close()
 
-	err = s.process.CreateSnapshot(
-		ctx,
-		snapfile.Path(),
-		memfile.Path(),
-	)
+	err = s.process.CreateSnapshot(ctx, snapfile.Path())
 	if err != nil {
 		return nil, fmt.Errorf("error creating snapshot: %w", err)
 	}
@@ -797,23 +753,32 @@ func (s *Sandbox) Pause(
 	if err != nil {
 		return nil, fmt.Errorf("failed to get original memfile: %w", err)
 	}
+
 	originalRootfs, err := s.Template.Rootfs()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get original rootfs: %w", err)
 	}
 
+	dirty, err := s.memory.Dirty(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get dirty pages: %w", err)
+	}
+
+	memoryView, err := s.process.Memory(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get memory view: %w", err)
+	}
+	defer memoryView.Close()
+
 	// Start POSTPROCESSING
 	memfileDiff, memfileDiffHeader, err := pauseProcessMemory(
 		ctx,
 		buildID,
 		originalMemfile.Header(),
 		&MemoryDiffCreator{
-			memfile:    memfile,
+			memory:     memoryView,
 			dirtyPages: dirty.BitSet(),
 			blockSize:  originalMemfile.BlockSize(),
-			doneHook: func(context.Context) error {
-				return memfile.Close()
-			},
 		},
 		s.config.DefaultCacheDir,
 	)

packages/orchestrator/internal/sandbox/uffd/fdexit/fdexit.go

@@ -7,6 +7,8 @@ import (
 	"sync"
 )
 
+var ErrFdExit = errors.New("fd exit signal")
+
 // FdExit is a wrapper around a pipe that allows to signal the exit of the uffd.
 type FdExit struct {
 	r    *os.File

packages/orchestrator/internal/sandbox/uffd/memory/mapping.go

@@ -2,14 +2,24 @@ package memory
 
 import (
 	"fmt"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/fc/models"
 )
 
 type AddressNotFoundError struct {
 	hostVirtAddr uintptr
 }
 
 func (e AddressNotFoundError) Error() string {
-	return fmt.Sprintf("address %d not found in any mapping", e.hostVirtAddr)
+	return fmt.Sprintf("host virtual address %d not found in any mapping", e.hostVirtAddr)
+}
+
+type OffsetNotFoundError struct {
+	offset int64
+}
+
+func (e OffsetNotFoundError) Error() string {
+	return fmt.Sprintf("offset %d not found in any mapping", e.offset)
 }
 
 type Mapping struct {
@@ -20,13 +30,43 @@ func NewMapping(regions []Region) *Mapping {
 	return &Mapping{Regions: regions}
 }
 
-// GetOffset returns the relative offset and the page size of the mapped range for a given address.
-func (m *Mapping) GetOffset(hostVirtAddr uintptr) (int64, uint64, error) {
+func NewMappingFromFCInfo(regions []*models.GuestMemoryRegionMapping) (*Mapping, error) {
+	r := make([]Region, len(regions))
+
+	for i, infoRegion := range regions {
+		if infoRegion.BaseHostVirtAddr == nil || infoRegion.Size == nil || infoRegion.Offset == nil || infoRegion.PageSize == nil {
+			return nil, fmt.Errorf("missing required fields for memory region %d", i)
+		}
+
+		r[i] = Region{
+			BaseHostVirtAddr: uintptr(*infoRegion.BaseHostVirtAddr),
+			Size:             uintptr(*infoRegion.Size),
+			Offset:           uintptr(*infoRegion.Offset),
+			PageSize:         uintptr(*infoRegion.PageSize),
+		}
+	}
+
+	return NewMapping(r), nil
+}
+
+// GetOffset returns the relative offset and the pagesize of the mapped range for a given address.
+func (m *Mapping) GetOffset(hostVirtAddr uintptr) (int64, uintptr, error) {
 	for _, r := range m.Regions {
 		if hostVirtAddr >= r.BaseHostVirtAddr && hostVirtAddr < r.endHostVirtAddr() {
-			return r.shiftedOffset(hostVirtAddr), uint64(r.PageSize), nil
+			return r.shiftedOffset(hostVirtAddr), r.PageSize, nil
 		}
 	}
 
 	return 0, 0, AddressNotFoundError{hostVirtAddr: hostVirtAddr}
 }
+
+// GetHostVirtAddr returns the host virtual address and size of the remaining contiguous mapped host range for the given offset.
+func (m *Mapping) GetHostVirtAddr(off int64) (uintptr, int64, error) {
+	for _, r := range m.Regions {
+		if off >= int64(r.Offset) && off < r.endOffset() {
+			return r.shiftedHostVirtAddr(off), r.endOffset() - off, nil
+		}
+	}
+
+	return 0, 0, OffsetNotFoundError{offset: off}
+}