Skip to content

Commit b5c7e29

Browse files
ValentaTomasValentaTomas
committed
feat(uffd): wire free-page-reporting through template build to FC balloon
Adds the FC-side integration plumbing for free page reporting on top of the UFFD REMOVE-event handling in #2520: - template-manager proto: optional bool freePageReporting (field 17). - TemplateConfig + sandbox.Config gain a FreePageReporting bool that flows from template create → build phases (base/steps/finalize) → sandbox factory → fc.Process.Create. - fc.apiClient.enableFreePageReporting calls PUT /balloon with free_page_reporting=true after entropy setup and before VM start. - fcversion.HasFreePageReporting gates rollout to FC v1.14+. - Adds free-page-reporting LaunchDarkly feature flag. - create-build CLI: --free-page-reporting flag, defaults to enabled when FC version supports it. - smoketest: opportunistically enables FPR when the FC version under test supports it. UFFD-side changes (REMOVE handling, page tracker, race tests, fix) remain in #2520; this PR is purely the production rollout path.
1 parent 08e94de commit b5c7e29

14 files changed

Lines changed: 239 additions & 143 deletions

File tree

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

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ func main() {
6363
memory := flag.Int("memory", 1024, "memory MB")
6464
disk := flag.Int("disk", 1024, "disk MB")
6565
hugePages := flag.Bool("hugepages", true, "use 2MB huge pages for memory (false = 4KB pages)")
66+
freePageReporting := flag.Bool("free-page-reporting", false, "enable free page reporting via balloon device (requires Firecracker v1.14+)")
6667
startCmd := flag.String("start-cmd", "", "start command")
6768
setupCmd := flag.String("setup-cmd", "", "setup command to run during build (e.g., install deps)")
6869
readyCmd := flag.String("ready-cmd", "", "ready check command")
@@ -100,7 +101,16 @@ func main() {
100101
log.Fatalf("network config: %v", err)
101102
}
102103

103-
err = doBuild(ctx, *templateID, *toBuild, *fromBuild, *kernel, *fc, *vcpu, *memory, *disk, *hugePages, *startCmd, *setupCmd, *readyCmd, localMode, *verbose, *timeout, builderConfig, networkConfig)
104+
// Detect if --free-page-reporting was explicitly set; if not, pass nil so
105+
// doBuild can default based on the Firecracker version.
106+
var fprOverride *bool
107+
flag.Visit(func(f *flag.Flag) {
108+
if f.Name == "free-page-reporting" {
109+
fprOverride = freePageReporting
110+
}
111+
})
112+
113+
err = doBuild(ctx, *templateID, *toBuild, *fromBuild, *kernel, *fc, *vcpu, *memory, *disk, *hugePages, fprOverride, *startCmd, *setupCmd, *readyCmd, localMode, *verbose, *timeout, builderConfig, networkConfig)
104114
if err != nil {
105115
log.Fatal(err)
106116
}
@@ -199,9 +209,10 @@ func setupEnv(ctx context.Context, storagePath, sandboxDir, kernel, fc string, l
199209

200210
func doBuild(
201211
parentCtx context.Context,
202-
templateID, buildID, fromBuild, kernel, fc string,
212+
templateID, buildID, fromBuild, kernel, fcVersion string,
203213
vcpu, memory, disk int,
204214
hugePages bool,
215+
freePageReporting *bool,
205216
startCmd, setupCmd, readyCmd string,
206217
localMode, verbose bool,
207218
timeout int,
@@ -341,6 +352,18 @@ func doBuild(
341352
})
342353
}
343354

355+
// Default FPR to enabled when the FC version supports it (v1.14+); explicit flag overrides.
356+
var fprEnabled bool
357+
if freePageReporting != nil {
358+
fprEnabled = *freePageReporting
359+
} else {
360+
versionOnly, _, _ := strings.Cut(fcVersion, "_")
361+
supported, err := utils.IsGTEVersion(versionOnly, "v1.14.0")
362+
if err == nil {
363+
fprEnabled = supported
364+
}
365+
}
366+
344367
tmpl := config.TemplateConfig{
345368
Version: templates.TemplateV2LatestVersion,
346369
TemplateID: templateID,
@@ -349,10 +372,11 @@ func doBuild(
349372
MemoryMB: int64(memory),
350373
DiskSizeMB: int64(disk),
351374
HugePages: hugePages,
375+
FreePageReporting: fprEnabled,
352376
StartCmd: startCmd,
353377
ReadyCmd: readyCmd,
354378
KernelVersion: kernel,
355-
FirecrackerVersion: fc,
379+
FirecrackerVersion: fcVersion,
356380
Steps: steps,
357381
}
358382

‎packages/orchestrator/cmd/smoketest/smoke_test.go‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ import (
4040
sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
4141
"github.com/e2b-dev/infra/packages/shared/pkg/storage"
4242
"github.com/e2b-dev/infra/packages/shared/pkg/templates"
43+
"github.com/e2b-dev/infra/packages/shared/pkg/utils"
4344
)
4445

4546
const (
@@ -71,11 +72,14 @@ func TestSmokeAllFCVersions(t *testing.T) { //nolint:paralleltest // subtests sh
7172
defer infra.close(ctx)
7273

7374
for fcMajor, fcVersion := range featureflags.FirecrackerVersionMap { //nolint:paralleltest // sequential by design
75+
versionOnly, _, _ := strings.Cut(fcVersion, "_")
76+
fpr, _ := utils.IsGTEVersion(versionOnly, "v1.14.0")
77+
7478
t.Run("fc-"+fcMajor, func(t *testing.T) {
7579
buildID := uuid.New().String()
7680

7781
// Phase 1: create build
78-
t.Logf("creating build %s with FC %s", buildID, fcVersion)
82+
t.Logf("creating build %s with FC %s (freePageReporting=%v)", buildID, fcVersion, fpr)
7983
force := true
8084
_, err := infra.builder.Build(
8185
ctx,
@@ -88,6 +92,7 @@ func TestSmokeAllFCVersions(t *testing.T) { //nolint:paralleltest // subtests sh
8892
MemoryMB: 512,
8993
DiskSizeMB: 512,
9094
HugePages: true,
95+
FreePageReporting: fpr,
9196
KernelVersion: featureflags.DefaultKernelVersion,
9297
FirecrackerVersion: fcVersion,
9398
FromImage: baseImage,
@@ -111,10 +116,11 @@ func TestSmokeAllFCVersions(t *testing.T) { //nolint:paralleltest // subtests sh
111116
ctx,
112117
tmpl,
113118
sandbox.NewConfig(sandbox.Config{
114-
BaseTemplateID: "smoke-" + fcMajor,
115-
Vcpu: 2,
116-
RamMB: 512,
117-
HugePages: true,
119+
BaseTemplateID: "smoke-" + fcMajor,
120+
Vcpu: 2,
121+
RamMB: 512,
122+
HugePages: true,
123+
FreePageReporting: fpr,
118124
Envd: sandbox.EnvdMetadata{
119125
Vars: map[string]string{},
120126
AccessToken: &token,

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -427,6 +427,30 @@ func (c *apiClient) startVM(ctx context.Context) error {
427427
return nil
428428
}
429429

430+
func (c *apiClient) enableFreePageReporting(ctx context.Context) error {
431+
ctx, span := tracer.Start(ctx, "enable-free-page-reporting")
432+
defer span.End()
433+
434+
amountMib := int64(0)
435+
deflateOnOom := false
436+
437+
balloonConfig := operations.PutBalloonParams{
438+
Context: ctx,
439+
Body: &models.Balloon{
440+
AmountMib: &amountMib,
441+
DeflateOnOom: &deflateOnOom,
442+
FreePageReporting: true,
443+
},
444+
}
445+
446+
_, err := c.client.Operations.PutBalloon(&balloonConfig)
447+
if err != nil {
448+
return fmt.Errorf("error setting up balloon device: %w", err)
449+
}
450+
451+
return nil
452+
}
453+
430454
func (c *apiClient) memoryMapping(ctx context.Context) (*memory.Mapping, error) {
431455
params := operations.GetMemoryMappingsParams{
432456
Context: ctx,

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,7 @@ func (p *Process) Create(
299299
vCPUCount int64,
300300
memoryMB int64,
301301
hugePages bool,
302+
freePageReporting bool,
302303
options ProcessOptions,
303304
txRateLimit RateLimiterConfig,
304305
driveRateLimit RateLimiterConfig,
@@ -438,6 +439,16 @@ func (p *Process) Create(
438439
}
439440
telemetry.ReportEvent(ctx, "set fc entropy config")
440441

442+
if freePageReporting {
443+
err = p.client.enableFreePageReporting(ctx)
444+
if err != nil {
445+
fcStopErr := p.Stop(ctx)
446+
447+
return errors.Join(fmt.Errorf("error enabling free page reporting: %w", err), fcStopErr)
448+
}
449+
telemetry.ReportEvent(ctx, "enabled free page reporting")
450+
}
451+
441452
err = p.client.startVM(ctx)
442453
if err != nil {
443454
fcStopErr := p.Stop(ctx)

‎packages/orchestrator/pkg/sandbox/sandbox.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,9 @@ type Config struct {
7171
TotalDiskSizeMB int64
7272
HugePages bool
7373

74+
// Enable free page reporting
75+
FreePageReporting bool
76+
7477
Envd EnvdMetadata
7578

7679
FirecrackerConfig fc.Config
@@ -495,6 +498,7 @@ func (f *Factory) CreateSandbox(
495498
config.Vcpu,
496499
config.RamMB,
497500
config.HugePages,
501+
config.FreePageReporting,
498502
processOptions,
499503
fc.RateLimiterConfig{
500504
Ops: fc.TokenBucketConfig(throttleConfig.Ops),

‎packages/orchestrator/pkg/template/build/config/config.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,9 @@ type TemplateConfig struct {
4141
// HugePages sets whether the VM use huge pages.
4242
HugePages bool
4343

44+
// FreePageReporting enables the corresponding feature in Firecracker
45+
FreePageReporting bool
46+
4447
// Command to run to check if the template is ready.
4548
ReadyCmd string
4649

‎packages/orchestrator/pkg/template/build/phases/base/builder.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -200,9 +200,10 @@ func (bb *BaseBuilder) buildLayerFromOCI(
200200

201201
// Allow sandbox internet access during provisioning (nil network = no restrictions).
202202
baseSbxConfig := sandbox.NewConfig(sandbox.Config{
203-
Vcpu: bb.Config.VCpuCount,
204-
RamMB: bb.Config.MemoryMB,
205-
HugePages: bb.Config.HugePages,
203+
Vcpu: bb.Config.VCpuCount,
204+
RamMB: bb.Config.MemoryMB,
205+
HugePages: bb.Config.HugePages,
206+
FreePageReporting: bb.Config.FreePageReporting,
206207

207208
Envd: sandbox.EnvdMetadata{
208209
Version: bb.EnvdVersion,

‎packages/orchestrator/pkg/template/build/phases/finalize/builder.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -149,9 +149,10 @@ func (ppb *PostProcessingBuilder) Build(
149149

150150
// Configure sandbox for final layer
151151
sbxConfig := sandbox.NewConfig(sandbox.Config{
152-
Vcpu: ppb.Config.VCpuCount,
153-
RamMB: ppb.Config.MemoryMB,
154-
HugePages: ppb.Config.HugePages,
152+
Vcpu: ppb.Config.VCpuCount,
153+
RamMB: ppb.Config.MemoryMB,
154+
HugePages: ppb.Config.HugePages,
155+
FreePageReporting: ppb.Config.FreePageReporting,
155156

156157
Envd: sandbox.EnvdMetadata{
157158
Version: ppb.EnvdVersion,

‎packages/orchestrator/pkg/template/build/phases/steps/builder.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -163,9 +163,10 @@ func (sb *StepBuilder) Build(
163163
step := sb.step
164164

165165
sbxConfig := sandbox.NewConfig(sandbox.Config{
166-
Vcpu: sb.Config.VCpuCount,
167-
RamMB: sb.Config.MemoryMB,
168-
HugePages: sb.Config.HugePages,
166+
Vcpu: sb.Config.VCpuCount,
167+
RamMB: sb.Config.MemoryMB,
168+
HugePages: sb.Config.HugePages,
169+
FreePageReporting: sb.Config.FreePageReporting,
169170

170171
Envd: sandbox.EnvdMetadata{
171172
Version: sb.EnvdVersion,

‎packages/orchestrator/pkg/template/server/create_template.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ func (s *ServerStore) TemplateCreate(ctx context.Context, templateRequest *templ
8888
ReadyCmd: cfg.GetReadyCommand(),
8989
DiskSizeMB: int64(cfg.GetDiskSizeMB()),
9090
HugePages: hugePages,
91+
FreePageReporting: cfg.GetFreePageReporting(),
9192
FromImage: cfg.GetFromImage(),
9293
FromTemplate: cfg.GetFromTemplate(),
9394
RegistryAuthProvider: authProvider,

0 commit comments

Comments
 (0)