feat(fc): drain virtio-balloon free-page-hinting before pause

Arm free-page-hinting on the existing balloon device (always set when
the balloon is installed; pure runtime toggle), and on pause do a
host-initiated hint+wait so MADV_DONTNEED-reclaimed pages are settled
before the snapshot. Pages reclaimed this way generate UFFD_EVENT_REMOVE,
which the orchestrator already tracks (parent FPR PR), so the snapshot
captures them as removed instead of zero-filled.

- fc/client.go: rename enableFreePageReporting -> installBalloon;
  always set FreePageHinting=true; add startBalloonHinting +
  describeBalloonHinting helpers.
- fc/process.go: track balloonInstalled; add DrainBalloon (start +
  poll guest_cmd >= host_cmd, with host>0 guard against transient
  nil/zero responses).
- sandbox.go: wire featureFlags into Sandbox; call DrainBalloon from
  Pause behind the flag. Failures are logged but non-fatal.

Gated by free-page-hinting-timeout-ms (LD int flag, ms; default 0 =
disabled). resume-build gains --fph-timeout-ms for local exercise.

Author: ValentaTomas

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

@@ -72,8 +72,15 @@ func main() {
 	cmdSignalPause := flag.String("cmd-signal-pause", "", "execute command in sandbox, then wait for SIGUSR1 before pausing")
 	optimize := flag.Bool("optimize", false, "collect fresh prefetch mapping after pause (resumes snapshot to record page faults)")
 
+	// Pause-time FPH override; 0 = use LD default (off).
+	fphTimeoutMs := flag.Int("fph-timeout-ms", 0, "override free-page-hinting-timeout-ms LD flag (0 = use LD default)")
+
 	flag.Parse()
 
+	if *fphTimeoutMs > 0 {
+		featureflags.NewIntFlag("free-page-hinting-timeout-ms", *fphTimeoutMs)
+	}
+
 	if *fromBuild == "" {
 		log.Fatal("-from-build required")
 	}

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

@@ -427,8 +427,11 @@ 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")
+// installBalloon installs the virtio-balloon pre-boot with target size 0.
+// FreePageHinting is always armed (pure runtime toggle, used by DrainBalloon);
+// FreePageReporting is set per template-build gate.
+func (c *apiClient) installBalloon(ctx context.Context, freePageReporting bool) error {
+	ctx, span := tracer.Start(ctx, "install-balloon")
 	defer span.End()
 
 	amountMib := int64(0)
@@ -439,7 +442,8 @@ func (c *apiClient) enableFreePageReporting(ctx context.Context) error {
 		Body: &models.Balloon{
 			AmountMib:         &amountMib,
 			DeflateOnOom:      &deflateOnOom,
-			FreePageReporting: true,
+			FreePageReporting: freePageReporting,
+			FreePageHinting:   true,
 		},
 	}
 
@@ -451,6 +455,33 @@ func (c *apiClient) enableFreePageReporting(ctx context.Context) error {
 	return nil
 }
 
+func (c *apiClient) startBalloonHinting(ctx context.Context, acknowledgeOnStop bool) error {
+	params := operations.StartBalloonHintingParams{
+		Context: ctx,
+		Body:    &models.BalloonStartCmd{AcknowledgeOnStop: acknowledgeOnStop},
+	}
+	_, err := c.client.Operations.StartBalloonHinting(&params)
+	if err != nil {
+		return fmt.Errorf("error starting balloon hinting: %w", err)
+	}
+
+	return nil
+}
+
+func (c *apiClient) describeBalloonHinting(ctx context.Context) (hostCmd, guestCmd int64, err error) {
+	params := operations.DescribeBalloonHintingParams{Context: ctx}
+	res, err := c.client.Operations.DescribeBalloonHinting(&params)
+	if err != nil {
+		return 0, 0, err
+	}
+	if res.Payload.HostCmd != nil {
+		hostCmd = *res.Payload.HostCmd
+	}
+	guestCmd = res.Payload.GuestCmd
+
+	return hostCmd, guestCmd, nil
+}
+
 func (c *apiClient) memoryMapping(ctx context.Context) (*memory.Mapping, error) {
 	params := operations.GetMemoryMappingsParams{
 		Context: ctx,

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

@@ -26,6 +26,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/rootfs"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/socket"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/template"
+	"github.com/e2b-dev/infra/packages/shared/pkg/fc/client/operations"
 	"github.com/e2b-dev/infra/packages/shared/pkg/keys"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
@@ -440,13 +441,12 @@ func (p *Process) Create(
 	telemetry.ReportEvent(ctx, "set fc entropy config")
 
 	if freePageReporting {
-		err = p.client.enableFreePageReporting(ctx)
-		if err != nil {
+		if err := p.client.installBalloon(ctx, freePageReporting); err != nil {
 			fcStopErr := p.Stop(ctx)
 
-			return errors.Join(fmt.Errorf("error enabling free page reporting: %w", err), fcStopErr)
+			return errors.Join(fmt.Errorf("error installing balloon device: %w", err), fcStopErr)
 		}
-		telemetry.ReportEvent(ctx, "enabled free page reporting")
+		telemetry.ReportEvent(ctx, "installed balloon device")
 	}
 
 	err = p.client.startVM(ctx)
@@ -710,6 +710,58 @@ func (p *Process) Pause(ctx context.Context) error {
 	return p.client.pauseVM(ctx)
 }
 
+// DrainBalloon triggers a free-page-hinting run and blocks until the guest
+// acknowledges or ctx fires. No-op when no balloon device is configured
+// (FC returns 400) so it survives snapshot/resume without local state.
+// Records `drain-balloon.outcome` on the span: ok | not-configured | timeout |
+// start-failed | describe-failed.
+func (p *Process) DrainBalloon(ctx context.Context) error {
+	ctx, span := tracer.Start(ctx, "drain-balloon")
+	outcome := "ok"
+	defer func() {
+		span.SetAttributes(attribute.String("drain-balloon.outcome", outcome))
+		span.End()
+	}()
+
+	if err := p.client.startBalloonHinting(ctx, true /* ackOnStop */); err != nil {
+		var notConfigured *operations.StartBalloonHintingBadRequest
+		if errors.As(err, &notConfigured) {
+			outcome = "not-configured"
+
+			return nil
+		}
+
+		outcome = "start-failed"
+
+		return fmt.Errorf("start balloon hinting: %w", err)
+	}
+
+	backoff := 5 * time.Millisecond
+	for {
+		select {
+		case <-ctx.Done():
+			outcome = "timeout"
+
+			return ctx.Err()
+		case <-time.After(backoff):
+		}
+
+		host, guest, err := p.client.describeBalloonHinting(ctx)
+		if err != nil {
+			outcome = "describe-failed"
+
+			return fmt.Errorf("balloon hinting status: %w", err)
+		}
+		// host_cmd is monotonic and we just called start, so host > 0
+		// after FC accepts it. Require it to guard against transient
+		// nil/zero responses returning a false-positive completion.
+		if host > 0 && guest >= host {
+			return nil
+		}
+		backoff = min(backoff*2, 50*time.Millisecond)
+	}
+}
+
 // CreateSnapshot VM needs to be paused before creating a snapshot.
 func (p *Process) CreateSnapshot(ctx context.Context, snapfilePath string) error {
 	ctx, childSpan := tracer.Start(ctx, "create-snapshot-fc")

packages/orchestrator/pkg/sandbox/sandbox.go

@@ -218,6 +218,8 @@ type Sandbox struct {
 	files   *storage.SandboxFiles
 	cleanup *Cleanup
 
+	featureFlags *featureflags.Client
+
 	process      *fc.Process
 	cgroupHandle *cgroup.CgroupHandle
 
@@ -458,7 +460,8 @@ func (f *Factory) CreateSandbox(
 		files:    sandboxFiles,
 		process:  fcHandle,
 
-		cleanup: cleanup,
+		cleanup:      cleanup,
+		featureFlags: f.featureFlags,
 
 		APIStoredConfig: apiConfigToStore,
 
@@ -799,7 +802,8 @@ func (f *Factory) ResumeSandbox(
 		files:    sandboxFiles,
 		process:  fcHandle,
 
-		cleanup: cleanup,
+		cleanup:      cleanup,
+		featureFlags: f.featureFlags,
 
 		APIStoredConfig: apiConfigToStore,
 		CABundle:        f.egressProxy.CABundle(),
@@ -1053,6 +1057,17 @@ func (s *Sandbox) Pause(
 	// Stop the health check before pausing the VM
 	s.Checks.Stop()
 
+	// Drain free-page-hinting before pause so the snapshot doesn't capture
+	// pages the guest already considers free. No-op when no balloon. Failures
+	// are logged but non-fatal. Timeout=0 disables the step.
+	if t := time.Duration(s.featureFlags.IntFlag(ctx, featureflags.FreePageHintingTimeoutMs)) * time.Millisecond; t > 0 {
+		drainCtx, cancel := context.WithTimeout(ctx, t)
+		if err := s.process.DrainBalloon(drainCtx); err != nil {
+			telemetry.ReportError(ctx, "balloon hinting drain failed (continuing pause)", err)
+		}
+		cancel()
+	}
+
 	if err := s.process.Pause(ctx); err != nil {
 		return nil, fmt.Errorf("failed to pause VM: %w", err)
 	}

packages/shared/pkg/featureflags/flags.go

@@ -158,7 +158,10 @@ var (
 	BestOfKMaxOvercommit          = NewIntFlag("best-of-k-max-overcommit", 400)              // Default R=4 (stored as percentage, max over-commit ratio)
 	BestOfKAlpha                  = NewIntFlag("best-of-k-alpha", 50)                        // Default Alpha=0.5 (stored as percentage for int flag, current usage weight)
 	EnvdInitTimeoutMilliseconds   = NewIntFlag("envd-init-request-timeout-milliseconds", 50) // Timeout for envd init request in milliseconds
-	HostStatsSamplingInterval     = NewIntFlag("host-stats-sampling-interval", 5000)         // Host stats sampling interval in milliseconds (default 5s)
+	// FreePageHintingTimeoutMs gates a pre-pause virtio-balloon free-page-hinting
+	// drain. 0 disables it. Operator opts in once the kernel has the FPH race fix.
+	FreePageHintingTimeoutMs      = NewIntFlag("free-page-hinting-timeout-ms", 0)
+	HostStatsSamplingInterval     = NewIntFlag("host-stats-sampling-interval", 5000) // Host stats sampling interval in milliseconds (default 5s)
 	MaxCacheWriterConcurrencyFlag = NewIntFlag("max-cache-writer-concurrency", 10)
 
 	// BuildCacheMaxUsagePercentage the maximum percentage of the cache disk storage