From fc987445529cbcf7d21ae61d286f92381070bee7 Mon Sep 17 00:00:00 2001
From: Nick <nick@myschuch.com>
Date: Fri, 4 Sep 2026 11:49:05 +1000
Subject: [PATCH] Cancel the root context on interrupt and termination signals

Both binaries called cobra's cmd.Execute(), which gives cmd.Context() a plain
context.Background() that is never cancelled. In the sidecar this left the
graceful-shutdown block (server.Shutdown with a 5s timeout) and the tracer
cleanup defers unreachable on a normal docker stop: Go terminates on SIGTERM
without running them.

Wire signal.NotifyContext(os.Interrupt, SIGTERM) into a root context and pass
it via cmd.ExecuteContext, so cmd.Context() cancels on a stop signal and the
existing errgroup-driven shutdown runs. The CLI is largely covered by
bubbletea already, but is switched too for consistency and deterministic
teardown of the collector goroutine.
---
 cmd/compass-sidecar/main.go | 9 ++++++++-
 cmd/compass/main.go         | 9 ++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/cmd/compass-sidecar/main.go b/cmd/compass-sidecar/main.go
index 55e3caf..730a2c6 100644
--- a/cmd/compass-sidecar/main.go
+++ b/cmd/compass-sidecar/main.go
@@ -10,6 +10,8 @@ import (
 	"log/slog"
 	"net/http"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/ilyakaznacheev/cleanenv"
@@ -123,6 +125,11 @@ type Options struct {
 }
 
 func main() {
+	// Cancel the root context on an interrupt or termination signal, so the
+	// HTTP server drains and the tracers run their cleanup on shutdown.
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
 	o := Options{}
 
 	cmd := &cobra.Command{
@@ -272,7 +279,7 @@ func main() {
 
 	// Cobra prints the error, so exit quietly rather than panicking with a
 	// stack trace over the top of it.
-	if err := cmd.Execute(); err != nil {
+	if err := cmd.ExecuteContext(ctx); err != nil {
 		os.Exit(1)
 	}
 }
diff --git a/cmd/compass/main.go b/cmd/compass/main.go
index a6db572..78acf39 100644
--- a/cmd/compass/main.go
+++ b/cmd/compass/main.go
@@ -5,8 +5,10 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"os/signal"
 	"regexp"
 	"strings"
+	"syscall"
 
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
@@ -60,6 +62,11 @@ func main() {
 		termenv.NewOutput(os.Stdout).EnvColorProfile(),
 	))
 
+	// Cancel the root context on an interrupt or termination signal, so the
+	// collector goroutine is torn down deterministically on shutdown.
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
 	o := Options{}
 
 	cmd := &cobra.Command{
@@ -141,7 +148,7 @@ func main() {
 
 	// Cobra prints the error, so exit quietly rather than panicking with a
 	// stack trace over the top of it.
-	if err := cmd.Execute(); err != nil {
+	if err := cmd.ExecuteContext(ctx); err != nil {
 		os.Exit(1)
 	}
 }