From e158d02a5bd9c20913895fc4113fc26ce802a191 Mon Sep 17 00:00:00 2001
From: Adrian Hesketh <adrianhesketh@hushmail.com>
Date: Sat, 29 Jan 2022 12:57:51 +0000
Subject: [PATCH] feat: add pprof flag to enable profiling at runtime of the
 LSP

---
 cmd/templ/lspcmd/main.go | 9 +++++++++
 cmd/templ/main.go        | 2 ++
 2 files changed, 11 insertions(+)

diff --git a/cmd/templ/lspcmd/main.go b/cmd/templ/lspcmd/main.go
index 11f77231c..8be51e3c3 100644
--- a/cmd/templ/lspcmd/main.go
+++ b/cmd/templ/lspcmd/main.go
@@ -3,6 +3,7 @@ package lspcmd
 import (
 	"context"
 	"log"
+	"net/http"
 	"os"
 	"os/signal"
 
@@ -10,12 +11,15 @@ import (
 	"github.com/sourcegraph/jsonrpc2"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
+
+	_ "net/http/pprof"
 )
 
 type Arguments struct {
 	Log           string
 	GoplsLog      string
 	GoplsRPCTrace bool
+	PPROF         bool
 }
 
 func Run(args Arguments) error {
@@ -27,6 +31,11 @@ func Run(args Arguments) error {
 		signal.Stop(signalChan)
 		cancel()
 	}()
+	if args.PPROF {
+		go func() {
+			http.ListenAndServe("localhost:9999", nil)
+		}()
+	}
 	go func() {
 		select {
 		case <-signalChan: // First signal, cancel context.
diff --git a/cmd/templ/main.go b/cmd/templ/main.go
index 84f758bd7..acaa18575 100644
--- a/cmd/templ/main.go
+++ b/cmd/templ/main.go
@@ -104,6 +104,7 @@ func lspCmd(args []string) {
 	goplsLog := cmd.String("goplsLog", "", "The file to log gopls output, or leave empty to disable logging.")
 	goplsRPCTrace := cmd.Bool("goplsRPCTrace", false, "Set gopls to log input and output messages.")
 	helpFlag := cmd.Bool("help", false, "Print help and exit.")
+	pprofFlag := cmd.Bool("pprof", false, "Enable pprof web server (default address is localhost:9999)")
 	err := cmd.Parse(args)
 	if err != nil || *helpFlag {
 		cmd.PrintDefaults()
@@ -113,6 +114,7 @@ func lspCmd(args []string) {
 		Log:           *log,
 		GoplsLog:      *goplsLog,
 		GoplsRPCTrace: *goplsRPCTrace,
+		PPROF:         *pprofFlag,
 	})
 	if err != nil {
 		fmt.Println(err.Error())