add -pprofDir profiling to LSP (#614)

Co-authored-by: Jake Bailey <[email protected]>

Author: mjames-c

_extension/package.json

@@ -30,6 +30,10 @@
                         ],
                         "default": "verbose",
                         "description": "Trace TypeScript Go server communication."
+                    },
+                    "typescript-go.pprofDir": {
+                        "type": "string",
+                        "description": "Directory to write pprof profiles to."
                     }
                 }
             }

_extension/src/extension.ts

@@ -26,15 +26,21 @@ export function activate(context: vscode.ExtensionContext) {
 
     output.appendLine(`Resolved to ${exe}`);
 
+    const config = vscode.workspace.getConfiguration("typescript-go");
+
+    // Get pprofDir
+    const pprofDir = config.get<string>("pprofDir");
+    const pprofArgs = pprofDir ? ["-pprofDir", pprofDir] : [];
+
     const serverOptions: ServerOptions = {
         run: {
             command: exe,
-            args: ["lsp"],
+            args: ["lsp", ...pprofArgs],
             transport: TransportKind.stdio,
         },
         debug: {
             command: exe,
-            args: ["lsp"],
+            args: ["lsp", ...pprofArgs],
             transport: TransportKind.stdio,
         },
     };

cmd/tsgo/lsp.go

@@ -10,12 +10,14 @@ import (
 	"github.com/microsoft/typescript-go/internal/bundled"
 	"github.com/microsoft/typescript-go/internal/core"
 	"github.com/microsoft/typescript-go/internal/lsp"
+	"github.com/microsoft/typescript-go/internal/pprof"
 	"github.com/microsoft/typescript-go/internal/vfs/osvfs"
 )
 
 func runLSP(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.")
 	pipe := flag.String("pipe", "", "use named pipe for communication")
 	_ = pipe
 	socket := flag.String("socket", "", "use socket for communication")
@@ -29,6 +31,12 @@ func runLSP(args []string) int {
 		return 1
 	}
 
+	if *pprofDir != "" {
+		fmt.Fprintf(os.Stderr, "pprof profiles will be written to: %v\n", *pprofDir)
+		profileSession := pprof.BeginProfiling(*pprofDir, os.Stderr)
+		defer profileSession.Stop()
+	}
+
 	fs := bundled.WrapFS(osvfs.FS())
 	defaultLibraryPath := bundled.LibPath()
 

cmd/tsgo/main.go

@@ -6,9 +6,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
-	"path/filepath"
 	"runtime"
-	"runtime/pprof"
 	"slices"
 	"strconv"
 	"strings"
@@ -21,6 +19,7 @@ import (
 	"github.com/microsoft/typescript-go/internal/core"
 	"github.com/microsoft/typescript-go/internal/diagnosticwriter"
 	"github.com/microsoft/typescript-go/internal/execute"
+	"github.com/microsoft/typescript-go/internal/pprof"
 	"github.com/microsoft/typescript-go/internal/scanner"
 	"github.com/microsoft/typescript-go/internal/tspath"
 	"github.com/microsoft/typescript-go/internal/vfs/osvfs"
@@ -135,8 +134,8 @@ func main() {
 	opts := parseArgs()
 
 	if opts.devel.pprofDir != "" {
-		profileSession := beginProfiling(opts.devel.pprofDir)
-		defer profileSession.stop()
+		profileSession := pprof.BeginProfiling(opts.devel.pprofDir, os.Stdout)
+		defer profileSession.Stop()
 	}
 
 	startTime := time.Now()
@@ -335,54 +334,3 @@ func printDiagnostics(diagnostics []*ast.Diagnostic, host ts.CompilerHost, compi
 		}
 	}
 }
-
-type profileSession struct {
-	cpuFilePath string
-	memFilePath string
-	cpuFile     *os.File
-	memFile     *os.File
-}
-
-func beginProfiling(profileDir string) *profileSession {
-	if err := os.MkdirAll(profileDir, 0o755); err != nil {
-		panic(err)
-	}
-
-	pid := os.Getpid()
-
-	cpuProfilePath := filepath.Join(profileDir, fmt.Sprintf("%d-cpuprofile.pb.gz", pid))
-	memProfilePath := filepath.Join(profileDir, fmt.Sprintf("%d-memprofile.pb.gz", pid))
-	cpuFile, err := os.Create(cpuProfilePath)
-	if err != nil {
-		panic(err)
-	}
-	memFile, err := os.Create(memProfilePath)
-	if err != nil {
-		panic(err)
-	}
-
-	if err := pprof.StartCPUProfile(cpuFile); err != nil {
-		panic(err)
-	}
-
-	return &profileSession{
-		cpuFilePath: cpuProfilePath,
-		memFilePath: memProfilePath,
-		cpuFile:     cpuFile,
-		memFile:     memFile,
-	}
-}
-
-func (p *profileSession) stop() {
-	pprof.StopCPUProfile()
-	err := pprof.Lookup("allocs").WriteTo(p.memFile, 0)
-	if err != nil {
-		panic(err)
-	}
-
-	p.cpuFile.Close()
-	p.memFile.Close()
-
-	fmt.Printf("CPU profile: %v\n", p.cpuFilePath)
-	fmt.Printf("Memory profile: %v\n", p.memFilePath)
-}

internal/pprof/pprof.go

@@ -0,0 +1,63 @@
+package pprof
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"runtime/pprof"
+)
+
+type profileSession struct {
+	cpuFilePath string
+	memFilePath string
+	cpuFile     *os.File
+	memFile     *os.File
+	logWriter   io.Writer
+}
+
+// BeginProfiling starts CPU and memory profiling, writing the profiles to the specified directory.
+func BeginProfiling(profileDir string, logWriter io.Writer) *profileSession {
+	if err := os.MkdirAll(profileDir, 0o755); err != nil {
+		panic(err)
+	}
+
+	pid := os.Getpid()
+
+	cpuProfilePath := filepath.Join(profileDir, fmt.Sprintf("%d-cpuprofile.pb.gz", pid))
+	memProfilePath := filepath.Join(profileDir, fmt.Sprintf("%d-memprofile.pb.gz", pid))
+	cpuFile, err := os.Create(cpuProfilePath)
+	if err != nil {
+		panic(err)
+	}
+	memFile, err := os.Create(memProfilePath)
+	if err != nil {
+		panic(err)
+	}
+
+	if err := pprof.StartCPUProfile(cpuFile); err != nil {
+		panic(err)
+	}
+
+	return &profileSession{
+		cpuFilePath: cpuProfilePath,
+		memFilePath: memProfilePath,
+		cpuFile:     cpuFile,
+		memFile:     memFile,
+		logWriter:   logWriter,
+	}
+}
+
+func (p *profileSession) Stop() {
+	pprof.StopCPUProfile()
+	err := pprof.Lookup("allocs").WriteTo(p.memFile, 0)
+	if err != nil {
+		panic(err)
+	}
+
+	p.cpuFile.Close()
+	p.memFile.Close()
+
+	fmt.Fprintf(p.logWriter, "CPU profile: %v\n", p.cpuFilePath)
+	fmt.Fprintf(p.logWriter, "Memory profile: %v\n", p.memFilePath)
+}