feat(orchestrator): drain sandboxes during shutdown#3005
Conversation
Introduce a shared draingate.Gate (counter plus notification channel) and use it in the sandbox factory and gRPC server to reject new sandbox starts while draining, wait for in-flight starts, and drain or force-stop live sandboxes before closers run. Forced shutdown preserves buffered close errors on context cancellation and avoids duplicate final-pass errors. Adds utils.WaitGroupWait to wait on a WaitGroup with context cancellation. The graceful drain phase is bounded by SHUTDOWN_DRAIN_TIMEOUT; when it expires the drain escalates to a forced sandbox shutdown. By default the drain waits forever, until sandboxes exit on their own or a force-stop API call empties the node.
PR SummaryMedium Risk Overview Graceful shutdown polls until live sandboxes reach zero and waits on lifecycle cleanup; Reviewed by Cursor Bugbot for commit 5d6b2e7. Bugbot is set up for automated code reviews on this repo. Configure here. |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 5d6b2e7. Configure here.
|
|
||
| if err := s.rejectIfDraining(ctx, "sandbox-create-before-start"); err != nil { | ||
| return nil, err | ||
| } |
There was a problem hiding this comment.
Draining checks abort in-flight starts
Medium Severity
Handlers that already acquired the server drain gate via enterSandboxStart can still fail when rejectIfDraining or waitForAcquire runs after shutdown calls StartDraining. That rejects work the drain gate was meant to let finish, including long create and checkpoint paths before ResumeSandbox.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 5d6b2e7. Configure here.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
Compilation errors exist in packages/orchestrator/pkg/server/main.go and packages/orchestrator/pkg/draingate/gate_test.go where wg.Go() is called on a sync.WaitGroup which does not have a Go method. In packages/orchestrator/pkg/server/main.go, calling context.WithoutCancel(ctx) strips the configured shutdown deadline, which can cause the orchestrator to hang indefinitely if a sandbox cleanup hangs. Additionally, the spin-loop using runtime.Gosched() in packages/shared/pkg/utils/waitgroup.go is flaky and can cause premature context cancellation errors under heavy CPU load.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| wg.Go(func() { | ||
| sbxLog := logger.L().With( | ||
| logger.WithSandboxID(sbx.Runtime.SandboxID), | ||
| logger.WithLifecycleID(sbx.LifecycleID), | ||
| logger.WithSandboxIP(sbx.Slot.HostIPString()), | ||
| ) | ||
| sbxLog.Warn(cleanupCtx, "force stopping sandbox during orchestrator shutdown") | ||
|
|
||
| marked := s.sandboxFactory.Sandboxes.MarkStopping(cleanupCtx, sbx.Runtime.SandboxID, sbx.LifecycleID) | ||
| if !marked { | ||
| sbxLog.Info(cleanupCtx, "sandbox was already removed from live map before force stop") | ||
| } | ||
|
|
||
| if err := sbx.Close(cleanupCtx); err != nil { | ||
| errCh <- fmt.Errorf("close sandbox %s/%s: %w", sbx.Runtime.SandboxID, sbx.LifecycleID, err) | ||
| sbxLog.Error(cleanupCtx, "failed to force close sandbox", zap.Error(err)) | ||
| } | ||
|
|
||
| sbxLog.Info(cleanupCtx, "forced sandbox cleanup complete") | ||
| }) |
There was a problem hiding this comment.
The sync.WaitGroup type does not have a Go method, which will cause a compilation error. To run the cleanup concurrently, you should use standard wg.Add(1) before the goroutine and defer wg.Done() inside it.
wg.Add(1)
go func() {
defer wg.Done()
sbxLog := logger.L().With(
logger.WithSandboxID(sbx.Runtime.SandboxID),
logger.WithLifecycleID(sbx.LifecycleID),
logger.WithSandboxIP(sbx.Slot.HostIPString()),
)
sbxLog.Warn(cleanupCtx, "force stopping sandbox during orchestrator shutdown")
marked := s.sandboxFactory.Sandboxes.MarkStopping(cleanupCtx, sbx.Runtime.SandboxID, sbx.LifecycleID)
if !marked {
sbxLog.Info(cleanupCtx, "sandbox was already removed from live map before force stop")
}
if err := sbx.Close(cleanupCtx); err != nil {
errCh <- fmt.Errorf("close sandbox %s/%s: %w", sbx.Runtime.SandboxID, sbx.LifecycleID, err)
sbxLog.Error(cleanupCtx, "failed to force close sandbox", zap.Error(err))
}
sbxLog.Info(cleanupCtx, "forced sandbox cleanup complete")
}()| for range 100 { | ||
| wg.Go(func() { | ||
| <-start | ||
| release, err := g.Enter() | ||
| if err != nil { | ||
| rejected <- err | ||
|
|
||
| return | ||
| } | ||
|
|
||
| entered <- release | ||
| }) | ||
| } |
There was a problem hiding this comment.
The sync.WaitGroup type does not have a Go method, which will cause a compilation error in this test. You should use standard wg.Add(1) and defer wg.Done() to manage the concurrent test goroutines.
| for range 100 { | |
| wg.Go(func() { | |
| <-start | |
| release, err := g.Enter() | |
| if err != nil { | |
| rejected <- err | |
| return | |
| } | |
| entered <- release | |
| }) | |
| } | |
| for range 100 { | |
| wg.Add(1) | |
| go func() { | |
| defer wg.Done() | |
| <-start | |
| release, err := g.Enter() | |
| if err != nil { | |
| rejected <- err | |
| return | |
| } | |
| entered <- release | |
| }() | |
| } |
| stopped := make(map[string]struct{}) | ||
| var errs []error | ||
|
|
||
| cleanupCtx := context.WithoutCancel(ctx) |
There was a problem hiding this comment.
Calling context.WithoutCancel(ctx) here strips the 30-second forceShutdownTimeout deadline that was explicitly configured during shutdown. This causes the subsequent sandbox close operations to run without any timeout, meaning the orchestrator can hang indefinitely if a sandbox cleanup hangs. You should use the passed ctx directly to ensure the force-shutdown phase is properly bounded.
| cleanupCtx := context.WithoutCancel(ctx) | |
| cleanupCtx := ctx |
| for range 10 { | ||
| select { | ||
| case <-done: | ||
| return nil | ||
| default: | ||
| runtime.Gosched() | ||
| } | ||
| } | ||
|
|
||
| select { | ||
| case <-ctx.Done(): | ||
| return fmt.Errorf("waiting for wait group: %w", ctx.Err()) | ||
| case <-done: | ||
| return nil | ||
| } |
There was a problem hiding this comment.
The spin-loop using runtime.Gosched() is flaky and does not guarantee that the goroutine closing done will be scheduled and executed within 10 iterations. Under heavy CPU load or thread starvation, the loop can easily fall through to the select block and return a context error even if the WaitGroup was already completed, making the associated test fragile. You should use a standard select block to wait for either the WaitGroup completion or context cancellation.
select {
case <-done:
return nil
case <-ctx.Done():
return fmt.Errorf("waiting for wait group: %w", ctx.Err())
}
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5d6b2e7976
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| defer func(g *errgroup.Group) { | ||
| err := g.Wait() | ||
| if err != nil { | ||
| if err != nil && !isServiceDoneError(err) { |
There was a problem hiding this comment.
When a service exits because f() returned an actual error (for example a non-benign grpcServer.Serve or cmuxServer.Serve failure), startService sends that error on serviceError but then returns only serviceDoneError; the serviceError select branch only logs it and never flips success to false. With this new guard, the deferred g.Wait() now also ignores the only error value it can see, so Run can report a clean shutdown and remove the lock file after a real service failure. Either preserve the original service error in the errgroup result or mark success = false when the serviceError branch receives a non-nil error.
Useful? React with 👍 / 👎.
1 participant
Introduce a shared draingate.Gate (counter plus notification channel) and use it in the sandbox factory and gRPC server to reject new sandbox starts while draining, wait for in-flight starts, and drain or force-stop live sandboxes before closers run. Forced shutdown preserves buffered close errors on context cancellation and avoids duplicate final-pass errors. Adds utils.WaitGroupWait to wait on a WaitGroup with context cancellation.
The graceful drain phase is bounded by SHUTDOWN_DRAIN_TIMEOUT; when it expires the drain escalates to a forced sandbox shutdown. By default the drain waits forever, until sandboxes exit on their own or a force-stop API call empties the node.