feat(orchestrator): graceful template-build drain on shutdown

[wj-e2b]

Actually honor config.ShutdownDrainTimeout. Before, wg.Wait() and time.Sleep(15s) ignored closeCtx entirely — the select only checked the context once at entry. So a configured shutdown timeout did nothing during the template drain: the process would block past the deadline until builds finished on their own (or indefinitely on a stuck build), then always burn the full 15s.

Also makes Wait return a meaningful error at the deadline. That return value is unused, but it's the signal the later commits wires into forceStopTemplateBuilds().

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[cursor]

PR Summary

Medium Risk
Changes orchestrator shutdown timing and error behavior; mis-handling could shorten drain or leave builds running, though behavior is covered by new tests.

Overview
Template manager shutdown drain now respects the shutdown context for the whole wait: in-flight builds are waited on via a new context-aware wait-group helper instead of blocking until they finish, and the post-build consumer grace period can exit early when the context is cancelled or times out instead of always sleeping 15 seconds. Wait returns wrapped context errors when the deadline hits during either phase.

^{Reviewed by Cursor Bugbot for commit dcf9cf8. Bugbot is set up for automated code reviews on this repo. Configure here.}

[codecov]

❌ 5 Tests Failed:

Tests completed	Failed	Passed	Skipped
3038	5	3033	8

View the top 3 failed test(s) by shortest run time

github.com/e2b-dev/infra/tests/integration/internal/tests/envd::TestListDir/depth_1_lists_root_directory

Stack Traces | 0.01s run time

=== RUN   TestListDir/depth_1_lists_root_directory
=== PAUSE TestListDir/depth_1_lists_root_directory
=== CONT  TestListDir/depth_1_lists_root_directory
    filesystem_test.go:96: 
        	Error Trace:	.../tests/envd/filesystem_test.go:96
        	Error:      	Received unexpected error:
        	            	unavailable: 502 Bad Gateway
        	Test:       	TestListDir/depth_1_lists_root_directory
--- FAIL: TestListDir/depth_1_lists_root_directory (0.01s)

github.com/e2b-dev/infra/tests/integration/internal/tests/envd::TestListDir/depth_2_lists_first_level_of_subdirectories_(in_this_case_the_root_directory)

Stack Traces | 0.01s run time

=== RUN   TestListDir/depth_2_lists_first_level_of_subdirectories_(in_this_case_the_root_directory)
=== PAUSE TestListDir/depth_2_lists_first_level_of_subdirectories_(in_this_case_the_root_directory)
=== CONT  TestListDir/depth_2_lists_first_level_of_subdirectories_(in_this_case_the_root_directory)
    filesystem_test.go:96: 
        	Error Trace:	.../tests/envd/filesystem_test.go:96
        	Error:      	Received unexpected error:
        	            	unavailable: 502 Bad Gateway
        	Test:       	TestListDir/depth_2_lists_first_level_of_subdirectories_(in_this_case_the_root_directory)
--- FAIL: TestListDir/depth_2_lists_first_level_of_subdirectories_(in_this_case_the_root_directory) (0.01s)

github.com/e2b-dev/infra/tests/integration/internal/tests/envd::TestListDir/depth_3_lists_all_directories_and_files

Stack Traces | 0.01s run time

=== RUN   TestListDir/depth_3_lists_all_directories_and_files
=== PAUSE TestListDir/depth_3_lists_all_directories_and_files
=== CONT  TestListDir/depth_3_lists_all_directories_and_files
    filesystem_test.go:96: 
        	Error Trace:	.../tests/envd/filesystem_test.go:96
        	Error:      	Received unexpected error:
        	            	unavailable: 502 Bad Gateway
        	Test:       	TestListDir/depth_3_lists_all_directories_and_files
--- FAIL: TestListDir/depth_3_lists_all_directories_and_files (0.01s)

github.com/e2b-dev/infra/tests/integration/internal/tests/envd::TestListDir/depth_0_lists_only_root_directory

Stack Traces | 0.02s run time

=== RUN   TestListDir/depth_0_lists_only_root_directory
=== PAUSE TestListDir/depth_0_lists_only_root_directory
=== CONT  TestListDir/depth_0_lists_only_root_directory
    filesystem_test.go:96: 
        	Error Trace:	.../tests/envd/filesystem_test.go:96
        	Error:      	Received unexpected error:
        	            	unavailable: 502 Bad Gateway
        	Test:       	TestListDir/depth_0_lists_only_root_directory
--- FAIL: TestListDir/depth_0_lists_only_root_directory (0.02s)

github.com/e2b-dev/infra/tests/integration/internal/tests/envd::TestListDir

Stack Traces | 3.66s run time

=== RUN   TestListDir
=== PAUSE TestListDir
=== CONT  TestListDir
--- FAIL: TestListDir (3.66s)

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[gemini-code-assist]

Code Review

Preserving the parent context's deadline for rollbacks in ApplyAllOrNone means that if the parent context has already timed out, the rollback context will be immediately canceled. This causes context-aware cleanup operations inside the rollbacks to fail instantly with a deadline exceeded error, preventing proper cleanup and leaving the system in an inconsistent state.

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.

[cursor]

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 b759930. Configure here.}

[rguliyev]

can we pass it via args or envs?

[wj-e2b]

None of the existing code hoists this to CLI env vars today. We can change it if it ever needs to be tweaked.