From 536f00a5e5e420acea4a8f09c168d736abf262ac Mon Sep 17 00:00:00 2001
From: Dane Everitt <dane@daneeveritt.com>
Date: Sun, 8 Nov 2020 11:59:04 -0800
Subject: [PATCH] Add support for both memory and cpu profiling

---
 .gitignore  |  2 ++
 cmd/root.go | 30 +++++++++++++++++++++---------
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/.gitignore b/.gitignore
index 15dbd0e..62d62dd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -47,3 +47,5 @@ test_*/
 debug
 data/.states.json
 .DS_Store
+*.pprof
+*.pdf
\ No newline at end of file
diff --git a/cmd/root.go b/cmd/root.go
index 706906e..029b1cd 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -30,12 +30,15 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var configPath = config.DefaultLocation
-var debug = false
-var shouldRunProfiler = false
-var useAutomaticTls = false
-var tlsHostname = ""
-var showVersion = false
+var (
+	configPath      = config.DefaultLocation
+	debug           = false
+	runCPUProfile   = false
+	runMemProfile   = false
+	useAutomaticTls = false
+	tlsHostname     = ""
+	showVersion     = false
+)
 
 var root = &cobra.Command{
 	Use:   "wings",
@@ -54,7 +57,8 @@ func init() {
 	root.PersistentFlags().BoolVar(&showVersion, "version", false, "show the version and exit")
 	root.PersistentFlags().StringVar(&configPath, "config", config.DefaultLocation, "set the location for the configuration file")
 	root.PersistentFlags().BoolVar(&debug, "debug", false, "pass in order to run wings in debug mode")
-	root.PersistentFlags().BoolVar(&shouldRunProfiler, "profile", false, "pass in order to profile wings")
+	root.PersistentFlags().BoolVar(&runCPUProfile, "profile-cpu", false, "pass in order to profile wings CPU usage")
+	root.PersistentFlags().BoolVar(&runMemProfile, "profile-mem", false, "pass in order to profile wings memory usage")
 	root.PersistentFlags().BoolVar(&useAutomaticTls, "auto-tls", false, "pass in order to have wings generate and manage it's own SSL certificates using Let's Encrypt")
 	root.PersistentFlags().StringVar(&tlsHostname, "tls-hostname", "", "required with --auto-tls, the FQDN for the generated SSL certificate")
 
@@ -89,8 +93,16 @@ func rootCmdRun(*cobra.Command, []string) {
 		os.Exit(0)
 	}
 
-	if shouldRunProfiler {
-		defer profile.Start().Stop()
+	if runCPUProfile {
+		defer profile.Start(func(p *profile.Profile) {
+			profile.CPUProfile(p)
+		}).Stop()
+	}
+
+	if runMemProfile {
+		defer profile.Start(func(p *profile.Profile) {
+			profile.MemProfile(p)
+		}).Stop()
 	}
 
 	// Only attempt configuration file relocation if a custom location has not