feat: Add support for pprof in coder agent

mafredri · mafredri · commit 490a226b3ddc · 2022-06-02T18:28:07.000+03:00
diff --git a/cli/agent.go b/cli/agent.go
@@ -3,6 +3,7 @@ package cli
 import (
 	"context"
 	"net/http"
+	_ "net/http/pprof"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -25,7 +26,9 @@ import (
 
 func workspaceAgent() *cobra.Command {
 	var (
-		auth string
+		auth         string
+		pprofEnabled bool
+		pprofAddress string
 	)
 	cmd := &cobra.Command{
 		Use: "agent",
@@ -49,6 +52,11 @@ func workspaceAgent() *cobra.Command {
 			logger := slog.Make(sloghuman.Sink(cmd.ErrOrStderr()), sloghuman.Sink(logWriter)).Leveled(slog.LevelDebug)
 			client := codersdk.New(coderURL)
 
+			if pprofEnabled {
+				//nolint:revive
+				defer serveHandler(cmd.Context(), logger, nil, pprofAddress, "pprof")()
+			}
+
 			// exchangeToken returns a session token.
 			// This is abstracted to allow for the same looping condition
 			// regardless of instance identity auth type.
@@ -139,5 +147,7 @@ func workspaceAgent() *cobra.Command {
 	}
 
 	cliflag.StringVarP(cmd.Flags(), &auth, "auth", "", "CODER_AGENT_AUTH", "token", "Specify the authentication type to use for the agent")
+	cliflag.BoolVarP(cmd.Flags(), &pprofEnabled, "pprof-enable", "", "CODER_AGENT_PPROF_ENABLE", false, "Enable serving pprof metrics on the address defined by --pprof-address.")
+	cliflag.StringVarP(cmd.Flags(), &pprofAddress, "pprof-address", "", "CODER_AGENT_PPROF_ADDRESS", "127.0.0.1:6060", "The address to serve pprof.")
 	return cmd
 }
