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) } }