Handle tsc watch and profiling signals

tsc/cmd/tsc/api.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"os"
 	"strings"
-	"syscall"
 
 	"github.com/microsoft/TypeScript/tsc/internal/api"
 	"github.com/microsoft/TypeScript/tsc/internal/bundled"
@@ -41,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
@@ -76,9 +75,6 @@ func runAPI(args []string) int {
 
 	s := api.NewStdioServer(options)
 
-	ctx, stop := 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

@@ -5,7 +5,6 @@ import (
 	"flag"
 	"fmt"
 	"os"
-	"syscall"
 	"time"
 
 	"github.com/microsoft/TypeScript/tsc/internal/bundled"
@@ -16,7 +15,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.")
@@ -45,7 +44,7 @@ func runLSP(args []string) int {
 	typingsLocation := osvfs.GetGlobalTypingsCacheLocation()
 	cwd := tspath.RootedDirectoryPathFromAbsolute(core.Must(getCurrentDirectory()))
 
-	ctx, stop := 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,20 +17,63 @@ 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)
+	stopSignals := notifySignals(signals, syscall.SIGINT, syscall.SIGTERM)
+	defer stopSignals()
+
+	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()
+			stopSignals()
+			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:])
 		}
 	}
 	// The language server and API take heap profiles on demand, so only the
 	// command line can skip sampling.
 	pprof.DisableMemProfiling()
-	ctx, stop := notifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
-	defer stop()
 	result := execute.CommandLine(ctx, newSystem(), args, nil)
 	return int(result.Status)
 }

tsc/cmd/tsc/platform.go

@@ -16,8 +16,9 @@ func getCurrentDirectory() (string, error) {
 	return os.Getwd()
 }
 
-func notifyContext(parent context.Context, signals ...os.Signal) (context.Context, context.CancelFunc) {
-	return signal.NotifyContext(parent, signals...)
+func notifySignals(ch chan<- os.Signal, signals ...os.Signal) func() {
+	signal.Notify(ch, signals...)
+	return func() { signal.Stop(ch) }
 }
 
 func getNpmInstall() func(ctx context.Context, cwd string, args []string) ([]byte, error) {

tsc/cmd/tsc/platform_wasip1.go

@@ -69,8 +69,8 @@ func getCurrentDirectory() (string, error) {
 	return cwd, nil
 }
 
-func notifyContext(parent context.Context, _ ...os.Signal) (context.Context, context.CancelFunc) {
-	return context.WithCancel(parent)
+func notifySignals(chan<- os.Signal, ...os.Signal) func() {
+	return func() {}
 }
 
 func getNpmInstall() func(context.Context, string, []string) ([]byte, error) {