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

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

@@ -69,11 +69,15 @@ func main() {
 	optimize := flag.Bool("optimize", false, "collect fresh prefetch mapping after pause (resumes snapshot to record page faults)")
 	shell := flag.Bool("shell", false, "attach an interactive PTY shell via envd (no sshd required in the sandbox)")
 
-	// Enables the pre-pause reclaim chain with sensible per-step caps.
+	fphTimeoutMs := flag.Int("fph-timeout-ms", 0, "override free-page-hinting-timeout-ms LD flag (0 = use LD default)")
 	reclaim := flag.Bool("reclaim", false, "enable pre-pause reclaim chain (fstrim 500ms, sync 500ms, drop_caches 200ms, compact 1s)")
 
 	flag.Parse()
 
+	if *fphTimeoutMs > 0 {
+		featureflags.NewIntFlag("free-page-hinting-timeout-ms", *fphTimeoutMs)
+	}
+
 	if *reclaim {
 		featureflags.NewJSONFlag("guest-pause-reclaim", ldvalue.FromJSONMarshal(map[string]int{
 			"sync":           500,

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

@@ -428,9 +428,9 @@ func (c *apiClient) startVM(ctx context.Context) error {
 }
 
 // installBalloon attaches a zero-MiB balloon device. Individual balloon
-// features (free-page-reporting today, free-page-hinting next) are toggled
-// via parameters so callers can opt in to any subset independently.
-func (c *apiClient) installBalloon(ctx context.Context, freePageReporting bool) error {
+// features (free-page-reporting, free-page-hinting) are toggled via
+// parameters so callers can opt in to any subset independently.
+func (c *apiClient) installBalloon(ctx context.Context, freePageReporting, freePageHinting bool) error {
 	ctx, span := tracer.Start(ctx, "install-balloon")
 	defer span.End()
 
@@ -443,6 +443,7 @@ func (c *apiClient) installBalloon(ctx context.Context, freePageReporting bool)
 			AmountMib:         &amountMib,
 			DeflateOnOom:      &deflateOnOom,
 			FreePageReporting: freePageReporting,
+			FreePageHinting:   freePageHinting,
 		},
 	}
 
@@ -454,6 +455,33 @@ func (c *apiClient) installBalloon(ctx context.Context, freePageReporting bool)
 	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/fph_gates.go

@@ -0,0 +1,17 @@
+package fc
+
+import (
+	"github.com/e2b-dev/infra/packages/shared/pkg/fcversion"
+)
+
+// FCSupportsFreePageHinting reports whether the FC version's API exposes
+// virtio-balloon free-page-hinting. Kernel-side eligibility (and the race-fix
+// requirement) is targeted via LaunchDarkly with kernel-version context.
+func FCSupportsFreePageHinting(fcVersion string) bool {
+	info, err := fcversion.New(fcVersion)
+	if err != nil {
+		return false
+	}
+
+	return info.HasFreePageHinting()
+}

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"
@@ -300,6 +301,7 @@ func (p *Process) Create(
 	memoryMB int64,
 	hugePages bool,
 	freePageReporting bool,
+	freePageHinting bool,
 	options ProcessOptions,
 	txRateLimit RateLimiterConfig,
 	driveRateLimit RateLimiterConfig,
@@ -442,14 +444,16 @@ func (p *Process) Create(
 	}
 	telemetry.ReportEvent(ctx, "set fc entropy config")
 
-	if freePageReporting {
-		err = p.client.installBalloon(ctx, freePageReporting)
-		if err != nil {
+	if freePageReporting || freePageHinting {
+		if err := p.client.installBalloon(ctx, freePageReporting, freePageHinting); err != nil {
 			fcStopErr := p.Stop(ctx)
 
 			return errors.Join(fmt.Errorf("error installing balloon device: %w", err), fcStopErr)
 		}
-		telemetry.ReportEvent(ctx, "installed balloon device")
+		telemetry.ReportEvent(ctx, "installed balloon device",
+			attribute.Bool("balloon.free_page_reporting", freePageReporting),
+			attribute.Bool("balloon.free_page_hinting", freePageHinting),
+		)
 	}
 
 	err = p.client.startVM(ctx)
@@ -713,6 +717,80 @@ 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 on FC < v1.14 and when no balloon is
+// configured (FC returns 400).
+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 !FCSupportsFreePageHinting(p.Versions.FirecrackerVersion) {
+		outcome = "fc-unsupported"
+
+		return nil
+	}
+
+	// Snapshot the host_cmd before starting so we can detect a true bump
+	// from this drain rather than reading stale counters carried over from a
+	// previous drain that was snapshotted then restored. host_cmd is monotonic.
+	hostBefore, _, err := p.client.describeBalloonHinting(ctx)
+	if err != nil {
+		var notConfigured *operations.DescribeBalloonHintingBadRequest
+		if errors.As(err, &notConfigured) {
+			outcome = "not-configured"
+
+			return nil
+		}
+
+		outcome = "describe-failed"
+
+		return fmt.Errorf("balloon hinting baseline: %w", err)
+	}
+
+	if err := p.client.startBalloonHinting(ctx, true); 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)
+		}
+		// Require a strict bump over the pre-start baseline to avoid
+		// a false-positive completion before FC's VMM thread has accepted
+		// the start (the API thread ack/VMM bump race), including the
+		// resume-from-snapshot case where counters restore non-zero.
+		if host > hostBefore && 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

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/launchdarkly/go-sdk-common/v3/ldcontext"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
@@ -167,6 +168,17 @@ type RuntimeMetadata struct {
 	SandboxType SandboxType
 }
 
+// sandboxLDContext builds an LD context with kernel/FC-version attributes so
+// flags evaluated for a sandbox can be targeted by guest kernel or FC version.
+func sandboxLDContext(runtime RuntimeMetadata, config *Config) ldcontext.Context {
+	return ldcontext.NewBuilder(runtime.SandboxID).
+		Kind(featureflags.SandboxKind).
+		SetString(featureflags.SandboxTemplateAttribute, runtime.TemplateID).
+		SetString(featureflags.SandboxKernelVersionAttribute, config.FirecrackerConfig.KernelVersion).
+		SetString(featureflags.SandboxFirecrackerVersionAttribute, config.FirecrackerConfig.FirecrackerVersion).
+		Build()
+}
+
 type Resources struct {
 	Slot   *network.Slot
 	rootfs rootfs.Provider
@@ -489,6 +501,9 @@ func (f *Factory) CreateSandbox(
 		return nil
 	})
 
+	freePageHinting := fc.FCSupportsFreePageHinting(config.FirecrackerConfig.FirecrackerVersion) &&
+		f.featureFlags.BoolFlag(ctx, featureflags.FreePageHintingInstallFlag, sandboxLDContext(runtime, config))
+
 	err = fcHandle.Create(
 		ctx,
 		sbxlogger.SandboxMetadata{
@@ -500,6 +515,7 @@ func (f *Factory) CreateSandbox(
 		config.RamMB,
 		config.HugePages,
 		config.FreePageReporting,
+		freePageHinting,
 		processOptions,
 		fc.RateLimiterConfig{
 			Ops:       fc.TokenBucketConfig(throttleConfig.Ops),
@@ -1062,6 +1078,18 @@ func (s *Sandbox) Pause(
 	// all default to 0 which disables the chain entirely. Non-fatal.
 	s.bestEffortReclaim(ctx)
 
+	// Drain free-page-hinting before pause so the snapshot doesn't capture
+	// pages the guest already considers free. Timeout=0 disables. Evaluated
+	// with kernel-version LD context so operators can roll out only on guests
+	// with the kernel race fix.
+	if t := time.Duration(s.featureFlags.IntFlag(ctx, featureflags.FreePageHintingTimeoutMs, sandboxLDContext(s.Runtime, s.Config))) * 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/fcversion/sandbox_features.go

@@ -11,3 +11,7 @@ func (v *Info) HasHugePages() bool {
 func (v *Info) HasFreePageReporting() bool {
 	return v.lastReleaseVersion.Major() > 1 || (v.lastReleaseVersion.Major() == 1 && v.lastReleaseVersion.Minor() >= 14)
 }
+
+func (v *Info) HasFreePageHinting() bool {
+	return v.lastReleaseVersion.Major() > 1 || (v.lastReleaseVersion.Major() == 1 && v.lastReleaseVersion.Minor() >= 14)
+}

packages/shared/pkg/featureflags/flags.go

@@ -124,6 +124,12 @@ var (
 	SandboxLabelBasedSchedulingFlag  = NewBoolFlag("sandbox-label-based-scheduling", false)
 	OptimisticResourceAccountingFlag = NewBoolFlag("sandbox-placement-optimistic-resource-accounting", false)
 	FreePageReportingFlag            = NewBoolFlag("free-page-reporting", false)
+	// FreePageHintingInstallFlag controls whether FreePageHinting=true is
+	// configured on the balloon at install time. Just having FPH on the
+	// balloon doesn't trigger the kernel race fixed in
+	// https://lore.kernel.org/lkml/20240429125100.7393-1-david@redhat.com/
+	// — that race is on the actual hinting flow, gated by FreePageHintingTimeoutMs.
+	FreePageHintingInstallFlag = NewBoolFlag("free-page-hinting-install", false)
 
 	NetworkTransformRulesFlag = NewBoolFlag("network-transform-rules", env.IsDevelopment())
 )
@@ -164,7 +170,11 @@ 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 the pre-pause balloon FPH drain. 0
+	// disables. Evaluated with sandbox/kernel-version LD context so operators
+	// can roll out only on guests with the kernel 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