feat(agent): use context.WithCancel to handle graceful shutdowns by dajneem23 · Pull Request #613 · uber/kraken

dajneem23 · 2026-05-04T07:26:54Z

What?

Replace the use of log.Fatal throughout the command runner with a context‑based cancellation mechanism, enabling the Run function to implement a proper “stopper” and perform graceful cleanup before exiting.
Fixes #510

Why?

Currently, any fatal error in cmd.Run or its sub‑functions calls log.Fatal, which immediately terminates the process with os.Exit(1). This prevents deferred cleanup (closing files, flushing logs, draining connections) and makes it impossible to implement a graceful shutdown routine.

The proposed pattern uses context.WithCancel to signal that the application should stop. The Run function can:

Pass the cancellable context to all downstream operations.
Cancel the context upon any fatal error (instead of calling log.Fatal).
Listen for context cancellation (<-ctx.Done()) and execute a controlled teardown sequence before allowing the process to exit normally.

How?

Create a root context with context.WithCancel at the beginning of the command’s Run method.
Propagate this context to all long‑running goroutines, servers, and API calls.
Where a fatal error is detected, call cancel() instead of log.Fatal.
In the main goroutine, block on <-ctx.Done() (or an errgroup) to wait for all workers to stop.
Perform any necessary cleanup (closing resources, stopping servers, etc.) after the context is cancelled.
Rely on a deferred cancel() to guarantee the context is released even if the function returns early.
This approach does not change the program’s exit code behavior—top‑level error handling still returns 1—but now all deferred cleanup has a chance to run.

Checklist?

Replace all log.Fatal calls inside cmd/ with cancel() + error return
Add context propagation to sub‑commands and shared utilities
Add signal handling (e.g., SIGINT, SIGTERM) that also triggers cancel()
Ensure any fire‑and‑forget goroutines respect ctx.Done()
Test graceful shutdown on both manual signals and simulated fatal errors

Test coverage

go test -coverprofile=coverage.out ./agent/cmd/...
go tool cover -func=coverage.out

github.com/uber/kraken/agent/cmd/cmd.go:59:     ParseFlags              100.0%
github.com/uber/kraken/agent/cmd/cmd.go:96:     WithConfig              100.0%
github.com/uber/kraken/agent/cmd/cmd.go:101:    WithMetrics             100.0%
github.com/uber/kraken/agent/cmd/cmd.go:106:    WithLogger              100.0%
github.com/uber/kraken/agent/cmd/cmd.go:111:    WithEffect              100.0%
github.com/uber/kraken/agent/cmd/cmd.go:116:    Run                     19.8%
github.com/uber/kraken/agent/cmd/cmd.go:305:    waitForShutdown         100.0%
github.com/uber/kraken/agent/cmd/cmd.go:321:    validateRequiredPorts   100.0%
github.com/uber/kraken/agent/cmd/cmd.go:343:    Chan                    0.0%
github.com/uber/kraken/agent/cmd/cmd.go:347:    Stop                    0.0%
github.com/uber/kraken/agent/cmd/cmd.go:353:    heartbeat               100.0%
total:                                          (statements)            38.9%

…ling; add logging context utilities

dajneem23 · 2026-05-04T07:30:05Z

+	}()
+
+	go func() {
+		if err := nginx.Run(config.Nginx, map[string]interface{}{


@sambhav-jain-16 Should we add a nginx.RunContext variant (accepting a context.Context) to support graceful shutdown ?

Copilot

Pull request overview

This PR attempts to make the agent daemon's startup/runtime path context-aware so fatal conditions and OS signals can trigger shutdown through agent/cmd.Run instead of calling log.Fatal from deep inside the command stack. It fits into the agent entrypoint and command runner, which are responsible for wiring up the agent HTTP server, registry, nginx, metrics, and supporting background work.

Changes:

Add signal-aware root-context handling in agent/main.go and switch cmd.Run to return an error.
Refactor agent/cmd.Run to use context.WithCancel, replace many internal log.Fatal calls with returned errors, and introduce waitForShutdown.
Add unit tests around shutdown coordination helpers and the HTTP server shutdown pattern.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.

File	Description
`agent/main.go`	Wraps agent startup in a signal-aware context and handles `cmd.Run` errors at the process entrypoint.
`agent/cmd/cmd.go`	Refactors the agent runner to use cancelable context flow, background goroutines, and shutdown/error coordination.
`agent/cmd/cmd_test.go`	Adds tests for `waitForShutdown` behavior and the HTTP server shutdown guard.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

sambhav-jain-16 · 2026-05-05T14:40:15Z

+	go func() {
+		<-ctx.Done()
+		if err := httpServer.Shutdown(context.Background()); err != nil {
+			log.Errorf("agent server shutdown: %s", err)


+1. Also, wouldn't it be better to use golang.org/x/sync/errgroup to run the goroutines this way you don't need to write a different function to wait for shutdown?

That make sense, but I'm concerned about the registry.ListenAndServe() comes from github.com/docker/distribution and the Registry.server field is unexported with no public Shutdown() method. The only shutdown hook in that package is an internal SIGTERM listener. So errgroup.Wait() would block forever on the registry goroutine during a clean shutdown.

For Nginx, I added nginx.RunContext which sends SIGQUIT on context cancellation, so nginx drains gracefully. httpServer is also covered via the synchronous httpServer.Shutdown() call after waitForShutdown returns.

The registry is the last piece. I'd suggest we track that as a follow-up: either add a Shutdown wrapper to our docker/distribution fork or expose the *http.Server so we can call Shutdown on it externally. Once that's in place, we can replace all of this with a clean errgroup.

Would you be happy if I opened a follow-up issue for the registry shutdown?

+	go func() {
+		if err := nginx.Run(config.Nginx, map[string]interface{}{
+			"allowed_cidrs": config.AllowedCidrs,
+			"port":          flags.AgentRegistryPort,
+			"registry_server": nginx.GetServer(


 	go func() {
 		if err := registry.ListenAndServe(); err != nil {
 			stopHeartbeat()
-			log.Fatal(err)
+			log.Errorf("registry exited: %s", err)
+			errCh <- err
+			cancel()


+	if err := ctx.Err(); err != nil && err != context.Canceled {
+		log.Fatalf("context error: %s", err)


+	go func() {
+		if err := nginx.Run(config.Nginx, map[string]interface{}{
+			"allowed_cidrs": config.AllowedCidrs,
+			"port":          flags.AgentRegistryPort,
+			"registry_server": nginx.GetServer(


dajneem23 · 2026-05-05T03:08:39Z

@sambhav-jain-16 to fix TestRingLocationsDistribution CI, should we increase sampleSize?

Error:      	Max difference between 0.5 and 0.555 allowed is 0.05, but difference was -0.05500000000000005

sambhav-jain-16

Good start.
Some comments

sambhav-jain-16 · 2026-05-05T13:36:52Z

 	go func() {
-		if err := http.ListenAndServe(addr, agentServer.Handler()); err != nil {
+		if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {


We should cancel on all errors, right?

@sambhav-jain-16

From https://pkg.go.dev/net/http#Server.Shutdown

" Shutdown does not attempt to close nor wait for hijacked connections such as WebSockets. [...] When Shutdown is called, Serve, ListenAndServe, and ListenAndServeTLS immediately return ErrServerClosed. Make sure the program doesn't exit and waits instead for Shutdown to return."

http.ErrServerClosed is returned by ListenAndServe exactly when Shutdown() or Close() is called on the server. Here's the flow in this code:

OS sends SIGTERM/SIGINT → signal.NotifyContext in main.go cancels the top-level ctx

ctx.Done() fires → the shutdown goroutine calls httpServer.Shutdown(context.Background())

Shutdown() drains in-flight requests, then causes ListenAndServe to return http.ErrServerClosed

sambhav-jain-16 · 2026-05-05T13:38:15Z

+	}()
+
+	go func() {
+		if err := nginx.Run(config.Nginx, map[string]interface{}{


sambhav-jain-16 · 2026-05-05T13:39:36Z

 func main() {
-	cmd.Run(cmd.ParseFlags(), cmd.WithEffect(func() {
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()


We shouldn't change the exit code mechanism.

We shouldn't change the exit code mechanism.

Agreed — and we didn't. The exit code mechanism is unchanged from master. log.Fatal(err) is still the last line in main, so any fatal error still terminates with exit code 1 via os.Exit(1) exactly as before.
cmd.Run needs a context. Context that is cancelled on SIGINT/SIGTERM.

cmd.Run now takes ctx—the context is passed down so internal goroutines (HTTP server, nginx) can shut down gracefully when a signal arrives, instead of being killed abruptly.

Is this acceptable?

sambhav-jain-16 · 2026-05-05T14:23:57Z

This should also be an error now, right rather than panic?

This should also be an error now, right rather than panic?

make sense

sambhav-jain-16 · 2026-05-05T14:40:15Z

+	go func() {
+		<-ctx.Done()
+		if err := httpServer.Shutdown(context.Background()); err != nil {
+			log.Errorf("agent server shutdown: %s", err)


+1. Also, wouldn't it be better to use golang.org/x/sync/errgroup to run the goroutines this way you don't need to write a different function to wait for shutdown?

sambhav-jain-16 · 2026-05-05T14:42:00Z

 	addr := fmt.Sprintf(":%d", flags.AgentServerPort)
-	log.Infof("Starting agent server on %s", addr)
+	log.InfoS("starting agent server", "addr", addr)
+	errCh := make(chan error, 2)


Use of buffered channel is not recommended, especially in case of errors, I recommend we use errgoup

feat(agent): enhance Run function with context support and error hand…

7867ba0

…ling; add logging context utilities

Copilot AI review requested due to automatic review settings May 4, 2026 07:26

Copilot started reviewing on behalf of dajneem23 May 4, 2026 07:27 View session

dajneem23 changed the title ~~Fix/510 graceful shutdown~~ Fix(Agent): graceful shutdown May 4, 2026

dajneem23 commented May 4, 2026

View reviewed changes

dajneem23 changed the title ~~Fix(Agent): graceful shutdown~~ Feat(agent): graceful shutdown May 4, 2026

Copilot AI reviewed May 4, 2026

View reviewed changes

dajneem23 changed the title ~~Feat(agent): graceful shutdown~~ feat(agent): use context.WithCancel to handle graceful shutdowns May 4, 2026

sambhav-jain-16 reviewed May 5, 2026

View reviewed changes

dajneem23 requested a review from sambhav-jain-16 May 6, 2026 01:40

feat(agent): use context.WithCancel to handle graceful shutdowns

6f60d13

dajneem23 force-pushed the fix/510-graceful-shutdown branch from 58508b9 to 6f60d13 Compare May 7, 2026 02:08

		if err := ctx.Err(); err != nil && err != context.Canceled {
		log.Fatalf("context error: %s", err)

Conversation

dajneem23 commented May 4, 2026

What?

Why?

How?

Checklist?

Test coverage

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dajneem23 commented May 5, 2026

Uh oh!

sambhav-jain-16 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dajneem23 May 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dajneem23 May 7, 2026 •

edited

Loading