From 16fe21138ffc7ac1d593c115d72dcc3314ca0e17 Mon Sep 17 00:00:00 2001
From: WithoutPants <53250216+WithoutPants@users.noreply.github.com>
Date: Sun, 16 May 2021 17:21:11 +1000
Subject: [PATCH] CPU profiling (#1371)

* Add cpuprofile flag
* Add notes to readme
---
 README.md                    | 12 ++++++++++++
 main.go                      | 14 +++++++++++++-
 pkg/manager/config/config.go | 10 +++++++++-
 pkg/manager/config/init.go   | 10 +++++++---
 pkg/manager/manager.go       | 18 ++++++++++++++++++
 5 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index ce1322be4..fd80fc78b 100644
--- a/README.md
+++ b/README.md
@@ -145,3 +145,15 @@ where the app can be cross-compiled.  This process is kicked off by CI via the `
 command to open a bash shell to the container to poke around:
 
 `docker run --rm --mount type=bind,source="$(pwd)",target=/stash -w /stash -i -t stashappdev/compiler:latest /bin/bash`
+
+## Profiling
+
+Stash can be profiled using the `--cpuprofile <output profile filename>` command line flag.
+
+The resulting file can then be used with pprof as follows: 
+
+`go tool pprof <path to binary> <path to profile filename>`
+
+With `graphviz` installed and in the path, a call graph can be generated with:
+
+`go tool pprof -svg <path to binary> <path to profile filename> > <output svg file>`
diff --git a/main.go b/main.go
index bc2a83c6f..e08271a7f 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,11 @@
 package main
 
 import (
+	"os"
+	"os/signal"
+	"runtime/pprof"
+	"syscall"
+
 	"github.com/stashapp/stash/pkg/api"
 	"github.com/stashapp/stash/pkg/manager"
 
@@ -12,9 +17,16 @@ import (
 func main() {
 	manager.Initialize()
 	api.Start()
+
+	// stop any profiling at exit
+	defer pprof.StopCPUProfile()
 	blockForever()
 }
 
 func blockForever() {
-	select {}
+	// handle signals
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
+
+	<-signals
 }
diff --git a/pkg/manager/config/config.go b/pkg/manager/config/config.go
index 440a1fded..6f3267f0b 100644
--- a/pkg/manager/config/config.go
+++ b/pkg/manager/config/config.go
@@ -142,7 +142,8 @@ func (e MissingConfigError) Error() string {
 }
 
 type Instance struct {
-	isNewSystem bool
+	cpuProfilePath string
+	isNewSystem    bool
 }
 
 var instance *Instance
@@ -162,6 +163,13 @@ func (i *Instance) SetConfigFile(fn string) {
 	viper.SetConfigFile(fn)
 }
 
+// GetCPUProfilePath returns the path to the CPU profile file to output
+// profiling info to. This is set only via a commandline flag. Returns an
+// empty string if not set.
+func (i *Instance) GetCPUProfilePath() string {
+	return i.cpuProfilePath
+}
+
 func (i *Instance) Set(key string, value interface{}) {
 	viper.Set(key, value)
 }
diff --git a/pkg/manager/config/init.go b/pkg/manager/config/init.go
index 8b6c2ff0e..49f168b66 100644
--- a/pkg/manager/config/init.go
+++ b/pkg/manager/config/init.go
@@ -17,18 +17,20 @@ var once sync.Once
 
 type flagStruct struct {
 	configFilePath string
+	cpuProfilePath string
 }
 
 func Initialize() (*Instance, error) {
 	var err error
 	once.Do(func() {
-		instance = &Instance{}
-
 		flags := initFlags()
+		instance = &Instance{
+			cpuProfilePath: flags.cpuProfilePath,
+		}
+
 		if err = initConfig(flags); err != nil {
 			return
 		}
-
 		initEnvs()
 
 		if instance.isNewSystem {
@@ -46,6 +48,7 @@ func Initialize() (*Instance, error) {
 }
 
 func initConfig(flags flagStruct) error {
+
 	// The config file is called config.  Leave off the file extension.
 	viper.SetConfigName("config")
 
@@ -98,6 +101,7 @@ func initFlags() flagStruct {
 	pflag.IP("host", net.IPv4(0, 0, 0, 0), "ip address for the host")
 	pflag.Int("port", 9999, "port to serve from")
 	pflag.StringVarP(&flags.configFilePath, "config", "c", "", "config file to use")
+	pflag.StringVar(&flags.cpuProfilePath, "cpuprofile", "", "write cpu profile to file")
 
 	pflag.Parse()
 	if err := viper.BindPFlags(pflag.CommandLine); err != nil {
diff --git a/pkg/manager/manager.go b/pkg/manager/manager.go
index 4aa05dcb4..68bb2647d 100644
--- a/pkg/manager/manager.go
+++ b/pkg/manager/manager.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"runtime/pprof"
 	"sync"
 	"time"
 
@@ -55,6 +56,7 @@ func Initialize() *singleton {
 		}
 
 		initLog()
+		initProfiling(cfg.GetCPUProfilePath())
 
 		instance = &singleton{
 			Config:        cfg,
@@ -92,6 +94,22 @@ func Initialize() *singleton {
 	return instance
 }
 
+func initProfiling(cpuProfilePath string) {
+	if cpuProfilePath == "" {
+		return
+	}
+
+	f, err := os.Create(cpuProfilePath)
+	if err != nil {
+		logger.Fatalf("unable to create cpu profile file: %s", err.Error())
+	}
+
+	logger.Infof("profiling to %s", cpuProfilePath)
+
+	// StopCPUProfile is defer called in main
+	pprof.StartCPUProfile(f)
+}
+
 func initFFMPEG() {
 	configDirectory := paths.GetStashHomeDirectory()
 	ffmpegPath, ffprobePath := ffmpeg.GetPaths(configDirectory)