From 352bbeb06cc3ac3fe6ae33c40f27b2ad72411b06 Mon Sep 17 00:00:00 2001
From: "Luis (LT) Carbonell" <lt.carbonell@hashicorp.com>
Date: Mon, 16 Sep 2024 11:00:08 -0400
Subject: [PATCH] VAULT-12732: Add Heap Profiling Option to Vault Server
 Command Line (#27033)

* Add pprof output option to CLI

* Add to docs

* Add changelog

* Update website/content/docs/commands/server.mdx

Co-authored-by: Sarah Chavis <62406755+schavis@users.noreply.github.com>

* some changes

* Update changelog/27033.txt

Co-authored-by: Josh Black <raskchanky@gmail.com>

---------

Co-authored-by: Sarah Chavis <62406755+schavis@users.noreply.github.com>
Co-authored-by: Josh Black <raskchanky@gmail.com>
---
 changelog/27033.txt                      |  3 ++
 command/server.go                        | 59 +++++++++++++++++++++++-
 website/content/docs/commands/server.mdx |  5 ++
 3 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 changelog/27033.txt

diff --git a/changelog/27033.txt b/changelog/27033.txt
new file mode 100644
index 0000000000..a06152dedc
--- /dev/null
+++ b/changelog/27033.txt
@@ -0,0 +1,3 @@
+```release-note:improvement
+command/server: Add support for dumping pprof files during startup using CLI option `pprof-dump-dir`
+```
diff --git a/command/server.go b/command/server.go
index 54e921bbab..5aba92cbd3 100644
--- a/command/server.go
+++ b/command/server.go
@@ -119,6 +119,7 @@ type ServerCommand struct {
 	flagConfigs            []string
 	flagRecovery           bool
 	flagExperiments        []string
+	flagCLIDump            string
 	flagDev                bool
 	flagDevTLS             bool
 	flagDevTLSCertDir      string
@@ -221,6 +222,13 @@ func (c *ServerCommand) Flags() *FlagSets {
 			"Valid experiments are: " + strings.Join(experiments.ValidExperiments(), ", "),
 	})
 
+	f.StringVar(&StringVar{
+		Name:       "pprof-dump-dir",
+		Target:     &c.flagCLIDump,
+		Completion: complete.PredictDirs("*"),
+		Usage:      "Directory where generated profiles are created. If left unset, files are not generated.",
+	})
+
 	f = set.NewFlagSet("Dev Options")
 
 	f.BoolVar(&BoolVar{
@@ -1593,6 +1601,11 @@ func (c *ServerCommand) Run(args []string) int {
 		coreShutdownDoneCh = core.ShutdownDone()
 	}
 
+	cliDumpCh := make(chan struct{})
+	if c.flagCLIDump != "" {
+		go func() { cliDumpCh <- struct{}{} }()
+	}
+
 	// Wait for shutdown
 	shutdownTriggered := false
 	retCode := 0
@@ -1707,7 +1720,6 @@ func (c *ServerCommand) Run(args []string) int {
 
 			// Notify systemd that the server has completed reloading config
 			c.notifySystemd(systemd.SdNotifyReady)
-
 		case <-c.SigUSR2Ch:
 			logWriter := c.logger.StandardWriter(&hclog.StandardLoggerOptions{})
 			pprof.Lookup("goroutine").WriteTo(logWriter, 2)
@@ -1759,6 +1771,51 @@ func (c *ServerCommand) Run(args []string) int {
 			}
 
 			c.logger.Info(fmt.Sprintf("Wrote pprof files to: %s", pprofPath))
+		case <-cliDumpCh:
+			path := c.flagCLIDump
+
+			if _, err := os.Stat(path); err != nil && !errors.Is(err, os.ErrNotExist) {
+				c.logger.Error("Checking cli dump path failed", "error", err)
+				continue
+			}
+
+			pprofPath := filepath.Join(path, "vault-pprof")
+			err := os.MkdirAll(pprofPath, os.ModePerm)
+			if err != nil {
+				c.logger.Error("Could not create temporary directory for pprof", "error", err)
+				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 startup pprof files to: %s", pprofPath))
 		}
 	}
 	// Notify systemd that the server is shutting down
diff --git a/website/content/docs/commands/server.mdx b/website/content/docs/commands/server.mdx
index ce0b987534..de61d671ef 100644
--- a/website/content/docs/commands/server.mdx
+++ b/website/content/docs/commands/server.mdx
@@ -92,6 +92,11 @@ flags](/vault/docs/commands) included on all commands.
   `VAULT_EXPERIMENTS` environment variable as a comma-separated list, or via the
   [`experiments`](/vault/docs/configuration#experiments) config key.
 
+- `-pprof-dump-dir` `(string: "")` - Directory where the generated profiles are
+  created. Vault does not generate profiles when `pprof-dump-dir` is unset.
+  Use `pprof-dump-dir` temporarily during debugging sessions. Do not use
+  `pprof-dump-dir` in regular production processes.
+
 - `VAULT_ALLOW_PENDING_REMOVAL_MOUNTS` `(bool: false)` - (environment variable)
   Allow Vault to be started with builtin engines which have the `Pending Removal`
   deprecation state. This is a temporary stopgap in place in order to perform an