Handle tsc watch and profiling signals

tsc/cmd/tsc/main.go

@@ -3,8 +3,6 @@ package main
 import (
 	"context"
 	"os"
-	"os/signal"
-	"syscall"
 
 	"github.com/microsoft/TypeScript/tsc/internal/core"
 	"github.com/microsoft/TypeScript/tsc/internal/execute"
@@ -26,8 +24,8 @@ func runMain() int {
 			return runAPI(args[1:])
 		}
 	}
-	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
-	defer stop()
-	result := execute.CommandLine(ctx, newSystem(), args, nil)
+	result := execute.CommandLineWithOptions(context.Background(), newSystem(), args, nil, execute.CommandLineOptions{
+		WatchContext: osutil.NotifyTerminationContext,
+	})
 	return int(result.Status)
 }

tsc/cmd/tsc/sys_unix_test.go

@@ -5,18 +5,29 @@ package main
 import (
 	"bufio"
 	"bytes"
+	"errors"
 	"fmt"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"syscall"
 	"testing"
 	"time"
 
+	"github.com/microsoft/TypeScript/tsc/internal/osutil"
 	"gotest.tools/v3/assert"
 )
 
+func TestMain(m *testing.M) {
+	if args := os.Getenv("TSGO_WATCH_SIGNAL_HELPER"); args != "" {
+		os.Args = append([]string{os.Args[0]}, strings.Fields(args)...)
+		os.Exit(runMain())
+	}
+	os.Exit(m.Run())
+}
+
 func TestChildProcessCloseDoesNotWaitForLauncherDescendants(t *testing.T) {
 	const (
 		launcherArg   = "child-process-launcher"
@@ -67,3 +78,84 @@ func TestChildProcessCloseDoesNotWaitForLauncherDescendants(t *testing.T) {
 		_ = syscall.Kill(descendantPID, syscall.SIGKILL)
 	}
 }
+
+func TestWatchTerminatesOnInterrupt(t *testing.T) {
+	t.Parallel()
+
+	executable, err := osutil.Executable()
+	assert.NilError(t, err)
+
+	for _, test := range []struct {
+		name string
+		args string
+	}{
+		{name: "watch", args: "--watch --project tsconfig.json"},
+		{name: "buildWatch", args: "--build --watch tsconfig.json"},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			projectDir := t.TempDir()
+			assert.NilError(t, os.WriteFile(filepath.Join(projectDir, "index.ts"), []byte("export const value = 1;\n"), 0o666))
+			assert.NilError(t, os.WriteFile(filepath.Join(projectDir, "tsconfig.json"), []byte(`{"compilerOptions":{"pretty":false},"files":["index.ts"]}`), 0o666))
+
+			output, outputErr := os.CreateTemp(t.TempDir(), "watch-output")
+			assert.NilError(t, outputErr)
+			defer output.Close()
+
+			cmd := exec.Command(executable, "-test.run=^TestWatchTerminatesOnInterrupt$")
+			cmd.Dir = projectDir
+			cmd.Env = append(os.Environ(), "TSGO_WATCH_SIGNAL_HELPER="+test.args)
+			cmd.Stdout = output
+			cmd.Stderr = output
+			assert.NilError(t, cmd.Start())
+			t.Cleanup(func() {
+				if cmd.ProcessState == nil {
+					_ = cmd.Process.Kill()
+					_ = cmd.Wait()
+				}
+			})
+
+			waitForWatchOutput(t, output.Name(), "Watching for file changes.")
+			assert.NilError(t, cmd.Process.Signal(os.Interrupt))
+
+			waitDone := make(chan error, 1)
+			go func() {
+				waitDone <- cmd.Wait()
+			}()
+			var waitErr error
+			select {
+			case result := <-waitDone:
+				waitErr = result
+			case <-time.After(10 * time.Second):
+				_ = cmd.Process.Kill()
+				<-waitDone
+				t.Fatal("timed out waiting for watch process to terminate")
+			}
+			var exitErr *exec.ExitError
+			if !errors.As(waitErr, &exitErr) {
+				t.Fatalf("watch process returned %v instead of terminating from SIGINT", waitErr)
+			}
+			status := exitErr.ProcessState.Sys().(syscall.WaitStatus)
+			if !status.Signaled() || status.Signal() != syscall.SIGINT {
+				t.Fatalf("watch process exited with %v instead of SIGINT", status)
+			}
+		})
+	}
+}
+
+func waitForWatchOutput(t *testing.T, path string, expected string) {
+	t.Helper()
+	deadline := time.Now().Add(10 * time.Second)
+	for {
+		output, err := os.ReadFile(path)
+		assert.NilError(t, err)
+		if strings.Contains(string(output), expected) {
+			return
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("timed out waiting for %q in output:\n%s", expected, output)
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+}

tsc/internal/execute/tsc.go

@@ -50,17 +50,25 @@ func stopTracing(sys tsc.System, tr *tracing.Tracing) {
 	}
 }
 
+type CommandLineOptions struct {
+	WatchContext func(context.Context) (context.Context, context.CancelFunc)
+}
+
 func CommandLine(ctx context.Context, sys tsc.System, commandLineArgs []string, testing tsc.CommandLineTesting) tsc.CommandLineResult {
+	return CommandLineWithOptions(ctx, sys, commandLineArgs, testing, CommandLineOptions{})
+}
+
+func CommandLineWithOptions(ctx context.Context, sys tsc.System, commandLineArgs []string, testing tsc.CommandLineTesting, options CommandLineOptions) tsc.CommandLineResult {
 	if len(commandLineArgs) > 0 {
 		switch strings.ToLower(commandLineArgs[0]) {
 		case "-b", "--b", "-build", "--build":
-			return tscBuildCompilation(ctx, sys, tsoptions.ParseBuildCommandLine(commandLineArgs, sys), testing)
+			return tscBuildCompilation(ctx, sys, tsoptions.ParseBuildCommandLine(commandLineArgs, sys), testing, options)
 			// case "-f":
 			// 	return fmtMain(sys, commandLineArgs[1], commandLineArgs[1])
 		}
 	}
 
-	return tscCompilation(ctx, sys, tsoptions.ParseCommandLine(commandLineArgs, sys), testing)
+	return tscCompilation(ctx, sys, tsoptions.ParseCommandLine(commandLineArgs, sys), testing, options)
 }
 
 func fmtMain(sys tsc.System, input, output string) tsc.ExitStatus {
@@ -88,9 +96,10 @@ func fmtMain(sys tsc.System, input, output string) tsc.ExitStatus {
 	return tsc.ExitStatusSuccess
 }
 
-func tscBuildCompilation(ctx context.Context, sys tsc.System, buildCommand *tsoptions.ParsedBuildCommandLine, testing tsc.CommandLineTesting) tsc.CommandLineResult {
+func tscBuildCompilation(ctx context.Context, sys tsc.System, buildCommand *tsoptions.ParsedBuildCommandLine, testing tsc.CommandLineTesting, options CommandLineOptions) tsc.CommandLineResult {
 	locale := buildCommand.Locale()
 	reportDiagnostic := tsc.CreateDiagnosticReporter(sys, sys.Writer(), locale, buildCommand.CompilerOptions)
+	profiled := false
 
 	if len(buildCommand.Errors) > 0 {
 		for _, err := range buildCommand.Errors {
@@ -103,6 +112,7 @@ func tscBuildCompilation(ctx context.Context, sys tsc.System, buildCommand *tsop
 		// !!! stderr?
 		profileSession := pprof.BeginProfiling(pprofDir, sys.Writer())
 		defer profileSession.Stop()
+		profiled = true
 	}
 
 	if buildCommand.CompilerOptions.Help.IsTrue() {
@@ -111,6 +121,11 @@ func tscBuildCompilation(ctx context.Context, sys tsc.System, buildCommand *tsop
 		return tsc.CommandLineResult{Status: tsc.ExitStatusSuccess}
 	}
 
+	if buildCommand.CompilerOptions.Watch.IsTrue() && !profiled && options.WatchContext != nil {
+		var stop context.CancelFunc
+		ctx, stop = options.WatchContext(ctx)
+		defer stop()
+	}
 	orchestrator := build.NewOrchestrator(build.Options{
 		Sys:     sys,
 		Command: buildCommand,
@@ -119,10 +134,11 @@ func tscBuildCompilation(ctx context.Context, sys tsc.System, buildCommand *tsop
 	return orchestrator.Start(ctx)
 }
 
-func tscCompilation(ctx context.Context, sys tsc.System, commandLine *tsoptions.ParsedCommandLine, testing tsc.CommandLineTesting) tsc.CommandLineResult {
+func tscCompilation(ctx context.Context, sys tsc.System, commandLine *tsoptions.ParsedCommandLine, testing tsc.CommandLineTesting, options CommandLineOptions) tsc.CommandLineResult {
 	configFileName := ""
 	locale := commandLine.Locale()
 	reportDiagnostic := tsc.CreateDiagnosticReporter(sys, sys.Writer(), locale, commandLine.CompilerOptions())
+	profiled := false
 
 	if len(commandLine.Errors) > 0 {
 		for _, e := range commandLine.Errors {
@@ -135,6 +151,7 @@ func tscCompilation(ctx context.Context, sys tsc.System, commandLine *tsoptions.
 		// !!! stderr?
 		profileSession := pprof.BeginProfiling(pprofDir, sys.Writer())
 		defer profileSession.Stop()
+		profiled = true
 	}
 
 	if commandLine.CompilerOptions().Init.IsTrue() {
@@ -231,6 +248,11 @@ func tscCompilation(ctx context.Context, sys tsc.System, commandLine *tsoptions.
 		return tsc.CommandLineResult{Status: tsc.ExitStatusSuccess}
 	}
 	if configForCompilation.CompilerOptions().Watch.IsTrue() {
+		if !profiled && options.WatchContext != nil {
+			var stop context.CancelFunc
+			ctx, stop = options.WatchContext(ctx)
+			defer stop()
+		}
 		watcher := createWatcher(
 			sys,
 			configForCompilation,

tsc/internal/osutil/osutil.go

@@ -1,5 +1,12 @@
 package osutil
 
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
 // Args returns the command-line arguments with platform-specific launcher details removed.
 func Args() []string {
 	return args()
@@ -9,3 +16,32 @@ func Args() []string {
 func Executable() (string, error) {
 	return executable()
 }
+
+// NotifyTerminationSignals registers for process termination signals.
+func NotifyTerminationSignals() (<-chan os.Signal, func()) {
+	ch := make(chan os.Signal, 1)
+	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
+	return ch, func() {
+		signal.Stop(ch)
+	}
+}
+
+// NotifyTerminationContext returns a context that terminates the process when a
+// process termination signal arrives.
+func NotifyTerminationContext(parent context.Context) (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithCancel(parent)
+	sigCh, stopSignals := NotifyTerminationSignals()
+	go func() {
+		select {
+		case sig := <-sigCh:
+			ReraiseSignal(sig)
+			cancel()
+		case <-ctx.Done():
+			return
+		}
+	}()
+	return ctx, func() {
+		stopSignals()
+		cancel()
+	}
+}

tsc/internal/osutil/signal_other.go

@@ -0,0 +1,8 @@
+//go:build !unix
+
+package osutil
+
+import "os"
+
+// ReraiseSignal is unsupported on this platform.
+func ReraiseSignal(sig os.Signal) {}

tsc/internal/osutil/signal_unix.go

@@ -0,0 +1,25 @@
+//go:build unix
+
+package osutil
+
+import (
+	"os"
+	"os/signal"
+	"runtime"
+	"syscall"
+)
+
+// ReraiseSignal restores the default handler and sends sig to the current process.
+func ReraiseSignal(sig os.Signal) {
+	syscallSignal, ok := sig.(syscall.Signal)
+	if !ok {
+		return
+	}
+	signal.Reset(syscallSignal)
+	if err := syscall.Kill(os.Getpid(), syscallSignal); err != nil {
+		return
+	}
+	for {
+		runtime.Gosched()
+	}
+}

tsc/internal/pprof/pprof.go

@@ -10,13 +10,18 @@ import (
 	"runtime/pprof"
 	"sync"
 	"time"
+
+	"github.com/microsoft/TypeScript/tsc/internal/osutil"
 )
 
 type ProfileSession struct {
 	cpuFilePath string
 	memFilePath string
 	cpuFile     *os.File
 	logWriter   io.Writer
+	stopSignals func()
+	done        chan struct{}
+	stopOnce    sync.Once
 }
 
 // BeginProfiling starts CPU and memory profiling, writing the profiles to the specified directory.
@@ -38,31 +43,55 @@ func BeginProfiling(profileDir string, logWriter io.Writer) *ProfileSession {
 		panic(err)
 	}
 
-	return &ProfileSession{
+	session := &ProfileSession{
 		cpuFilePath: cpuProfilePath,
 		memFilePath: memProfilePath,
 		cpuFile:     cpuFile,
 		logWriter:   logWriter,
-	}
+		done:        make(chan struct{}),
+	}
+	sigCh, stopSignals := osutil.NotifyTerminationSignals()
+	session.stopSignals = stopSignals
+	go func() {
+		select {
+		case sig := <-sigCh:
+			defer func() {
+				osutil.ReraiseSignal(sig)
+				os.Exit(1)
+			}()
+			session.Stop()
+		case <-session.done:
+			return
+		}
+	}()
+	return session
 }
 
 func (p *ProfileSession) Stop() {
-	pprof.StopCPUProfile()
-	p.cpuFile.Close()
-
-	if p.memFilePath != "" {
-		memFile, err := os.Create(p.memFilePath)
-		if err != nil {
-			panic(err)
+	p.stopOnce.Do(func() {
+		if p.stopSignals != nil {
+			p.stopSignals()
 		}
-		if err := pprof.Lookup("allocs").WriteTo(memFile, 0); err != nil {
-			panic(err)
+		if p.done != nil {
+			close(p.done)
+		}
+		pprof.StopCPUProfile()
+		p.cpuFile.Close()
+
+		if p.memFilePath != "" {
+			memFile, err := os.Create(p.memFilePath)
+			if err != nil {
+				panic(err)
+			}
+			if err := pprof.Lookup("allocs").WriteTo(memFile, 0); err != nil {
+				panic(err)
+			}
+			memFile.Close()
+			fmt.Fprintf(p.logWriter, "Memory profile: %v\n", p.memFilePath)
 		}
-		memFile.Close()
-		fmt.Fprintf(p.logWriter, "Memory profile: %v\n", p.memFilePath)
-	}
 
-	fmt.Fprintf(p.logWriter, "CPU profile: %v\n", p.cpuFilePath)
+		fmt.Fprintf(p.logWriter, "CPU profile: %v\n", p.cpuFilePath)
+	})
 }
 
 // CPUProfiler manages on-demand CPU profiling.