feat(sandbox): pre-pause guest reclaim via envd

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

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"log"
 	"math"
-	"net/http"
 	"os"
 	"os/signal"
 	"path/filepath"
@@ -16,11 +15,11 @@ import (
 	"syscall"
 	"time"
 
-	"connectrpc.com/connect"
 	"github.com/containernetworking/plugins/pkg/ns"
 	"github.com/coreos/go-iptables/iptables"
 	"github.com/google/uuid"
 	"github.com/launchdarkly/go-sdk-common/v3/ldlog"
+	"github.com/launchdarkly/go-sdk-common/v3/ldvalue"
 	"github.com/vishvananda/netlink"
 	"golang.org/x/sys/unix"
 
@@ -39,11 +38,8 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/tcpfirewall"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/core/rootfs"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/metadata"
-	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
-	"github.com/e2b-dev/infra/packages/shared/pkg/grpc"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/envd/process"
-	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/envd/process/processconnect"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
@@ -73,8 +69,20 @@ 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.
+	reclaim := flag.Bool("reclaim", false, "enable pre-pause reclaim chain (fstrim 500ms, sync 500ms, drop_caches 200ms, compact 1s)")
+
 	flag.Parse()
 
+	if *reclaim {
+		featureflags.NewJSONFlag("guest-pause-reclaim", ldvalue.FromJSONMarshal(map[string]int{
+			"sync":           500,
+			"drop_caches":    200,
+			"compact_memory": 1000,
+			"fstrim":         500,
+		}))
+	}
+
 	if *fromBuild == "" {
 		log.Fatal("-from-build required")
 	}
@@ -1184,32 +1192,15 @@ func printTemplateInfo(ctx context.Context, tmpl template.Template, meta metadat
 	}
 }
 
-// runCommandInSandbox runs a command inside the sandbox via envd
-func runCommandInSandbox(ctx context.Context, sbx *sandbox.Sandbox, command string) error {
-	// Connect directly to envd on the sandbox
-	envdURL := fmt.Sprintf("http://%s:%d", sbx.Slot.HostIPString(), consts.DefaultEnvdServerPort)
-
-	hc := http.Client{
-		Timeout:   10 * time.Minute,
-		Transport: sandbox.SandboxHttpTransport,
-	}
-
-	processC := processconnect.NewProcessClient(&hc, envdURL)
+// runCommandInSandboxTimeout caps how long a single resume-build command may
+// run before envd kills it. Restores the prior 10-minute upper bound that the
+// shared http.Client used to enforce, so a stuck command can't block the CLI.
+const runCommandInSandboxTimeout = 10 * time.Minute
 
-	req := connect.NewRequest(&process.StartRequest{
-		Process: &process.ProcessConfig{
-			Cmd:  "/bin/bash",
-			Args: []string{"-l", "-c", command},
-		},
-	})
-	grpc.SetUserHeader(req.Header(), "root")
-
-	// Set access token if available
-	if sbx.Config.Envd.AccessToken != nil {
-		req.Header().Set("X-Access-Token", *sbx.Config.Envd.AccessToken)
-	}
-
-	stream, err := processC.Start(ctx, req)
+// runCommandInSandbox runs a command inside the sandbox via envd as a
+// login shell so /etc/profile is sourced.
+func runCommandInSandbox(ctx context.Context, sbx *sandbox.Sandbox, command string) error {
+	stream, err := sbx.StartEnvdShell(ctx, "/bin/bash", []string{"-l", "-c", command}, "root", runCommandInSandboxTimeout)
 	if err != nil {
 		return fmt.Errorf("failed to start process: %w", err)
 	}

packages/orchestrator/pkg/sandbox/envd_process.go

@@ -0,0 +1,43 @@
+package sandbox
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"strconv"
+	"time"
+
+	"connectrpc.com/connect"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
+	"github.com/e2b-dev/infra/packages/shared/pkg/grpc"
+	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/envd/process"
+	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/envd/process/processconnect"
+)
+
+// StartEnvdShell opens a streaming Process.Start call against the sandbox's
+// envd. timeout > 0 sets Connect-Timeout-Ms so envd kills the process at
+// the deadline. Caller owns the returned stream.
+func (s *Sandbox) StartEnvdShell(
+	ctx context.Context,
+	shell string,
+	shellArgs []string,
+	user string,
+	timeout time.Duration,
+) (*connect.ServerStreamForClient[process.StartResponse], error) {
+	addr := fmt.Sprintf("http://%s:%d", s.Slot.HostIPString(), consts.DefaultEnvdServerPort)
+	pc := processconnect.NewProcessClient(&http.Client{Transport: sandboxHttpClient.Transport}, addr)
+
+	req := connect.NewRequest(&process.StartRequest{
+		Process: &process.ProcessConfig{Cmd: shell, Args: shellArgs},
+	})
+	if timeout > 0 {
+		req.Header().Set("Connect-Timeout-Ms", strconv.FormatInt(timeout.Milliseconds(), 10))
+	}
+	if s.Config.Envd.AccessToken != nil {
+		req.Header().Set("X-Access-Token", *s.Config.Envd.AccessToken)
+	}
+	grpc.SetUserHeader(req.Header(), user)
+
+	return pc.Start(ctx, req)
+}

packages/orchestrator/pkg/sandbox/reclaim.go

@@ -0,0 +1,95 @@
+package sandbox
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+)
+
+// Slack covers shell start + envd round-trip overhead.
+const reclaimOuterSlack = 500 * time.Millisecond
+
+// Order: fstrim → sync → drop_caches → compact_memory. fstrim runs first so
+// that the fs metadata it pulls into the page cache and the superblock dirties
+// (last-trim timestamps) get flushed by sync and evicted by drop_caches in the
+// same pass; compact_memory then consolidates the minimal RSS so the snapshot
+// has long contiguous zero runs that compress well. Each step is disabled at
+// sub-ms cap. Returns ("", 0) when every step is disabled.
+func (s *Sandbox) buildReclaimScript(ctx context.Context) (string, time.Duration) {
+	cfg := featureflags.GetReclaimConfig(ctx, s.featureFlags,
+		featureflags.SandboxContext(s.Runtime.SandboxID),
+		featureflags.TeamContext(s.Runtime.TeamID),
+		featureflags.TemplateContext(s.Runtime.TemplateID),
+	)
+
+	steps := []struct {
+		cap time.Duration
+		cmd string
+	}{
+		{cfg.Fstrim, "fstrim -av"},
+		{cfg.Sync, "sync"},
+		{cfg.DropCaches, "echo 3 > /proc/sys/vm/drop_caches"},
+		{cfg.CompactMemory, "echo 1 > /proc/sys/vm/compact_memory"},
+	}
+
+	var (
+		parts []string
+		sum   time.Duration
+	)
+	for _, st := range steps {
+		// %.3f at <1ms renders as 0.000 → GNU timeout reads as "no timeout".
+		if st.cap < time.Millisecond {
+			continue
+		}
+		parts = append(parts, fmt.Sprintf("timeout -s KILL %.3f sh -c %q >/dev/null 2>&1 || rc=$?", st.cap.Seconds(), st.cmd))
+		sum += st.cap
+	}
+	if len(parts) == 0 {
+		return "", 0
+	}
+
+	return "rc=0; " + strings.Join(parts, "; ") + "; exit $rc", sum + reclaimOuterSlack
+}
+
+// bestEffortReclaim runs the reclaim chain via envd before pause.
+func (s *Sandbox) bestEffortReclaim(ctx context.Context) {
+	script, timeout := s.buildReclaimScript(ctx)
+	if script == "" {
+		return
+	}
+
+	ctx, span := tracer.Start(ctx, "envd-reclaim")
+	defer span.End()
+
+	rcCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	stream, err := s.StartEnvdShell(rcCtx, "/bin/sh", []string{"-c", script}, "root", timeout)
+	if err != nil {
+		logger.L().Warn(ctx, "envd reclaim failed", logger.WithSandboxID(s.Runtime.SandboxID), zap.Error(err))
+
+		return
+	}
+	defer stream.Close()
+
+	var exitCode int32
+	for stream.Receive() {
+		if end := stream.Msg().GetEvent().GetEnd(); end != nil {
+			exitCode = end.GetExitCode()
+		}
+	}
+	if err := stream.Err(); err != nil {
+		logger.L().Warn(ctx, "envd reclaim stream error", logger.WithSandboxID(s.Runtime.SandboxID), zap.Error(err))
+
+		return
+	}
+	if exitCode != 0 {
+		logger.L().Warn(ctx, "envd reclaim non-zero exit", logger.WithSandboxID(s.Runtime.SandboxID), zap.Int32("exit_code", exitCode))
+	}
+}

packages/orchestrator/pkg/sandbox/sandbox.go

@@ -217,6 +217,8 @@ type Sandbox struct {
 	files   *storage.SandboxFiles
 	cleanup *Cleanup
 
+	featureFlags *featureflags.Client
+
 	process      *fc.Process
 	cgroupHandle *cgroup.CgroupHandle
 
@@ -457,7 +459,8 @@ func (f *Factory) CreateSandbox(
 		files:    sandboxFiles,
 		process:  fcHandle,
 
-		cleanup: cleanup,
+		cleanup:      cleanup,
+		featureFlags: f.featureFlags,
 
 		APIStoredConfig: apiConfigToStore,
 
@@ -797,7 +800,8 @@ func (f *Factory) ResumeSandbox(
 		files:    sandboxFiles,
 		process:  fcHandle,
 
-		cleanup: cleanup,
+		cleanup:      cleanup,
+		featureFlags: f.featureFlags,
 
 		APIStoredConfig: apiConfigToStore,
 		CABundle:        f.egressProxy.CABundle(),
@@ -1051,6 +1055,11 @@ func (s *Sandbox) Pause(
 	// Stop the health check before pausing the VM
 	s.Checks.Stop()
 
+	// Best-effort pre-pause guest reclaim (fstrim, sync, drop_caches,
+	// compact_memory) on the live VM via envd. Per-step caps are LD-flag-driven;
+	// all default to 0 which disables the chain entirely. Non-fatal.
+	s.bestEffortReclaim(ctx)
+
 	if err := s.process.Pause(ctx); err != nil {
 		return nil, fmt.Errorf("failed to pause VM: %w", err)
 	}

packages/shared/pkg/featureflags/flags.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/launchdarkly/go-sdk-common/v3/ldcontext"
 	"github.com/launchdarkly/go-sdk-common/v3/ldvalue"
@@ -213,6 +214,32 @@ var (
 	MinChunkerReadSizeKB = NewIntFlag("min-chunker-read-size-kb", 16)
 )
 
+// ReclaimConfigFlag holds per-step caps in milliseconds for the pre-pause
+// reclaim chain. Missing/zero/negative values disable the step.
+// Example: {"sync":500,"drop_caches":200,"compact_memory":1000,"fstrim":500}
+var ReclaimConfigFlag = NewJSONFlag("guest-pause-reclaim", ldvalue.Null())
+
+type ReclaimConfig struct {
+	Sync          time.Duration
+	DropCaches    time.Duration
+	CompactMemory time.Duration
+	Fstrim        time.Duration
+}
+
+func GetReclaimConfig(ctx context.Context, ff *Client, contexts ...ldcontext.Context) ReclaimConfig {
+	v := ff.JSONFlag(ctx, ReclaimConfigFlag, contexts...)
+	ms := func(key string) time.Duration {
+		return time.Duration(v.GetByKey(key).IntValue()) * time.Millisecond
+	}
+
+	return ReclaimConfig{
+		Sync:          ms("sync"),
+		DropCaches:    ms("drop_caches"),
+		CompactMemory: ms("compact_memory"),
+		Fstrim:        ms("fstrim"),
+	}
+}
+
 type StringFlag struct {
 	name     string
 	fallback string