uffd: add support for UFFD_EVENT_REMOVE events

packages/api/internal/template-manager/create_template.go

@@ -16,6 +16,7 @@ import (
 	"github.com/e2b-dev/infra/packages/db/pkg/types"
 	"github.com/e2b-dev/infra/packages/db/queries"
 	"github.com/e2b-dev/infra/packages/shared/pkg/fcversion"
+	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	templatemanagergrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
 	"github.com/e2b-dev/infra/packages/shared/pkg/id"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
@@ -109,6 +110,8 @@ func (tm *TemplateManager) CreateTemplate(
 		return fmt.Errorf("failed to convert image registry: %w", err)
 	}
 
+	freePageReporting := features.HasFreePageReporting() && tm.featureFlags.BoolFlag(ctx, featureflags.FreePageReportingFlag, featureflags.TeamContext(teamID.String()))
+
 	// TODO(ENG-3852): Remove later. KernelVersion and FirecrackerVersion are deprecated on
 	// template-manager selects its own versions and reports the ones it actually
 	// used via TemplateBuildMetadata. They are still populated here for
@@ -123,6 +126,7 @@ func (tm *TemplateManager) CreateTemplate(
 		KernelVersion:      kernelVersion,
 		FirecrackerVersion: firecrackerVersion,
 		HugePages:          features.HasHugePages(),
+		FreePageReporting:  &freePageReporting,
 		StartCommand:       startCmd,
 		ReadyCommand:       readyCmd,
 		Force:              force,

packages/orchestrator/cmd/create-build/main.go

@@ -63,6 +63,7 @@ func main() {
 	memory := flag.Int("memory", 1024, "memory MB")
 	disk := flag.Int("disk", 1024, "disk MB")
 	hugePages := flag.Bool("hugepages", true, "use 2MB huge pages for memory (false = 4KB pages)")
+	freePageReporting := flag.Bool("free-page-reporting", false, "enable free page reporting via balloon device (requires Firecracker v1.14+)")
 	startCmd := flag.String("start-cmd", "", "start command")
 	setupCmd := flag.String("setup-cmd", "", "setup command to run during build (e.g., install deps)")
 	readyCmd := flag.String("ready-cmd", "", "ready check command")
@@ -100,7 +101,16 @@ func main() {
 		log.Fatalf("network config: %v", err)
 	}
 
-	err = doBuild(ctx, *templateID, *toBuild, *fromBuild, *kernel, *fc, *vcpu, *memory, *disk, *hugePages, *startCmd, *setupCmd, *readyCmd, localMode, *verbose, *timeout, builderConfig, networkConfig)
+	// Detect if --free-page-reporting was explicitly set; if not, pass nil so
+	// doBuild can default based on the Firecracker version.
+	var fprOverride *bool
+	flag.Visit(func(f *flag.Flag) {
+		if f.Name == "free-page-reporting" {
+			fprOverride = freePageReporting
+		}
+	})
+
+	err = doBuild(ctx, *templateID, *toBuild, *fromBuild, *kernel, *fc, *vcpu, *memory, *disk, *hugePages, fprOverride, *startCmd, *setupCmd, *readyCmd, localMode, *verbose, *timeout, builderConfig, networkConfig)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -199,9 +209,10 @@ func setupEnv(ctx context.Context, storagePath, sandboxDir, kernel, fc string, l
 
 func doBuild(
 	parentCtx context.Context,
-	templateID, buildID, fromBuild, kernel, fc string,
+	templateID, buildID, fromBuild, kernel, fcVersion string,
 	vcpu, memory, disk int,
 	hugePages bool,
+	freePageReporting *bool,
 	startCmd, setupCmd, readyCmd string,
 	localMode, verbose bool,
 	timeout int,
@@ -341,6 +352,18 @@ func doBuild(
 		})
 	}
 
+	// Default FPR to enabled when the FC version supports it (v1.14+); explicit flag overrides.
+	var fprEnabled bool
+	if freePageReporting != nil {
+		fprEnabled = *freePageReporting
+	} else {
+		versionOnly, _, _ := strings.Cut(fcVersion, "_")
+		supported, err := utils.IsGTEVersion(versionOnly, "v1.14.0")
+		if err == nil {
+			fprEnabled = supported
+		}
+	}
+
 	tmpl := config.TemplateConfig{
 		Version:            templates.TemplateV2LatestVersion,
 		TemplateID:         templateID,
@@ -349,10 +372,11 @@ func doBuild(
 		MemoryMB:           int64(memory),
 		DiskSizeMB:         int64(disk),
 		HugePages:          hugePages,
+		FreePageReporting:  fprEnabled,
 		StartCmd:           startCmd,
 		ReadyCmd:           readyCmd,
 		KernelVersion:      kernel,
-		FirecrackerVersion: fc,
+		FirecrackerVersion: fcVersion,
 		Steps:              steps,
 	}
 

packages/orchestrator/cmd/smoketest/smoke_test.go

@@ -40,6 +40,7 @@ import (
 	sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 	"github.com/e2b-dev/infra/packages/shared/pkg/templates"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
 const (
@@ -71,11 +72,14 @@ func TestSmokeAllFCVersions(t *testing.T) { //nolint:paralleltest // subtests sh
 	defer infra.close(ctx)
 
 	for fcMajor, fcVersion := range featureflags.FirecrackerVersionMap { //nolint:paralleltest // sequential by design
+		versionOnly, _, _ := strings.Cut(fcVersion, "_")
+		fpr, _ := utils.IsGTEVersion(versionOnly, "v1.14.0")
+
 		t.Run("fc-"+fcMajor, func(t *testing.T) {
 			buildID := uuid.New().String()
 
 			// Phase 1: create build
-			t.Logf("creating build %s with FC %s", buildID, fcVersion)
+			t.Logf("creating build %s with FC %s (freePageReporting=%v)", buildID, fcVersion, fpr)
 			force := true
 			_, err := infra.builder.Build(
 				ctx,
@@ -88,6 +92,7 @@ func TestSmokeAllFCVersions(t *testing.T) { //nolint:paralleltest // subtests sh
 					MemoryMB:           512,
 					DiskSizeMB:         512,
 					HugePages:          true,
+					FreePageReporting:  fpr,
 					KernelVersion:      featureflags.DefaultKernelVersion,
 					FirecrackerVersion: fcVersion,
 					FromImage:          baseImage,
@@ -111,10 +116,11 @@ func TestSmokeAllFCVersions(t *testing.T) { //nolint:paralleltest // subtests sh
 				ctx,
 				tmpl,
 				sandbox.NewConfig(sandbox.Config{
-					BaseTemplateID: "smoke-" + fcMajor,
-					Vcpu:           2,
-					RamMB:          512,
-					HugePages:      true,
+					BaseTemplateID:    "smoke-" + fcMajor,
+					Vcpu:              2,
+					RamMB:             512,
+					HugePages:         true,
+					FreePageReporting: fpr,
 					Envd: sandbox.EnvdMetadata{
 						Vars:        map[string]string{},
 						AccessToken: &token,

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

@@ -427,6 +427,30 @@ func (c *apiClient) startVM(ctx context.Context) error {
 	return nil
 }
 
+func (c *apiClient) enableFreePageReporting(ctx context.Context) error {
+	ctx, span := tracer.Start(ctx, "enable-free-page-reporting")
+	defer span.End()
+
+	amountMib := int64(0)
+	deflateOnOom := false
+
+	balloonConfig := operations.PutBalloonParams{
+		Context: ctx,
+		Body: &models.Balloon{
+			AmountMib:         &amountMib,
+			DeflateOnOom:      &deflateOnOom,
+			FreePageReporting: true,
+		},
+	}
+
+	_, err := c.client.Operations.PutBalloon(&balloonConfig)
+	if err != nil {
+		return fmt.Errorf("error setting up balloon device: %w", err)
+	}
+
+	return nil
+}
+
 func (c *apiClient) memoryMapping(ctx context.Context) (*memory.Mapping, error) {
 	params := operations.GetMemoryMappingsParams{
 		Context: ctx,

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

@@ -299,6 +299,7 @@ func (p *Process) Create(
 	vCPUCount int64,
 	memoryMB int64,
 	hugePages bool,
+	freePageReporting bool,
 	options ProcessOptions,
 	txRateLimit RateLimiterConfig,
 	driveRateLimit RateLimiterConfig,
@@ -438,6 +439,16 @@ func (p *Process) Create(
 	}
 	telemetry.ReportEvent(ctx, "set fc entropy config")
 
+	if freePageReporting {
+		err = p.client.enableFreePageReporting(ctx)
+		if err != nil {
+			fcStopErr := p.Stop(ctx)
+
+			return errors.Join(fmt.Errorf("error enabling free page reporting: %w", err), fcStopErr)
+		}
+		telemetry.ReportEvent(ctx, "enabled free page reporting")
+	}
+
 	err = p.client.startVM(ctx)
 	if err != nil {
 		fcStopErr := p.Stop(ctx)

packages/orchestrator/pkg/sandbox/sandbox.go

@@ -71,6 +71,9 @@ type Config struct {
 	TotalDiskSizeMB int64
 	HugePages       bool
 
+	// Enable free page reporting
+	FreePageReporting bool
+
 	Envd EnvdMetadata
 
 	FirecrackerConfig fc.Config
@@ -495,6 +498,7 @@ func (f *Factory) CreateSandbox(
 		config.Vcpu,
 		config.RamMB,
 		config.HugePages,
+		config.FreePageReporting,
 		processOptions,
 		fc.RateLimiterConfig{
 			Ops:       fc.TokenBucketConfig(throttleConfig.Ops),