feat(orchestrator): drain sandboxes during shutdown#3005
Conversation
PR SummaryHigh Risk Overview Reviewed by Cursor Bugbot for commit b2e186e. 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.
| 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".
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
1 similar comment
|
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.
Gate new template build, delete, and layer upload requests behind the shared drain gate, cancel or await in-flight builds via the build cache, and split ServerStore shutdown into Wait (graceful) and ForceStop so callers choose the escalation explicitly. If the graceful template drain is cut short (for example by SHUTDOWN_DRAIN_TIMEOUT), shutdown escalates to a bounded template force stop instead of abandoning in-flight builds.
The sandbox server kept its own drain gate in addition to the shared factory start gate, and drained them in sequence (server gate, then factory gate) so an admitted checkpoint's internal resume would not be rejected when the factory drained first. Collapse to a single factory gate. Create/Checkpoint now enter the factory gate at the handler boundary and mark the context with WithHeldStartGate so the nested ResumeSandbox does not re-enter (or get rejected by) the gate. This keeps the checkpoint's remove-then-resume atomic with respect to drain without a second gate. DrainSandboxes / ForceStopSandboxes wait only on the factory gate, and the starting-limit refresher stops on the factory's drain Done channel. Note: not built/tested locally (cgo userfaultfd can't cross-compile from darwin and Docker is unavailable here); needs linux CI.
TemplateBuildDelete was gated by the drain gate, so once a node entered drain it rejected deletes with Unavailable. Delete is the cancel/kill path (it fails a running build then removes artifacts), so it must keep working while draining. Ungate delete and stop tracking it on the build wait group; adding to the wg after a graceful drain's wg.Wait has started would trip the sync.WaitGroup misuse panic. In-flight delete RPCs are drained by grpcServer.GracefulStop during shutdown instead.
InitLayerFileUpload only mints a signed upload URL and checks existence in shared, content-addressed build storage; the orchestrator is not in the upload data path. A call landing on a draining node is therefore harmless: the client's upload to storage is unaffected by the node draining, and the cached layer is usable by a build on any node. Drop the rejectIfDraining guard (and the now-unused helper) so the upload-init is not rejected mid-drain. Any in-flight RPC is drained by the gRPC server's GracefulStop during shutdown. Note: not built/tested locally (cgo userfaultfd can't cross-compile from darwin and Docker is unavailable here); needs linux CI.
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 b2e186e. Configure here.
|
|
||
| err = forceStopTemplateBuilds() | ||
| } | ||
| } |
There was a problem hiding this comment.
Factory drain precedes template wait
High Severity
Graceful shutdown calls orchestratorService.StartDraining on the shared sandbox factory before tmpl.Wait finishes in-flight template builds. Template build steps call CreateSandbox/ResumeSandbox on that same factory without WithHeldStartGate, so mid-build layer work gets ErrFactoryDraining while shutdown still expects builds to complete gracefully.
Reviewed by Cursor Bugbot for commit b2e186e. 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.