Gracefully Shutdown Foreground Server on Interrupt (#2863)

The solution here avoids the competing signal handlers by:

Removing the os.Exit(0) from the default handler.
When the default handler receives a signal, cancel the context that is plumbed to the run command.
Remove the run signal handler and instead rely on context cancellation.
This solution is not without risk: Commands that don’t check for context cancellation and relied on the os.Exit(0) may now continue running until they hit code that does respect the context cancellation. Based on my understanding of the code, thv run --foreground is one of the few commands that can be long running, so I think this risk is relatively small.

Beyond the scope of this pull request, I think we should prefer to rely on context cancellation rather than using signal.Notify deep in the implementation of commands. This “competing signal handlers” problem likely exists elsewhere: I count 10 uses of signal.Notify.

---------

Signed-off-by: Jeremy Drouillard <jeremy@stacklok.com>

Author: jerm-dro

cmd/thv/main.go

@@ -2,6 +2,7 @@
 package main
 
 import (
+	"context"
 	"os"
 	"os/signal"
 	"syscall"
@@ -23,7 +24,7 @@ func main() {
 	logger.Initialize()
 
 	// Setup signal handling for graceful cleanup
-	setupSignalHandler()
+	ctx := setupSignalHandler()
 
 	// Clean up stale lock files on startup
 	cleanupStaleLockFiles()
@@ -47,8 +48,10 @@ func main() {
 		migration.CheckAndPerformDefaultGroupMigration()
 	}
 
+	cmd := app.NewRootCmd(!app.IsCompletionCommand(os.Args) && !runtime.IsKubernetesRuntime())
+
 	// Skip update check for completion command or if we are running in kubernetes
-	if err := app.NewRootCmd(!app.IsCompletionCommand(os.Args) && !runtime.IsKubernetesRuntime()).Execute(); err != nil {
+	if err := cmd.ExecuteContext(ctx); err != nil {
 		// Clean up any remaining lock files on error exit
 		lockfile.CleanupAllLocks()
 		os.Exit(1)
@@ -59,16 +62,19 @@ func main() {
 }
 
 // setupSignalHandler configures signal handling to ensure lock files are cleaned up
-func setupSignalHandler() {
+func setupSignalHandler() context.Context {
 	sigCh := make(chan os.Signal, 1)
 	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
 
+	ctx, cancel := context.WithCancel(context.Background())
 	go func() {
 		<-sigCh
 		logger.Debugf("Received signal, cleaning up lock files...")
 		lockfile.CleanupAllLocks()
-		os.Exit(0)
+		cancel()
 	}()
+
+	return ctx
 }
 
 // cleanupStaleLockFiles removes stale lock files from known directories on startup

pkg/runner/runner.go

@@ -8,9 +8,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
-	"os/signal"
 	"strings"
-	"syscall"
 	"time"
 
 	"golang.org/x/oauth2"
@@ -354,10 +352,6 @@ func (r *Runner) Run(ctx context.Context) error {
 		logger.Info("Press Ctrl+C to stop or wait for container to exit")
 	}
 
-	// Set up signal handling
-	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
-
 	// Create a done channel to signal when the server has been stopped
 	doneCh := make(chan struct{})
 
@@ -399,8 +393,8 @@ func (r *Runner) Run(ctx context.Context) error {
 
 	// Wait for either a signal or the done channel to be closed
 	select {
-	case sig := <-sigCh:
-		stopMCPServer(fmt.Sprintf("Received signal %s", sig))
+	case <-ctx.Done():
+		stopMCPServer("Context cancelled")
 	case <-doneCh:
 		// The transport has already been stopped (likely by the container exit)
 		// Clean up the PID file and state