feat(orchestrator): drain sandboxes during shutdown#3005
Conversation
PR SummaryHigh Risk Overview A reusable Shutdown wiring adds optional Reviewed by Cursor Bugbot for commit 1aa9982. 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}
![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.
| 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())
}
❌ 1 Tests Failed:
View the top 1 failed test(s) by shortest run time
To view more test analytics, go to the Test Analytics Dashboard |
![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 👍 / 👎.
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
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.
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
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 1aa9982. Configure here.
| s.StartDraining(ctx) | ||
| // The sandbox factory is shared by API sandboxes and template-build sandboxes. | ||
| // Drain it before waiting so no new starts can enter while shutdown proceeds. | ||
| s.sandboxFactory.StartDraining(ctx) |
There was a problem hiding this comment.
Force stop drains factory early
Medium Severity
ForceStopSandboxes calls sandboxFactory.StartDraining before in-flight gRPC sandbox starts finish, while DrainSandboxes waits on the server drain gate first so checkpoint resume can still enter the factory. A FORCE_STOP shutdown (or any path that only uses force stop) can reject ResumeSandbox mid-checkpoint or mid-create after the server gate already admitted the RPC.
Reviewed by Cursor Bugbot for commit 1aa9982. Configure here.
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.