Handle tsc watch and profiling signals

tsc/cmd/tsc/api.go

@@ -5,9 +5,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
-	"os/signal"
 	"strings"
-	"syscall"
 
 	"github.com/microsoft/TypeScript/tsc/internal/api"
 	"github.com/microsoft/TypeScript/tsc/internal/bundled"
@@ -42,7 +40,7 @@ func parseAPIFlags(args []string) (apiFlags, error) {
 	return result, nil
 }
 
-func runAPI(args []string) int {
+func runAPI(ctx context.Context, args []string) int {
 	flags, err := parseAPIFlags(args)
 	if err != nil {
 		return 2
@@ -77,9 +75,6 @@ func runAPI(args []string) int {
 
 	s := api.NewStdioServer(options)
 
-	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
-	defer stop()
-
 	if err := s.Run(ctx); err != nil {
 		fmt.Fprintln(os.Stderr, err)
 		return 1

tsc/cmd/tsc/lsp.go

@@ -6,8 +6,6 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
-	"os/signal"
-	"syscall"
 	"time"
 
 	"github.com/microsoft/TypeScript/tsc/internal/bundled"
@@ -18,7 +16,7 @@ import (
 	"github.com/microsoft/TypeScript/tsc/internal/vfs/osvfs"
 )
 
-func runLSP(args []string) int {
+func runLSP(ctx context.Context, args []string) int {
 	flag := flag.NewFlagSet("lsp", flag.ContinueOnError)
 	stdio := flag.Bool("stdio", false, "use stdio for communication")
 	pprofDir := flag.String("pprofDir", "", "Generate pprof CPU/memory profiles to the given directory.")
@@ -47,7 +45,7 @@ func runLSP(args []string) int {
 	typingsLocation := osvfs.GetGlobalTypingsCacheLocation()
 	cwd := tspath.RootedDirectoryPathFromAbsolute(core.Must(os.Getwd()))
 
-	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	ctx, stop := context.WithCancel(ctx)
 	defer stop()
 
 	s := lsp.NewServer(&lsp.ServerOptions{

tsc/cmd/tsc/main.go

@@ -17,17 +17,60 @@ func main() {
 
 func runMain() int {
 	core.ApplyDebugStackLimit()
+	return runWithSignals(run)
+}
+
+// The first signal requests cancellation and waits for the command's cleanup,
+// including profile flushing. A second signal terminates without waiting.
+func runWithSignals(run func(context.Context) int) int {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	signals := make(chan os.Signal, 2)
+	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
+	defer signal.Stop(signals)
+
+	done := make(chan int, 1)
+	go func() {
+		done <- run(ctx)
+	}()
+
+	var interrupted os.Signal
+	for {
+		select {
+		case sig := <-signals:
+			if interrupted != nil {
+				osutil.ExitWithSignal(sig)
+			}
+			interrupted = sig
+			cancel()
+		case status := <-done:
+			cancel()
+			signal.Stop(signals)
+			if interrupted == nil {
+				// Account for a signal queued as the command finished.
+				select {
+				case sig := <-signals:
+					interrupted = sig
+				default:
+					return status
+				}
+			}
+			osutil.ExitWithSignal(interrupted)
+			return status
+		}
+	}
+}
+
+func run(ctx context.Context) int {
 	args := osutil.Args()[1:]
 	if len(args) > 0 {
 		switch args[0] {
 		case "--lsp":
-			return runLSP(args[1:])
+			return runLSP(ctx, args[1:])
 		case "--api":
-			return runAPI(args[1:])
+			return runAPI(ctx, args[1:])
 		}
 	}
-	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
-	defer stop()
 	result := execute.CommandLine(ctx, newSystem(), args, nil)
 	return int(result.Status)
 }