feat: Add support for pprof in coder agent (#1985)

mafredri · web-flow · commit 59a68269208d · 2022-06-06T16:38:33.000+03:00
* feat: Allow USR1 signal to start pprof
diff --git 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
@@ -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
@@ -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() {}
+}
