From 3d57b09e05c3ff5052738d7531e2dccfa278b3c8 Mon Sep 17 00:00:00 2001 From: lorenzozanee Date: Thu, 8 Oct 2026 20:49:13 +0800 Subject: [PATCH] fix(orchestrator): preserve creation-time env vars across live envd upgrades --- docs/ARCHITECTURE.md | 4 +- packages/orchestrator/pkg/sandbox/envd.go | 42 ++++++++ .../pkg/sandbox/envd_upgrade_test.go | 96 +++++++++++++++++++ .../server/envd_upgrade_cache_gate_test.go | 47 +++++++++ packages/orchestrator/pkg/server/sandboxes.go | 7 ++ 5 files changed, 195 insertions(+), 1 deletion(-) diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 84d4989810..00f5893fa6 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -636,7 +636,9 @@ sequenceDiagram `packages/shared/pkg/featureflags`), via envd's `POST /upgrade` (see the envd section). It is best-effort — a delivery failure before the `exec` leaves the old envd serving — except an unrecoverable post-`exec` failure (the new envd never re-initializes), which fails the resume - rather than return a permanently unusable sandbox. + rather than return a permanently unusable sandbox. Before a live upgrade, the orchestrator reads + `/envs` into the resumed sandbox config so the post-upgrade `/init` can restore defaults that + were not retained in resume metadata; if that read fails, it defers the upgrade. - **Envd offline-upgrade on cold-boot resume**: reaches envd too old for the live `/upgrade` handover (below `MinEnvdVersionForUpgrade`). When a *filesystem-only* snapshot cold-boots (`RebootSandbox`), the orchestrator rewrites `/usr/bin/envd` in the rootfs **before** the VM diff --git a/packages/orchestrator/pkg/sandbox/envd.go b/packages/orchestrator/pkg/sandbox/envd.go index 82ec2259cc..4514de4fc7 100644 --- a/packages/orchestrator/pkg/sandbox/envd.go +++ b/packages/orchestrator/pkg/sandbox/envd.go @@ -254,6 +254,48 @@ func (s *Sandbox) postEnvd(ctx context.Context, timeout time.Duration, path stri return nil } +// CaptureEnvVarsForUpgrade reads the running envd's defaults so the post-upgrade +// /init can restore them when resume metadata did not retain the original values. +func (s *Sandbox) CaptureEnvVarsForUpgrade(ctx context.Context) error { + request, err := http.NewRequestWithContext(ctx, http.MethodGet, s.envdServerURL()+"/envs", nil) + if err != nil { + return fmt.Errorf("build env vars request: %w", err) + } + if s.Config.Envd.AccessToken != nil { + request.Header.Set("X-Access-Token", *s.Config.Envd.AccessToken) + } + + response, err := sandboxHttpClient.Do(request) + if err != nil { + return fmt.Errorf("get env vars before envd upgrade: %w", err) + } + defer response.Body.Close() + + if response.StatusCode != http.StatusOK { + return fmt.Errorf("get env vars before envd upgrade returned %d", response.StatusCode) + } + + var envVars map[string]string + if err := json.NewDecoder(response.Body).Decode(&envVars); err != nil { + return fmt.Errorf("decode env vars before envd upgrade: %w", err) + } + if envVars == nil { + return errors.New("get env vars before envd upgrade returned null") + } + // /envs flattens user and internal values; the new envd rebuilds these keys itself. + for _, key := range [...]string{ + "E2B_SANDBOX", + "E2B_SANDBOX_ID", + "E2B_TEMPLATE_ID", + "E2B_EVENTS_ADDRESS", + } { + delete(envVars, key) + } + + s.Config.Envd.Vars = envVars + return nil +} + // CallEnvdUpgrade triggers envd's POST /upgrade — the orchestrator-driven // live-upgrade trigger. It streams the new envd binary // from localSrcPath as the (authenticated) request body; envd writes it to diff --git a/packages/orchestrator/pkg/sandbox/envd_upgrade_test.go b/packages/orchestrator/pkg/sandbox/envd_upgrade_test.go index c8c8886483..f0cc11ae61 100644 --- a/packages/orchestrator/pkg/sandbox/envd_upgrade_test.go +++ b/packages/orchestrator/pkg/sandbox/envd_upgrade_test.go @@ -4,15 +4,111 @@ package sandbox import ( "context" + "encoding/json" "errors" "io" "net" + "net/http" + "net/http/httptest" "syscall" "testing" + "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) +func TestCaptureEnvVarsForUpgradeReplaysValuesOnInit(t *testing.T) { + t.Parallel() + + const ( + accessToken = "test-access-token" + sentinel = "creation-time-value" + ) + + type initRequest struct { + envVars map[string]string `json:"envVars"` + } + initRequests := make(chan initRequest, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("X-Access-Token") != accessToken { + w.WriteHeader(http.StatusUnauthorized) + return + } + + switch r.URL.Path { + case "/envs": + if r.Method != http.MethodGet { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"RESUME_SENTINEL":"`+sentinel+`","E2B_CUSTOM":"user-value","E2B_SANDBOX":"old","E2B_SANDBOX_ID":"old","E2B_TEMPLATE_ID":"old","E2B_EVENTS_ADDRESS":"old"}`) + case "/init": + if r.Method != http.MethodPost { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + var body struct { + EnvVars map[string]string `json:"envVars"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + initRequests <- initRequest{envVars: body.EnvVars} + w.WriteHeader(http.StatusNoContent) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + defer server.Close() + + accessTokenValue := accessToken + sbx := &Sandbox{Metadata: &Metadata{ + internalConfig: internalConfig{EnvdInitRequestTimeout: 5 * time.Second, envdServerURLOverride: server.URL}, + Config: NewConfig(Config{Envd: EnvdMetadata{AccessToken: &accessTokenValue}}), + }} + + require.NoError(t, sbx.CaptureEnvVarsForUpgrade(t.Context())) + resp, _, err := sbx.doRequestWithInfiniteRetries(t.Context(), http.MethodPost, server.URL+"/init") + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusNoContent, resp.StatusCode) + + var request initRequest + select { + case request = <-initRequests: + case <-time.After(5 * time.Second): + t.Fatal("envd did not receive the /init request") + } + assert.Equal(t, sentinel, request.envVars["RESUME_SENTINEL"]) + assert.Equal(t, "user-value", request.envVars["E2B_CUSTOM"]) + assert.NotContains(t, request.envVars, "E2B_SANDBOX") + assert.NotContains(t, request.envVars, "E2B_SANDBOX_ID") + assert.NotContains(t, request.envVars, "E2B_TEMPLATE_ID") + assert.NotContains(t, request.envVars, "E2B_EVENTS_ADDRESS") +} + +func TestCaptureEnvVarsForUpgradeLeavesConfigUnchangedOnFailure(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + })) + defer server.Close() + + existing := map[string]string{"EXISTING": "value"} + sbx := &Sandbox{Metadata: &Metadata{ + internalConfig: internalConfig{envdServerURLOverride: server.URL}, + Config: NewConfig(Config{Envd: EnvdMetadata{Vars: existing}}), + }} + + err := sbx.CaptureEnvVarsForUpgrade(t.Context()) + require.Error(t, err) + assert.Equal(t, existing, sbx.Config.Envd.Vars) +} + // TestIsUpgradeDeliveryFailure guards the distinction CallEnvdUpgrade relies on: // a request that never reached a running envd (so no upgrade happened) is a // failure, while the expected post-send connection drop when envd execs diff --git a/packages/orchestrator/pkg/server/envd_upgrade_cache_gate_test.go b/packages/orchestrator/pkg/server/envd_upgrade_cache_gate_test.go index 472d751ec4..60c75436bc 100644 --- a/packages/orchestrator/pkg/server/envd_upgrade_cache_gate_test.go +++ b/packages/orchestrator/pkg/server/envd_upgrade_cache_gate_test.go @@ -5,6 +5,8 @@ package server import ( "context" "fmt" + "net" + "net/http" "os" "os/exec" "path/filepath" @@ -22,6 +24,8 @@ import ( "github.com/e2b-dev/infra/packages/orchestrator/pkg/cfg" "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox" "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/envdbin" + "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/network" + "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/sandboxtypes" "github.com/e2b-dev/infra/packages/shared/pkg/telemetry" @@ -247,6 +251,49 @@ func TestUpgradeDefersWhenTheBinaryIsNotCached(t *testing.T) { "a deferral delivered nothing, so it must not appear as an attempt") } +func TestUpgradeDefersWhenEnvVarsCannotBeCaptured(t *testing.T) { + dir := warmScratchDir(t) + src := filepath.Join(dir, "envd") + require.NoError(t, os.WriteFile(src, []byte("#!/bin/sh\necho 0.6.14\n"), 0o755)) + + listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", consts.DefaultEnvdServerPort)) + if err != nil { + t.Skipf("envd port is unavailable: %v", err) + } + upgradeCalls := make(chan struct{}, 1) + envdServer := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/envs": + w.WriteHeader(http.StatusServiceUnavailable) + case "/upgrade": + upgradeCalls <- struct{}{} + w.WriteHeader(http.StatusInternalServerError) + default: + http.NotFound(w, r) + } + })} + go func() { _ = envdServer.Serve(listener) }() + t.Cleanup(func() { _ = envdServer.Close() }) + + orchestrator, reader := upgradeGateServer(t, src, filepath.Join(dir, "base"), false) + sbx := upgradeSandbox("0.6.13") + user := "sandbox" + sbx.Config.Envd.DefaultUser = &user + sbx.Resources = &sandbox.Resources{Slot: &network.Slot{HostIP: net.ParseIP("127.0.0.1")}} + + upgraded, err := orchestrator.maybeUpgradeEnvd(t.Context(), sbx) + require.NoError(t, err) + require.False(t, upgraded) + select { + case <-upgradeCalls: + t.Fatal("envd upgrade was triggered without captured environment defaults") + default: + } + require.Equal(t, map[string]int64{"env_vars_unavailable": 1}, + counterPoints(t, reader, telemetry.OrchestratorEnvdUpgradeGated, "reason")) + require.Empty(t, counterPoints(t, reader, telemetry.OrchestratorEnvdUpgradeAttempts, "result")) +} + // TestUpgradeWithTheCacheFlagOffProbesTheSourceItself is the direction that shows // the gate is not inverted. With the flag off the live path must behave exactly as // it did before the cache existed: it execs the source to learn its version, and diff --git a/packages/orchestrator/pkg/server/sandboxes.go b/packages/orchestrator/pkg/server/sandboxes.go index 70095ffbff..1bc6addb97 100644 --- a/packages/orchestrator/pkg/server/sandboxes.go +++ b/packages/orchestrator/pkg/server/sandboxes.go @@ -2447,6 +2447,13 @@ func (s *Server) maybeUpgradeEnvd(ctx context.Context, sbx *sandbox.Sandbox) (up return false, nil } + if err := sbx.CaptureEnvVarsForUpgrade(ctx); err != nil { + s.envdUpgradeGated.Add(ctx, 1, metric.WithAttributes(attribute.String("reason", "env_vars_unavailable"))) + sbxlogger.I(sbx).Warn(ctx, "envd auto-upgrade: could not capture default env vars; deferring", + zap.Error(err)) + + return false, nil + } attempted = true start = time.Now()