feat(orch): collapse envd heap before pause behind a feature flag

Mirror the user-cgroup-freeze pre-pause path: when the new collapse-envd-heap
LaunchDarkly flag is set, the orchestrator calls envd's POST /collapse just
before pause so envd consolidates its scattered heap into hugepages and faults
fewer distinct frames on resume. Gate on envd version (>= 0.6.4, which exposes
the endpoint) with a dedicated 10s timeout, and keep it best-effort so a failure
never blocks pause. The flag defaults off and is rolled out via LaunchDarkly.

Observe the collapse on a dedicated envd-collapse span and two metrics: the
orchestrator.sandbox.envd.collapse.duration histogram (round-trip ms, with a
success attribute) for the pause-path cost, and the
orchestrator.sandbox.envd.collapse.chunks counter split by result
(collapsed|skipped) so chunk attempts (total) and successes (collapsed) are both
queryable. The per-call stats are attached as span attributes and logged on
success ("envd heap collapsed") so efficacy is visible in Loki when the metric
counters reset on restart. Add a resume-build -collapse flag that overrides the
LD flag on for local testing.

Signed-off-by: Nikita Kalyazin <nikita.kalyazin@e2b.dev>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: kalyazin

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

@@ -75,6 +75,7 @@ func main() {
 
 	fphTimeoutMs := flag.Int("fph-timeout-ms", 0, "override free-page-hinting-config pause timeout 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)")
+	collapse := flag.Bool("collapse", false, "collapse envd's heap before pause (overrides the collapse-envd-heap flag)")
 
 	fphBench := flag.Bool("fph-bench", false, "compare pause memfile size with vs without FPH; requires -cmd-pause workload, uses -iterations (default 3), forces FPR on")
 	fphBenchDelay := flag.Duration("fph-bench-delay", 0, "wait this long between workload completion and pause (lets FPR settle)")
@@ -101,6 +102,10 @@ func main() {
 		featureflags.OverrideBoolFlag(featureflags.UseMemFdFlag, true)
 	}
 
+	if *collapse {
+		featureflags.OverrideBoolFlag(featureflags.CollapseEnvdHeapFlag, true)
+	}
+
 	if *memfileDiffDedup {
 		featureflags.OverrideJSONFlag(featureflags.MemfileDiffDedupFlag, ldvalue.FromJSONMarshal(map[string]any{
 			"enabled": true,

packages/orchestrator/pkg/sandbox/envd.go

@@ -108,28 +108,72 @@ func (s *Sandbox) callEnvdUnfreeze(ctx context.Context, timeout time.Duration) e
 }
 
 func (s *Sandbox) callEnvdCgroupOp(ctx context.Context, timeout time.Duration, op envdCgroupOp) error {
+	return s.postEnvd(ctx, timeout, string(op))
+}
+
+// callEnvdCollapse calls envd's native POST /collapse endpoint, which compacts
+// envd's own anonymous heap into 2 MiB hugepages before pause so it faults
+// fewer distinct frames on resume. Unlike freeze/unfreeze it returns a body:
+// the per-call collapse stats, which the caller records as metrics and span
+// attributes.
+func (s *Sandbox) callEnvdCollapse(ctx context.Context, timeout time.Duration) (envd.CollapseResult, error) {
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	address := fmt.Sprintf("http://%s:%d/collapse", s.Slot.HostIPString(), consts.DefaultEnvdServerPort)
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, address, nil)
+	if err != nil {
+		return envd.CollapseResult{}, fmt.Errorf("build collapse request: %w", err)
+	}
+	if s.Config.Envd.AccessToken != nil {
+		req.Header.Set("X-Access-Token", *s.Config.Envd.AccessToken)
+	}
+
+	resp, err := sandboxHttpClient.Do(req)
+	if err != nil {
+		return envd.CollapseResult{}, fmt.Errorf("collapse request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+
+		return envd.CollapseResult{}, fmt.Errorf("collapse returned %d: %s", resp.StatusCode, utils.Truncate(string(body), 100))
+	}
+
+	var result envd.CollapseResult
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return envd.CollapseResult{}, fmt.Errorf("decode collapse result: %w", err)
+	}
+
+	return result, nil
+}
+
+// postEnvd issues an authenticated POST to envd's /<path> endpoint with a tight,
+// dedicated deadline and expects 204 No Content.
+func (s *Sandbox) postEnvd(ctx context.Context, timeout time.Duration, path string) error {
 	ctx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
-	address := fmt.Sprintf("http://%s:%d/%s", s.Slot.HostIPString(), consts.DefaultEnvdServerPort, op)
+	address := fmt.Sprintf("http://%s:%d/%s", s.Slot.HostIPString(), consts.DefaultEnvdServerPort, path)
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, address, nil)
 	if err != nil {
-		return fmt.Errorf("build %s request: %w", op, err)
+		return fmt.Errorf("build %s request: %w", path, err)
 	}
 	if s.Config.Envd.AccessToken != nil {
 		req.Header.Set("X-Access-Token", *s.Config.Envd.AccessToken)
 	}
 
 	resp, err := sandboxHttpClient.Do(req)
 	if err != nil {
-		return fmt.Errorf("%s request: %w", op, err)
+		return fmt.Errorf("%s request: %w", path, err)
 	}
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusNoContent {
 		body, _ := io.ReadAll(resp.Body)
 
-		return fmt.Errorf("%s returned %d: %s", op, resp.StatusCode, utils.Truncate(string(body), 100))
+		return fmt.Errorf("%s returned %d: %s", path, resp.StatusCode, utils.Truncate(string(body), 100))
 	}
 
 	return nil

packages/orchestrator/pkg/sandbox/envd/envd.gen.go

@@ -8,6 +8,9 @@ import (
 	"strings"
 	"time"
 
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/metric"
 	"go.uber.org/zap"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
@@ -23,6 +26,12 @@ const reclaimOuterSlack = 500 * time.Millisecond
 // must stay independent of the reclaim shell deadline.
 const freezeTimeout = 2 * time.Second
 
+// collapseTimeout bounds the native POST /collapse call. Collapsing migrates
+// envd's scattered heap pages into hugepages, which is heavier than the freeze
+// sysfs write, so it gets a larger, independent budget. Best-effort: if it is
+// cut short, the partially-collapsed heap is still an improvement.
+const collapseTimeout = 10 * time.Second
+
 // buildReclaimScript builds the fstrim/sync/drop_caches/compact_memory chain.
 // Returns ("", 0) when every step is disabled.
 func (s *Sandbox) buildReclaimScript(cfg featureflags.ReclaimConfig) (string, time.Duration) {
@@ -76,6 +85,14 @@ func (s *Sandbox) bestEffortReclaim(ctx context.Context) {
 		s.bestEffortFreeze(ctx)
 	}
 
+	if s.featureFlags.BoolFlag(ctx, featureflags.CollapseEnvdHeapFlag,
+		featureflags.SandboxContext(s.Runtime.SandboxID),
+		featureflags.TeamContext(s.Runtime.TeamID),
+		featureflags.TemplateContext(s.Runtime.TemplateID),
+	) {
+		s.bestEffortCollapse(ctx)
+	}
+
 	script, timeout := s.buildReclaimScript(cfg)
 	if script == "" {
 		return
@@ -122,6 +139,73 @@ func (s *Sandbox) envdSupportsCgroupFreeze(ctx context.Context) bool {
 	return ok
 }
 
+// envdSupportsHeapCollapse reports whether the sandbox's envd exposes the native
+// /collapse endpoint. Bad version strings log and return false so we never call
+// an unsupported endpoint.
+func (s *Sandbox) envdSupportsHeapCollapse(ctx context.Context) bool {
+	ok, err := utils.IsGTEVersion(s.Config.Envd.Version, utils.MinEnvdVersionForHeapCollapse)
+	if err != nil {
+		logger.L().Warn(ctx, "heap collapse version gate: bad envd version", logger.WithSandboxID(s.Runtime.SandboxID), zap.String("envd_version", s.Config.Envd.Version), zap.Error(err))
+
+		return false
+	}
+
+	return ok
+}
+
+// bestEffortCollapse asks envd to collapse its own heap into hugepages before
+// pause, so on resume envd touches fewer distinct frames. Gated on envd version;
+// failures are logged but never block pause.
+func (s *Sandbox) bestEffortCollapse(ctx context.Context) {
+	if !s.envdSupportsHeapCollapse(ctx) {
+		return
+	}
+
+	ctx, span := tracer.Start(ctx, "envd-collapse")
+	defer span.End()
+
+	start := time.Now()
+	stats, err := s.callEnvdCollapse(ctx, collapseTimeout)
+	elapsedMs := time.Since(start).Milliseconds()
+	success := err == nil
+
+	// Record the round-trip duration whether or not it succeeded: a timed-out or
+	// failed collapse still spends time on the pause path and must be visible.
+	envdCollapseDurationHistogram.Record(ctx, elapsedMs, metric.WithAttributes(attribute.Bool("success", success)))
+	span.SetAttributes(
+		attribute.Bool("collapse.success", success),
+		attribute.Int64("collapse.duration_ms", elapsedMs),
+	)
+
+	if err != nil {
+		span.SetStatus(codes.Error, err.Error())
+		logger.L().Warn(ctx, "envd heap collapse failed", logger.WithSandboxID(s.Runtime.SandboxID), zap.Error(err))
+
+		return
+	}
+
+	// Chunk-level efficacy: attempts = collapsed + skipped, successful = collapsed.
+	// One counter split by result so attempts (sum) and successes (collapsed) are
+	// both queryable.
+	envdCollapseChunks.Add(ctx, int64(stats.Collapsed), metric.WithAttributes(attribute.String("result", "collapsed")))
+	envdCollapseChunks.Add(ctx, int64(stats.Skipped), metric.WithAttributes(attribute.String("result", "skipped")))
+	span.SetAttributes(
+		attribute.Int("collapse.regions", stats.Regions),
+		attribute.Int("collapse.chunks", stats.Chunks),
+		attribute.Int("collapse.collapsed", stats.Collapsed),
+		attribute.Int("collapse.skipped", stats.Skipped),
+	)
+
+	logger.L().Info(ctx, "envd heap collapsed",
+		logger.WithSandboxID(s.Runtime.SandboxID),
+		zap.Int("regions", stats.Regions),
+		zap.Int("chunks", stats.Chunks),
+		zap.Int("collapsed", stats.Collapsed),
+		zap.Int("skipped", stats.Skipped),
+		zap.Int64("duration_ms", elapsedMs),
+	)
+}
+
 // bestEffortFreeze calls envd's native /freeze endpoint with a tight, freeze-
 // only deadline so it doesn't share a timeout budget with the rest of reclaim.
 // Gated on envd version; failures are logged but never block pause.

packages/orchestrator/pkg/sandbox/reclaim.go

@@ -45,9 +45,11 @@ import (
 )
 
 var (
-	meter                        = otel.Meter("github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox")
-	envdInitCalls                = utils.Must(telemetry.GetCounter(meter, telemetry.EnvdInitCalls))
-	waitForEnvdDurationHistogram = utils.Must(telemetry.GetHistogram(meter, telemetry.WaitForEnvdDurationHistogramName))
+	meter                         = otel.Meter("github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox")
+	envdInitCalls                 = utils.Must(telemetry.GetCounter(meter, telemetry.EnvdInitCalls))
+	waitForEnvdDurationHistogram  = utils.Must(telemetry.GetHistogram(meter, telemetry.WaitForEnvdDurationHistogramName))
+	envdCollapseDurationHistogram = utils.Must(telemetry.GetHistogram(meter, telemetry.EnvdCollapseDurationHistogramName))
+	envdCollapseChunks            = utils.Must(telemetry.GetCounter(meter, telemetry.EnvdCollapseChunks))
 
 	uffdStartupPagesHistogram       = utils.Must(telemetry.GetHistogram(meter, telemetry.UffdStartupPagesHistogramName))
 	uffdStartupSourcePagesHistogram = utils.Must(telemetry.GetHistogram(meter, telemetry.UffdStartupSourcePagesHistogramName))

packages/orchestrator/pkg/sandbox/sandbox.go

@@ -160,6 +160,10 @@ var (
 	OptimisticResourceAccountingFlag = NewBoolFlag("sandbox-placement-optimistic-resource-accounting", false)
 	FreePageReportingFlag            = NewBoolFlag("free-page-reporting", false)
 	FreezeUserCgroupFlag             = NewBoolFlag("freeze-user-cgroup", env.IsDevelopment())
+	// CollapseEnvdHeapFlag makes the orchestrator ask envd to collapse its own
+	// anonymous heap into 2 MiB hugepages just before pause, reducing the number
+	// of distinct frames envd faults on resume. Off by default; rolled out via LD.
+	CollapseEnvdHeapFlag = NewBoolFlag("collapse-envd-heap", false)
 
 	NetworkTransformRulesFlag = NewBoolFlag("network-transform-rules", env.IsDevelopment())
 

packages/shared/pkg/featureflags/flags.go

@@ -26,6 +26,11 @@ const (
 	TeamSandboxCreated CounterType = "e2b.team.sandbox.created"
 
 	EnvdInitCalls CounterType = "orchestrator.sandbox.envd.init.calls"
+
+	// 2 MiB chunks the pre-pause envd heap collapse attempted, split by the
+	// result attribute (collapsed|skipped): attempts = total, successful =
+	// collapsed.
+	EnvdCollapseChunks CounterType = "orchestrator.sandbox.envd.collapse.chunks"
 	// Incremented by the balance_dirty_pages thread count at every 200 ms poll
 	// for the lifetime of the process. rate() shows dirty-page throttle
 	// intensity in real-time; 0 when no stalls are occurring.
@@ -75,6 +80,11 @@ const (
 	OrchestratorSandboxCreateDurationName HistogramType = "orchestrator.sandbox.create.duration"
 	WaitForEnvdDurationHistogramName      HistogramType = "orchestrator.sandbox.envd.init.duration"
 
+	// Pre-pause envd heap collapse round-trip duration (the pause-path cost of
+	// POST /collapse: network plus envd's madvise work), recorded once per pause
+	// when the collapse-envd-heap flag is on.
+	EnvdCollapseDurationHistogramName HistogramType = "orchestrator.sandbox.envd.collapse.duration"
+
 	// Sandbox startup working-set histograms: demand-fault pages/bytes a guest
 	// needed to reach a successful envd init, recorded once per start. Sampled
 	// per start (not per fault), so histogram_quantile yields per-sandbox
@@ -186,6 +196,7 @@ var counterDesc = map[CounterType]string{
 	TeamSandboxCreated:                       "Counter of started sandboxes for the team in the interval",
 	OrchestratorHostBalanceDirtyPagesThreads: "Cumulative stalled thread-polls during sandbox resume; rate() gives throttle intensity",
 	EnvdInitCalls:                            "Number of envd initialization calls",
+	EnvdCollapseChunks:                       "2 MiB chunks the pre-pause envd heap collapse attempted, by result",
 	OrchestratorSandboxKilledCounterName:     "Number of sandboxes killed, labeled by kill reason",
 	TCPFirewallConnectionsTotal:              "Total number of TCP firewall connections processed",
 	TCPFirewallErrorsTotal:                   "Total number of TCP firewall errors",
@@ -213,6 +224,7 @@ var counterUnits = map[CounterType]string{
 	TeamSandboxCreated:                       "{sandbox}",
 	OrchestratorHostBalanceDirtyPagesThreads: "{thread}",
 	EnvdInitCalls:                            "1",
+	EnvdCollapseChunks:                       "{chunk}",
 	OrchestratorSandboxKilledCounterName:     "{sandbox}",
 	TCPFirewallConnectionsTotal:              "{connection}",
 	TCPFirewallErrorsTotal:                   "{error}",
@@ -389,6 +401,7 @@ var histogramDesc = map[HistogramType]string{
 	BuildRootfsSizeHistogramName:          "Size of the built template rootfs in bytes",
 	OrchestratorSandboxCreateDurationName: "Time taken to create a sandbox",
 	WaitForEnvdDurationHistogramName:      "Time taken for Envd to initialize successfully",
+	EnvdCollapseDurationHistogramName:     "Time taken for the pre-pause envd heap collapse round-trip",
 
 	UffdStartupPagesHistogramName:       "Demand-fault pages a guest needed to reach a successful envd init, per start",
 	UffdStartupSourcePagesHistogramName: "Subset of startup demand-fault pages pulled from the source (e.g. GCS), per start",
@@ -434,6 +447,7 @@ var histogramUnits = map[HistogramType]string{
 	BuildRootfsSizeHistogramName:                  "{By}",
 	OrchestratorSandboxCreateDurationName:         "ms",
 	WaitForEnvdDurationHistogramName:              "ms",
+	EnvdCollapseDurationHistogramName:             "ms",
 	UffdStartupPagesHistogramName:                 "{page}",
 	UffdStartupSourcePagesHistogramName:           "{page}",
 	UffdStartupBytesHistogramName:                 "{By}",

packages/shared/pkg/telemetry/meters.go

@@ -14,6 +14,11 @@ const MinEnvdVersionForSnapshot = "0.5.0"
 // sandbox doesn't end up permanently frozen after resume on older envds.
 const MinEnvdVersionForCgroupFreeze = "0.6.0"
 
+// MinEnvdVersionForHeapCollapse is the first envd that exposes a native
+// POST /collapse endpoint, which compacts envd's own anonymous heap into 2 MiB
+// hugepages before pause to reduce the frames it faults on resume.
+const MinEnvdVersionForHeapCollapse = "0.6.4"
+
 func sanitizeVersion(version string) string {
 	if len(version) > 0 && version[0] != 'v' {
 		version = "v" + version