feat: properly manage interrupt signal on foreground process (#1185)

Co-authored-by: taskbot <[email protected]>

Author: yrobla

cmd/thv/app/run.go

@@ -5,12 +5,16 @@ import (
 	"fmt"
 	"net"
 	"os"
+	"os/signal"
+	"syscall"
+	"time"
 
 	"github.com/spf13/cobra"
 
 	"github.com/stacklok/toolhive/pkg/container"
 	"github.com/stacklok/toolhive/pkg/groups"
 	"github.com/stacklok/toolhive/pkg/logger"
+	"github.com/stacklok/toolhive/pkg/process"
 	"github.com/stacklok/toolhive/pkg/runner"
 	"github.com/stacklok/toolhive/pkg/workloads"
 )
@@ -77,6 +81,27 @@ func init() {
 	AddOIDCFlags(runCmd)
 }
 
+func cleanupAndWait(workloadManager workloads.Manager, name string, cancel context.CancelFunc, errCh <-chan error) {
+	cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cleanupCancel()
+
+	group, err := workloadManager.DeleteWorkloads(cleanupCtx, []string{name})
+	if err != nil {
+		logger.Warnf("Failed to delete workload %q: %v", name, err)
+	} else if group != nil {
+		if err := group.Wait(); err != nil {
+			logger.Warnf("DeleteWorkloads group error for %q: %v", name, err)
+		}
+	}
+
+	cancel()
+	select {
+	case <-errCh:
+	case <-time.After(5 * time.Second):
+		logger.Warnf("Timeout waiting for workload to stop")
+	}
+}
+
 func runCmdFunc(cmd *cobra.Command, args []string) error {
 	ctx := cmd.Context()
 
@@ -107,6 +132,8 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 		workloadName = serverOrImage
 	}
 
+	isForeground := runFlags.Foreground
+
 	if runFlags.Group != "" {
 		groupManager, err := groups.NewManager()
 		if err != nil {
@@ -145,18 +172,35 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 	}
 	workloadManager := workloads.NewManagerFromRuntime(rt)
 
-	if runFlags.Foreground {
-		err = workloadManager.RunWorkload(ctx, runnerConfig)
-	} else {
-		// Run the workload in detached mode
-		err = workloadManager.RunWorkloadDetached(ctx, runnerConfig)
-	}
-	if err != nil {
-		return err
+	if isForeground {
+		return runForeground(ctx, workloadManager, runnerConfig)
 	}
+	return workloadManager.RunWorkloadDetached(ctx, runnerConfig)
+}
 
-	return nil
+func runForeground(ctx context.Context, workloadManager workloads.Manager, runnerConfig *runner.RunConfig) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
+	defer signal.Stop(sigCh)
+
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- workloadManager.RunWorkload(ctx, runnerConfig)
+	}()
+
+	select {
+	case sig := <-sigCh:
+		if !process.IsDetached() {
+			logger.Infof("Received signal: %v, stopping server %q", sig, runnerConfig.BaseName)
+			cleanupAndWait(workloadManager, runnerConfig.BaseName, cancel, errCh)
+		}
+		return nil
+	case err := <-errCh:
+		return err
+	}
 }
 
 // parseCommandArguments processes command-line arguments to find everything after the -- separator