feat(orchestrator): graceful sandbox drain on shutdown

packages/orchestrator/pkg/factories/run.go

@@ -885,6 +885,16 @@ func run(config cfg.Config, opts Options) (success bool) {
 		}
 	}
 
+	// Gracefully wait for live sandboxes to exit before closing the services they
+	// depend on. The forced-stop path skips this and tears sandboxes down later.
+	if !config.ForceStop {
+		logger.L().Info(ctx, "Starting sandbox drain phase", zap.Int("sandbox_count", sandboxes.Count()))
+		if err := orchestratorService.DrainSandboxes(closeCtx); err != nil {
+			logger.L().Error(ctx, "error while draining sandboxes", zap.Error(err))
+			success = false
+		}
+	}
+
 	slices.Reverse(closers)
 	for _, closer := range closers {
 		clog := globalLogger.With(zap.String("service", closer.name), zap.Bool("forced", config.ForceStop))

packages/orchestrator/pkg/server/drain_test.go

@@ -0,0 +1,101 @@
+//go:build linux
+
+package server
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox"
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/network"
+)
+
+func TestDrainSandboxesReturnsWhenEmpty(t *testing.T) {
+	t.Parallel()
+
+	s := drainTestServer()
+
+	require.NoError(t, s.DrainSandboxes(t.Context()))
+}
+
+func TestDrainSandboxesBlocksWhileLiveAndReturnsOnCancel(t *testing.T) {
+	t.Parallel()
+
+	s := drainTestServer()
+	sbx := drainTestSandbox(t, "lifecycle-1")
+	s.sandboxFactory.Sandboxes.MarkRunning(t.Context(), sbx)
+
+	ctx, cancel := context.WithCancel(t.Context())
+	done := make(chan error, 1)
+	go func() {
+		done <- s.DrainSandboxes(ctx)
+	}()
+
+	// A live sandbox keeps the drain blocked.
+	select {
+	case err := <-done:
+		require.Failf(t, "DrainSandboxes returned while a sandbox was live", "err: %v", err)
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	cancel()
+
+	select {
+	case err := <-done:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("DrainSandboxes did not return after context cancellation")
+	}
+}
+
+func TestDrainSandboxesCompletesAfterSandboxLeaves(t *testing.T) {
+	t.Parallel()
+
+	s := drainTestServer()
+	sbx := drainTestSandbox(t, "lifecycle-1")
+	s.sandboxFactory.Sandboxes.MarkRunning(t.Context(), sbx)
+
+	// Remove the sandbox before draining so the first poll observes an empty
+	// node and the drain completes without waiting on the poll interval.
+	require.True(t, s.sandboxFactory.Sandboxes.MarkStopping(t.Context(), sbx.Runtime.SandboxID, sbx.LifecycleID))
+	s.sandboxFactory.Sandboxes.MarkStopped(t.Context(), sbx)
+
+	done := make(chan error, 1)
+	go func() {
+		done <- s.DrainSandboxes(t.Context())
+	}()
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("DrainSandboxes did not complete after the node emptied")
+	}
+}
+
+func drainTestServer() *Server {
+	return &Server{
+		sandboxFactory: &sandbox.Factory{
+			Sandboxes: sandbox.NewSandboxesMap(),
+		},
+	}
+}
+
+func drainTestSandbox(t *testing.T, lifecycleID string) *sandbox.Sandbox {
+	t.Helper()
+
+	slot, err := network.NewSlot("test", 1, network.Config{}, network.NoopEgressProxy{})
+	require.NoError(t, err)
+
+	return &sandbox.Sandbox{
+		LifecycleID: lifecycleID,
+		Metadata: &sandbox.Metadata{
+			Config:  sandbox.NewConfig(sandbox.Config{}),
+			Runtime: sandbox.RuntimeMetadata{SandboxID: "sandbox-1"},
+		},
+		Resources: &sandbox.Resources{Slot: slot},
+	}
+}

packages/orchestrator/pkg/server/main.go

@@ -43,6 +43,23 @@ const startingSandboxesLimitRefreshInterval = 30 * time.Second
 // in-flight snapshot uploads to finish during shutdown.
 const uploadDrainLogInterval = 10 * time.Second
 
+// sandboxDrainPollInterval is how often the graceful sandbox drain re-checks
+// the live sandbox count during shutdown.
+const sandboxDrainPollInterval = 5 * time.Second
+
+// sandboxDrainLogInterval backs off how often the graceful sandbox drain logs
+// progress so a long-lived node draining for hours does not spam the logs.
+func sandboxDrainLogInterval(elapsed time.Duration) time.Duration {
+	switch {
+	case elapsed < time.Minute:
+		return 5 * time.Second
+	case elapsed < time.Hour:
+		return time.Minute
+	default:
+		return 15 * time.Minute
+	}
+}
+
 type Server struct {
 	orchestrator.UnimplementedSandboxServiceServer
 	orchestrator.UnimplementedChunkServiceServer
@@ -268,6 +285,54 @@ func (s *Server) drainUploads(ctx context.Context, uploadsDone <-chan struct{})
 	}
 }
 
+// DrainSandboxes waits for the live sandboxes on this node to exit on their own
+// during a graceful shutdown, then waits for their lifecycle cleanup to finish.
+// It does not reject new sandbox starts; that admission gating is layered in
+// separately. It returns ctx.Err() if ctx is cancelled before the node empties.
+func (s *Server) DrainSandboxes(ctx context.Context) error {
+	live := s.sandboxFactory.Sandboxes.Count()
+	logger.L().Info(ctx, "starting graceful sandbox drain", zap.Int("live_sandboxes", live))
+
+	ticker := time.NewTicker(sandboxDrainPollInterval)
+	defer ticker.Stop()
+	startedAt := time.Now()
+	lastLoggedAt := startedAt
+
+	for {
+		remaining := s.sandboxFactory.Sandboxes.Count()
+		if remaining == 0 {
+			logger.L().Info(ctx, "graceful sandbox drain complete", zap.Int("live_sandboxes", remaining))
+
+			return s.waitSandboxLifecycles(ctx)
+		}
+
+		select {
+		case <-ctx.Done():
+			logger.L().Warn(ctx, "graceful sandbox drain timed out",
+				zap.Int("remaining_sandboxes", remaining),
+				zap.Error(ctx.Err()),
+			)
+
+			return ctx.Err()
+		case <-ticker.C:
+			now := time.Now()
+			remaining = s.sandboxFactory.Sandboxes.Count()
+			elapsed := now.Sub(startedAt)
+			if remaining > 0 && now.Sub(lastLoggedAt) >= sandboxDrainLogInterval(elapsed) {
+				logger.L().Info(ctx, "waiting for sandbox drain",
+					zap.Int("remaining_sandboxes", remaining),
+					zap.Duration("elapsed", elapsed),
+				)
+				lastLoggedAt = now
+			}
+		}
+	}
+}
+
+func (s *Server) waitSandboxLifecycles(ctx context.Context) error {
+	return s.sandboxFactory.Sandboxes.WaitLifecycles(ctx)
+}
+
 func (s *Server) refreshStartingSandboxesLimit(ctx context.Context) {
 	ticker := time.NewTicker(startingSandboxesLimitRefreshInterval)
 	defer ticker.Stop()