Add memory profiling for custom builds (#5584)

Makefile

@@ -37,6 +37,16 @@ dev-ui: prep
 dev-dynamic: prep
 	@CGO_ENABLED=1 BUILD_TAGS='$(BUILD_TAGS)' VAULT_DEV_BUILD=1 sh -c "'$(CURDIR)/scripts/build.sh'"
 
+# *-mem variants will enable memory profiling which will write snapshots of heap usage
+# to $TMP/vaultprof every 5 minutes. These can be analyzed using `$ go tool pprof <profile_file>`.
+# Note that any build can have profiling added via: `$ BUILD_TAGS=memprofiler make ...`
+dev-mem: BUILD_TAGS+=memprofiler
+dev-mem: dev
+dev-ui-mem: BUILD_TAGS+=memprofiler
+dev-ui-mem: dev-ui
+dev-dynamic-mem: BUILD_TAGS+=memprofiler
+dev-dynamic-mem: dev-dynamic
+
 testtravis: BUILD_TAGS+=travis
 testtravis: test
 

command/server.go

@@ -20,12 +20,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/mitchellh/cli"
-	testing "github.com/mitchellh/go-testing-interface"
-	"github.com/posener/complete"
-
-	"google.golang.org/grpc/grpclog"
-
 	"github.com/armon/go-metrics"
 	"github.com/armon/go-metrics/circonus"
 	"github.com/armon/go-metrics/datadog"
@@ -49,11 +43,17 @@ import (
 	"github.com/hashicorp/vault/vault"
 	vaultseal "github.com/hashicorp/vault/vault/seal"
 	"github.com/hashicorp/vault/version"
+	"github.com/mitchellh/cli"
+	testing "github.com/mitchellh/go-testing-interface"
+	"github.com/posener/complete"
+	"google.golang.org/grpc/grpclog"
 )
 
 var _ cli.Command = (*ServerCommand)(nil)
 var _ cli.CommandAutocomplete = (*ServerCommand)(nil)
 
+var memProfilerEnabled = false
+
 const storageMigrationLock = "core/migration"
 
 type ServerCommand struct {
@@ -431,6 +431,10 @@ func (c *ServerCommand) Run(args []string) int {
 		log:    os.Getenv("VAULT_GRPC_LOGGING") != "",
 	})
 
+	if memProfilerEnabled {
+		c.startMemProfiler()
+	}
+
 	// Ensure that a backend is provided
 	if config.Storage == nil {
 		c.UI.Output("A storage backend must be specified")

command/server_noprofile.go

@@ -0,0 +1,6 @@
+// +build !memprofiler
+
+package command
+
+func (c *ServerCommand) startMemProfiler() {
+}

command/server_profile.go

@@ -0,0 +1,40 @@
+// +build memprofiler
+
+package command
+
+import (
+	"os"
+	"path/filepath"
+	"runtime"
+	"runtime/pprof"
+	"time"
+)
+
+func init() {
+	memProfilerEnabled = true
+}
+
+func (c *ServerCommand) startMemProfiler() {
+	profileDir := filepath.Join(os.TempDir(), "vaultprof")
+	if err := os.MkdirAll(profileDir, 0700); err != nil {
+		c.logger.Debug("could not create profile directory", "error", err)
+		return
+	}
+
+	go func() {
+		for {
+			filename := filepath.Join(profileDir, time.Now().UTC().Format("20060102_150405")) + ".pprof"
+			f, err := os.Create(filename)
+			if err != nil {
+				c.logger.Debug("could not create memory profile", "error", err)
+			}
+			runtime.GC()
+			if err := pprof.WriteHeapProfile(f); err != nil {
+				c.logger.Debug("could not write memory profile", "error", err)
+			}
+			f.Close()
+			c.logger.Debug("wrote memory profile", "filename", filename)
+			time.Sleep(5 * time.Minute)
+		}
+	}()
+}