Add SIGUSR2 pprof to agent and proxy (#27510)

* Add SIGUSR2 pprof to agent

* changelog

* Update command/agent.go

Co-authored-by: Violet Hynes <violet.hynes@hashicorp.com>

* Update command/agent.go

Co-authored-by: Violet Hynes <violet.hynes@hashicorp.com>

* Add to proxy, update tests

* Fix path

* Changelog

* dry

* choose one error style

---------

Co-authored-by: Violet Hynes <violet.hynes@hashicorp.com>

changelog/27510.txt

@@ -0,0 +1,6 @@
+```release-note:improvement
+agent: Add the ability to dump pprof to the filesystem using SIGUSR2
+```
+```release-note:improvement
+proxy: Add the ability to dump pprof to the filesystem using SIGUSR2
+```

command/agent.go

@@ -13,6 +13,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"path/filepath"
 	"sort"
 	"strings"
 	"sync"
@@ -74,6 +75,7 @@ type AgentCommand struct {
 
 	ShutdownCh chan struct{}
 	SighupCh   chan struct{}
+	SigUSR2Ch  chan struct{}
 
 	tlsReloadFuncsLock sync.RWMutex
 	tlsReloadFuncs     []reloadutil.ReloadFunc
@@ -758,6 +760,16 @@ func (c *AgentCommand) Run(args []string) int {
 				case c.reloadedCh <- struct{}{}:
 				default:
 				}
+			case <-c.SigUSR2Ch:
+				pprofPath := filepath.Join(os.TempDir(), "vault-agent-pprof")
+				cpuProfileDuration := time.Second * 1
+				err := WritePprofToFile(pprofPath, cpuProfileDuration)
+				if err != nil {
+					c.logger.Error(err.Error())
+					continue
+				}
+
+				c.logger.Info(fmt.Sprintf("Wrote pprof files to: %s", pprofPath))
 			case <-ctx.Done():
 				return nil
 			}

command/agent_test.go

@@ -91,6 +91,7 @@ func testAgentCommand(tb testing.TB, logger hclog.Logger) (*cli.MockUi, *AgentCo
 		},
 		ShutdownCh: MakeShutdownCh(),
 		SighupCh:   MakeSighupCh(),
+		SigUSR2Ch:  MakeSigUSR2Ch(),
 		logger:     logger,
 		startedCh:  make(chan struct{}, 5),
 		reloadedCh: make(chan struct{}, 5),

command/commands.go

@@ -192,6 +192,7 @@ func initCommands(ui, serverCmdUi cli.Ui, runOpts *RunOptions) map[string]cli.Co
 				},
 				ShutdownCh: MakeShutdownCh(),
 				SighupCh:   MakeSighupCh(),
+				SigUSR2Ch:  MakeSigUSR2Ch(),
 			}, nil
 		},
 		"agent generate-config": func() (cli.Command, error) {
@@ -576,6 +577,7 @@ func initCommands(ui, serverCmdUi cli.Ui, runOpts *RunOptions) map[string]cli.Co
 				},
 				ShutdownCh: MakeShutdownCh(),
 				SighupCh:   MakeSighupCh(),
+				SigUSR2Ch:  MakeSigUSR2Ch(),
 			}, nil
 		},
 		"policy": func() (cli.Command, error) {

command/proxy.go

@@ -12,6 +12,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"path/filepath"
 	"sort"
 	"strings"
 	"sync"
@@ -71,6 +72,7 @@ type ProxyCommand struct {
 
 	ShutdownCh chan struct{}
 	SighupCh   chan struct{}
+	SigUSR2Ch  chan struct{}
 
 	tlsReloadFuncsLock sync.RWMutex
 	tlsReloadFuncs     []reloadutil.ReloadFunc
@@ -715,6 +717,16 @@ func (c *ProxyCommand) Run(args []string) int {
 				case c.reloadedCh <- struct{}{}:
 				default:
 				}
+			case <-c.SigUSR2Ch:
+				pprofPath := filepath.Join(os.TempDir(), "vault-proxy-pprof")
+				cpuProfileDuration := time.Second * 1
+				err := WritePprofToFile(pprofPath, cpuProfileDuration)
+				if err != nil {
+					c.logger.Error(err.Error())
+					continue
+				}
+
+				c.logger.Info(fmt.Sprintf("Wrote pprof files to: %s", pprofPath))
 			case <-ctx.Done():
 				return nil
 			}

command/proxy_test.go

@@ -44,6 +44,7 @@ func testProxyCommand(tb testing.TB, logger hclog.Logger) (*cli.MockUi, *ProxyCo
 		},
 		ShutdownCh: MakeShutdownCh(),
 		SighupCh:   MakeSighupCh(),
+		SigUSR2Ch:  MakeSigUSR2Ch(),
 		logger:     logger,
 		startedCh:  make(chan struct{}, 5),
 		reloadedCh: make(chan struct{}, 5),

command/server.go

@@ -1782,41 +1782,13 @@ func (c *ServerCommand) Run(args []string) int {
 			// into a state where it cannot process requests so we can get pprof outputs
 			// via SIGUSR2.
 			pprofPath := filepath.Join(os.TempDir(), "vault-pprof")
-			err := os.MkdirAll(pprofPath, os.ModePerm)
+			cpuProfileDuration := time.Second * 1
+			err := WritePprofToFile(pprofPath, cpuProfileDuration)
 			if err != nil {
-				c.logger.Error("Could not create temporary directory for pprof", "error", err)
+				c.logger.Error(err.Error())
 				continue
 			}
 
-			dumps := []string{"goroutine", "heap", "allocs", "threadcreate", "profile"}
-			for _, dump := range dumps {
-				pFile, err := os.Create(filepath.Join(pprofPath, dump))
-				if err != nil {
-					c.logger.Error("error creating pprof file", "name", dump, "error", err)
-					break
-				}
-
-				if dump != "profile" {
-					err = pprof.Lookup(dump).WriteTo(pFile, 0)
-					if err != nil {
-						c.logger.Error("error generating pprof data", "name", dump, "error", err)
-						pFile.Close()
-						break
-					}
-				} else {
-					// CPU profiles need to run for a duration so we're going to run it
-					// just for one second to avoid blocking here.
-					if err := pprof.StartCPUProfile(pFile); err != nil {
-						c.logger.Error("could not start CPU profile: ", err)
-						pFile.Close()
-						break
-					}
-					time.Sleep(time.Second * 1)
-					pprof.StopCPUProfile()
-				}
-				pFile.Close()
-			}
-
 			c.logger.Info(fmt.Sprintf("Wrote pprof files to: %s", pprofPath))
 		}
 	}

command/util.go

@@ -8,6 +8,8 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"path/filepath"
+	"runtime/pprof"
 	"testing"
 	"time"
 
@@ -201,3 +203,40 @@ func (r *recordingRoundTripper) RoundTrip(req *http.Request) (*http.Response, er
 		StatusCode: 200,
 	}, nil
 }
+
+// WritePprofToFile will create a temporary directory at the specified path
+// and generate pprof files at that location. CPU requires polling over a
+// duration. For most situations 1 second is enough.
+func WritePprofToFile(path string, cpuProfileDuration time.Duration) error {
+	err := os.MkdirAll(path, os.ModePerm)
+	if err != nil {
+		return fmt.Errorf("could not create temporary directory for pprof: %v", err)
+	}
+
+	dumps := []string{"goroutine", "heap", "allocs", "threadcreate", "profile"}
+	for _, dump := range dumps {
+		pFile, err := os.Create(filepath.Join(path, dump))
+		if err != nil {
+			return fmt.Errorf("error creating pprof file %s: %v", dump, err)
+		}
+
+		if dump != "profile" {
+			err = pprof.Lookup(dump).WriteTo(pFile, 0)
+			if err != nil {
+				pFile.Close()
+				return fmt.Errorf("error generating pprof data for %s: %v", dump, err)
+			}
+		} else {
+			// CPU profiles need to run for a duration so we're going to run it
+			// just for one second to avoid blocking here.
+			if err := pprof.StartCPUProfile(pFile); err != nil {
+				pFile.Close()
+				return fmt.Errorf("could not start CPU profile: %v", err)
+			}
+			time.Sleep(cpuProfileDuration)
+			pprof.StopCPUProfile()
+		}
+		pFile.Close()
+	}
+	return nil
+}