refactor(envd): share one workload freezer across freeze paths

The live-upgrade handover added a second workload freeze/unfreeze path
(process.Service.FreezeWorkload/UnfreezeWorkload) that swept the user/pty
cgroups directly, independent of the HTTP API's existing /freeze, /unfreeze
and /init-deferred-thaw. Those API paths serialize their sweep through
API.freezeLock precisely so a pause freeze, a rollback unfreeze and the
resume thaw can't interleave and strand the workload frozen; the upgrade path
took no lock at all, so an upgrade freeze racing the resume thaw could leave
the sandbox frozen.

Extract the sweep and its lock into a single cgroups.WorkloadFreezer that
owns the manager plus a serializing semaphore and exposes Freeze/Unfreeze
(best-effort, joined errors; Unfreeze detaches from ctx cancellation so a
thaw always lands). main constructs one instance and passes it to both the
HTTP API and the process service, so every freeze/unfreeze caller now shares
the same lock. The canonical user/pty cgroup list moves to
cgroups.WorkloadProcessTypes.

The handover now holds that shared lock across the WHOLE window (freeze ->
serialize -> execve) via a new FreezeHold/FreezeWorkloadHold, not just the
freeze sweep, so a concurrent /init or /unfreeze thaw blocks on the lock until
the handover finishes and cannot thaw the workload mid-handover.

No behavior change to the existing endpoints: PostFreeze still returns 503 on
a cancelled request and 500 on a sweep failure, PostUnfreeze/​/init still
attempt every cgroup best-effort. Per-cgroup error logs collapse into one
joined log line per call.

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

Author: kalyazin

packages/envd/internal/api/compose_test.go

@@ -33,7 +33,7 @@ func newComposeTestAPI(t *testing.T) (*API, *user.User) {
 		User:    currentUser.Username,
 	}
 
-	return New(&logger, defaults, nil, false, cgroups.NewNoopManager()), currentUser
+	return New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager())), currentUser
 }
 
 func writeSourceFile(t *testing.T, dir string, name string, data []byte) string {

packages/envd/internal/api/download_test.go

@@ -96,7 +96,7 @@ func TestGetFilesContentDisposition(t *testing.T) {
 				EnvVars: utils.NewEnvVars(),
 				User:    currentUser.Username,
 			}
-			api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+			api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 			// Create request and response recorder
 			req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/files?path="+url.QueryEscape(tempFile), nil)
@@ -145,7 +145,7 @@ func TestGetFilesContentDispositionWithNestedPath(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	// Create request and response recorder
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/files?path="+url.QueryEscape(tempFile), nil)
@@ -188,7 +188,7 @@ func TestGetFiles_GzipEncoding_ExplicitIdentityOffWithRange(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	// Create request and response recorder
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/files?path="+url.QueryEscape(tempFile), nil)
@@ -229,7 +229,7 @@ func TestGetFiles_GzipDownload(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/files?path="+url.QueryEscape(tempFile), nil)
 	req.Header.Set("Accept-Encoding", "gzip")
@@ -294,7 +294,7 @@ func TestPostFiles_GzipUpload(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), &gzBuf)
 	req.Header.Set("Content-Type", mpWriter.FormDataContentType())
@@ -334,7 +334,7 @@ func TestPostFiles_RawBodyUpload(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), bytes.NewReader(originalContent))
 	req.Header.Set("Content-Type", "application/octet-stream")
@@ -372,7 +372,7 @@ func TestPostFiles_RawBodyUploadCreatesDirectories(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), bytes.NewReader(originalContent))
 	req.Header.Set("Content-Type", "application/octet-stream")
@@ -405,7 +405,7 @@ func TestPostFiles_RawBodyUploadRequiresPath(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files", bytes.NewReader([]byte("some content")))
 	req.Header.Set("Content-Type", "application/octet-stream")
@@ -440,7 +440,7 @@ func TestPostFiles_RawBodyUploadOverwritesExisting(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), bytes.NewReader(newContent))
 	req.Header.Set("Content-Type", "application/octet-stream")
@@ -486,7 +486,7 @@ func TestPostFiles_RawBodyGzipUpload(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), &gzBuf)
 	req.Header.Set("Content-Type", "application/octet-stream")
@@ -520,7 +520,7 @@ func TestPostFiles_UnsupportedContentType(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	tempDir := t.TempDir()
 	destPath := filepath.Join(tempDir, "test.txt")
@@ -566,7 +566,7 @@ func TestPostFiles_MultipartStillWorksWithoutContentType(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), &multipartBuf)
 	req.Header.Set("Content-Type", mpWriter.FormDataContentType())
@@ -624,7 +624,7 @@ func TestGzipUploadThenGzipDownload(t *testing.T) {
 		EnvVars: utils.NewEnvVars(),
 		User:    currentUser.Username,
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 
 	uploadReq := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/files?path="+url.QueryEscape(destPath), &gzBuf)
 	uploadReq.Header.Set("Content-Type", mpWriter.FormDataContentType())

packages/envd/internal/api/init.go

@@ -22,7 +22,6 @@ import (
 	"github.com/e2b-dev/infra/packages/envd/internal/host"
 	"github.com/e2b-dev/infra/packages/envd/internal/logs"
 	"github.com/e2b-dev/infra/packages/envd/internal/logs/ratelimit"
-	"github.com/e2b-dev/infra/packages/envd/internal/services/cgroups"
 	"github.com/e2b-dev/infra/packages/envd/pkg"
 	"github.com/e2b-dev/infra/packages/shared/pkg/keys"
 )
@@ -268,11 +267,6 @@ func (a *API) SetData(ctx context.Context, logger zerolog.Logger, data PostInitJ
 }
 
 // userCgroupsToFreeze is the cgroup set frozen pre-pause and thawed on /init.
-var userCgroupsToFreeze = []cgroups.ProcessType{
-	cgroups.ProcessTypeUser,
-	cgroups.ProcessTypePTY,
-}
-
 // PostFreeze freezes user/pty cgroups directly (no Process.Start / shell).
 // Orchestrator calls this just before pause; the frozen state persists into the
 // snapshot and /init thaws on resume. Best-effort: tries every cgroup even if
@@ -282,22 +276,16 @@ func (a *API) PostFreeze(w http.ResponseWriter, r *http.Request) {
 
 	logger := a.logger.With().Str(string(logs.OperationIDKey), logs.AssignOperationID()).Logger()
 
-	if err := a.freezeLock.Acquire(r.Context(), 1); err != nil {
-		w.WriteHeader(http.StatusServiceUnavailable)
-
-		return
-	}
-	defer a.freezeLock.Release(1)
+	if err := a.workloadFreezer.Freeze(r.Context()); err != nil {
+		// A failed lock acquire means the request ctx was cancelled; a failed
+		// sweep leaves ctx intact.
+		if r.Context().Err() != nil {
+			w.WriteHeader(http.StatusServiceUnavailable)
 
-	var errs []error
-	for _, pt := range userCgroupsToFreeze {
-		if err := a.cgroupManager.Freeze(pt); err != nil {
-			logger.Error().Err(err).Msgf("freeze %s cgroup", pt)
-			errs = append(errs, fmt.Errorf("freeze %s cgroup: %w", pt, err))
+			return
 		}
-	}
-	if len(errs) > 0 {
-		jsonError(w, http.StatusInternalServerError, errors.Join(errs...))
+		logger.Error().Err(err).Msg("freeze workload cgroups")
+		jsonError(w, http.StatusInternalServerError, err)
 
 		return
 	}
@@ -313,25 +301,11 @@ func (a *API) PostFreeze(w http.ResponseWriter, r *http.Request) {
 func (a *API) PostUnfreeze(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
 
-	ctx := r.Context()
 	logger := a.logger.With().Str(string(logs.OperationIDKey), logs.AssignOperationID()).Logger()
 
-	if err := a.freezeLock.Acquire(context.WithoutCancel(ctx), 1); err != nil {
-		w.WriteHeader(http.StatusServiceUnavailable)
-
-		return
-	}
-	defer a.freezeLock.Release(1)
-
-	var errs []error
-	for _, pt := range userCgroupsToFreeze {
-		if err := a.cgroupManager.Unfreeze(pt); err != nil {
-			logger.Error().Err(err).Msgf("unfreeze %s cgroup", pt)
-			errs = append(errs, fmt.Errorf("unfreeze %s cgroup: %w", pt, err))
-		}
-	}
-	if len(errs) > 0 {
-		jsonError(w, http.StatusInternalServerError, errors.Join(errs...))
+	if err := a.workloadFreezer.Unfreeze(r.Context()); err != nil {
+		logger.Error().Err(err).Msg("unfreeze workload cgroups")
+		jsonError(w, http.StatusInternalServerError, err)
 
 		return
 	}
@@ -341,15 +315,11 @@ func (a *API) PostUnfreeze(w http.ResponseWriter, r *http.Request) {
 }
 
 // unfreezeUserCgroups unfreezes user/pty cgroups (idempotent if not frozen).
-// Wraps the context with WithoutCancel so the unfreeze always completes.
+// The freezer detaches the wait from ctx cancellation so the unfreeze always
+// completes.
 func (a *API) unfreezeUserCgroups(ctx context.Context, logger zerolog.Logger) {
-	_ = a.freezeLock.Acquire(context.WithoutCancel(ctx), 1)
-	defer a.freezeLock.Release(1)
-
-	for _, pt := range userCgroupsToFreeze {
-		if err := a.cgroupManager.Unfreeze(pt); err != nil {
-			logger.Warn().Err(err).Msgf("unfreeze %s cgroup", pt)
-		}
+	if err := a.workloadFreezer.Unfreeze(ctx); err != nil {
+		logger.Warn().Err(err).Msg("unfreeze workload cgroups")
 	}
 }
 

packages/envd/internal/api/init_test.go

@@ -148,7 +148,7 @@ func newTestAPI(accessToken *SecureToken, mmdsClient MMDSClient) *API {
 	defaults := &execcontext.Defaults{
 		EnvVars: utils.NewEnvVars(),
 	}
-	api := New(&logger, defaults, nil, false, cgroups.NewNoopManager())
+	api := New(&logger, defaults, nil, false, cgroups.NewWorkloadFreezer(cgroups.NewNoopManager()))
 	if accessToken != nil {
 		api.accessToken.TakeFrom(accessToken)
 	}
@@ -637,7 +637,7 @@ func (f *fakeCgroupManager) Close() error { return nil }
 func newAPIWithCgroupManager(mgr cgroups.Manager) *API {
 	logger := zerolog.Nop()
 
-	return New(&logger, &execcontext.Defaults{EnvVars: utils.NewEnvVars()}, nil, false, mgr)
+	return New(&logger, &execcontext.Defaults{EnvVars: utils.NewEnvVars()}, nil, false, cgroups.NewWorkloadFreezer(mgr))
 }
 
 func TestPostFreeze(t *testing.T) {
@@ -654,7 +654,7 @@ func TestPostFreeze(t *testing.T) {
 		api.PostFreeze(rec, req)
 
 		require.Equal(t, http.StatusNoContent, rec.Code)
-		assert.Equal(t, userCgroupsToFreeze, mgr.frozen)
+		assert.Equal(t, cgroups.WorkloadProcessTypes, mgr.frozen)
 	})
 
 	t.Run("returns 500 on freeze error", func(t *testing.T) {
@@ -686,7 +686,7 @@ func TestPostUnfreeze(t *testing.T) {
 		api.PostUnfreeze(rec, req)
 
 		require.Equal(t, http.StatusNoContent, rec.Code)
-		assert.Equal(t, userCgroupsToFreeze, mgr.unfrozen)
+		assert.Equal(t, cgroups.WorkloadProcessTypes, mgr.unfrozen)
 	})
 
 	t.Run("returns 500 but attempts every cgroup on unfreeze error", func(t *testing.T) {
@@ -701,7 +701,7 @@ func TestPostUnfreeze(t *testing.T) {
 
 		assert.Equal(t, http.StatusInternalServerError, rec.Code)
 		assert.Empty(t, mgr.unfrozen)
-		assert.Equal(t, userCgroupsToFreeze, mgr.unfreezeAttempts)
+		assert.Equal(t, cgroups.WorkloadProcessTypes, mgr.unfreezeAttempts)
 	})
 }
 
@@ -732,7 +732,7 @@ func TestPostInit_UnfreezeOnStaleTimestamp(t *testing.T) {
 	require.Equal(t, http.StatusNoContent, rec.Code)
 	_, ok := api.defaults.EnvVars.Load("SHOULD_NOT_BE_SET")
 	assert.False(t, ok, "stale /init should not apply EnvVars")
-	assert.Equal(t, userCgroupsToFreeze, mgr.unfrozen, "stale /init must still unfreeze")
+	assert.Equal(t, cgroups.WorkloadProcessTypes, mgr.unfrozen, "stale /init must still unfreeze")
 }
 
 // Unauthorized /init must NOT thaw cgroups.

packages/envd/internal/api/store.go

@@ -43,14 +43,13 @@ type API struct {
 	initLock    *semaphore.Weighted
 
 	caCertInstaller *host.CACertInstaller
-	cgroupManager   cgroups.Manager
-	// freezeLock serializes the per-cgroup sweep across /freeze, /unfreeze
-	// and the /init deferred unfreeze. PostFreeze acquires with the request
-	// ctx; unfreeze paths acquire with Background so they always land
-	// regardless of HTTP-client cancellation.
-	freezeLock    *semaphore.Weighted
-	isMountingNFS atomic.Bool
-	mountedPaths  sync.Map // map[path]lifecycleID - tracks which lifecycle each path was mounted for
+	// workloadFreezer freezes/thaws the user+pty cgroups. Shared with the process
+	// service (the live-upgrade handover) so every freeze/unfreeze caller — this
+	// API's /freeze, /unfreeze and /init deferred thaw, plus the upgrade — is
+	// serialized through one lock.
+	workloadFreezer *cgroups.WorkloadFreezer
+	isMountingNFS   atomic.Bool
+	mountedPaths    sync.Map // map[path]lifecycleID - tracks which lifecycle each path was mounted for
 
 	// fsFreezer freezes/thaws the guest rootfs for filesystem-only pauses;
 	// fsFreezeLock serializes /fsfreeze and /fsthaw.
@@ -101,7 +100,7 @@ func (a *API) SetHandoverResult(procs, procsFailed, retained, retainedFailed, wa
 	}
 }
 
-func New(l *zerolog.Logger, defaults *execcontext.Defaults, mmdsChan chan *host.MMDSOpts, isNotFC bool, cgroupManager cgroups.Manager) *API {
+func New(l *zerolog.Logger, defaults *execcontext.Defaults, mmdsChan chan *host.MMDSOpts, isNotFC bool, workloadFreezer *cgroups.WorkloadFreezer) *API {
 	return &API{
 		logger:          l,
 		defaults:        defaults,
@@ -111,9 +110,8 @@ func New(l *zerolog.Logger, defaults *execcontext.Defaults, mmdsChan chan *host.
 		lastSetTime:     utils.NewAtomicMax(),
 		accessToken:     &SecureToken{},
 		caCertInstaller: host.NewCACertInstaller(l),
-		cgroupManager:   cgroupManager,
+		workloadFreezer: workloadFreezer,
 		initLock:        semaphore.NewWeighted(1),
-		freezeLock:      semaphore.NewWeighted(1),
 		fsFreezer:       fsfreeze.New(),
 		fsFreezeLock:    semaphore.NewWeighted(1),
 	}