diff --git a/packages/orchestrator/pkg/factories/run.go b/packages/orchestrator/pkg/factories/run.go index ad622e66f6..28a4b4fbd0 100644 --- a/packages/orchestrator/pkg/factories/run.go +++ b/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)) diff --git a/packages/orchestrator/pkg/server/drain_test.go b/packages/orchestrator/pkg/server/drain_test.go new file mode 100644 index 0000000000..3ee5dfd2cc --- /dev/null +++ b/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}, + } +} diff --git a/packages/orchestrator/pkg/server/main.go b/packages/orchestrator/pkg/server/main.go index 0ddb3f105f..7164823386 100644 --- a/packages/orchestrator/pkg/server/main.go +++ b/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()