From 59a68269208d7af7f6deff12484bcc541e99b978 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Mon, 6 Jun 2022 16:38:33 +0300
Subject: [PATCH] feat: Add support for pprof in `coder agent` (#1985)

* feat: Allow USR1 signal to start pprof
---
 cli/agent.go         | 17 ++++++++++++++++-
 cli/agent_unix.go    | 38 ++++++++++++++++++++++++++++++++++++++
 cli/agent_windows.go | 12 ++++++++++++
 3 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 cli/agent_unix.go
 create mode 100644 cli/agent_windows.go

diff --git a/cli/agent.go b/cli/agent.go
index 238a1e8589..287f18f7fb 100644
--- a/cli/agent.go
+++ b/cli/agent.go
@@ -3,6 +3,7 @@ package cli
 import (
 	"context"
 	"net/http"
+	_ "net/http/pprof" //nolint: gosec
 	"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,16 @@ func workspaceAgent() *cobra.Command {
 			logger := slog.Make(sloghuman.Sink(cmd.ErrOrStderr()), sloghuman.Sink(logWriter)).Leveled(slog.LevelDebug)
 			client := codersdk.New(coderURL)
 
+			if pprofEnabled {
+				srvClose := serveHandler(cmd.Context(), logger, nil, pprofAddress, "pprof")
+				defer srvClose()
+			} else {
+				// If pprof wasn't enabled at startup, allow a
+				// `kill -USR1 $agent_pid` to start it (on Unix).
+				srvClose := agentStartPPROFOnUSR1(cmd.Context(), logger, pprofAddress)
+				defer srvClose()
+			}
+
 			// exchangeToken returns a session token.
 			// This is abstracted to allow for the same looping condition
 			// regardless of instance identity auth type.
@@ -139,5 +152,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
 }
diff --git a/cli/agent_unix.go b/cli/agent_unix.go
new file mode 100644
index 0000000000..1a7dd3d854
--- /dev/null
+++ b/cli/agent_unix.go
@@ -0,0 +1,38 @@
+//go:build !windows
+
+package cli
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"cdr.dev/slog"
+)
+
+func agentStartPPROFOnUSR1(ctx context.Context, logger slog.Logger, pprofAddress string) (srvClose func()) {
+	ctx, cancel := context.WithCancel(ctx)
+
+	usr1 := make(chan os.Signal, 1)
+	signal.Notify(usr1, syscall.SIGUSR1)
+	go func() {
+		defer close(usr1)
+		defer signal.Stop(usr1)
+
+		select {
+		case <-usr1:
+			signal.Stop(usr1)
+			srvClose := serveHandler(ctx, logger, nil, pprofAddress, "pprof")
+			defer srvClose()
+		case <-ctx.Done():
+			return
+		}
+		<-ctx.Done() // Prevent defer close until done.
+	}()
+
+	return func() {
+		cancel()
+		<-usr1 // Wait until usr1 is closed, ensures srvClose was run.
+	}
+}
diff --git a/cli/agent_windows.go b/cli/agent_windows.go
new file mode 100644
index 0000000000..6635e350cf
--- /dev/null
+++ b/cli/agent_windows.go
@@ -0,0 +1,12 @@
+package cli
+
+import (
+	"context"
+
+	"cdr.dev/slog"
+)
+
+// agentStartPPROFOnUSR1 is no-op on Windows (no SIGUSR1 signal).
+func agentStartPPROFOnUSR1(ctx context.Context, logger slog.Logger, pprofAddress string) (srvClose func()) {
+	return func() {}
+}