1 files changed, 60 insertions, 1 deletions

diff --git a/mattermost.go b/mattermost.go
index 8d4880504..7b9f6c0c1 100644
--- a/mattermost.go
+++ b/mattermost.go
@@ -14,6 +14,7 @@ import (
 	"os/exec"
 	"os/signal"
 	"runtime"
+	"runtime/pprof"
 	"strconv"
 	"strings"
 	"syscall"
@@ -83,6 +84,10 @@ var flagChannelHeader string
 var flagChannelPurpose string
 var flagUserSetInactive bool
 var flagImportArchive string
+var flagCpuProfile bool
+var flagMemProfile bool
+var flagBlockProfile bool
+var flagHttpProfiler bool
 
 func doLoadConfig(filename string) (err string) {
 	defer func() {
@@ -122,7 +127,26 @@ func main() {
 
 	cmdUpdateDb30()
 
-	api.NewServer()
+	if flagCpuProfile {
+		f, err := os.Create(utils.GetLogFileLocation(utils.Cfg.LogSettings.FileLocation) + ".cpu.prof")
+		if err != nil {
+			l4g.Error("Error creating cpu profile log: " + err.Error())
+		}
+
+		l4g.Info("CPU Profiler is logging to " + utils.GetLogFileLocation(utils.Cfg.LogSettings.FileLocation) + ".cpu.prof")
+		pprof.StartCPUProfile(f)
+	}
+
+	if flagBlockProfile {
+		l4g.Info("Block Profiler is logging to " + utils.GetLogFileLocation(utils.Cfg.LogSettings.FileLocation) + ".blk.prof")
+		runtime.SetBlockProfileRate(1)
+	}
+
+	if flagMemProfile {
+		l4g.Info("Memory Profiler is logging to " + utils.GetLogFileLocation(utils.Cfg.LogSettings.FileLocation) + ".mem.prof")
+	}
+
+	api.NewServer(flagHttpProfiler)
 	api.InitApi()
 	web.InitWeb()
 
@@ -169,6 +193,37 @@ func main() {
 		}
 
 		api.StopServer()
+
+		if flagCpuProfile {
+			l4g.Info("Closing CPU Profiler")
+			pprof.StopCPUProfile()
+		}
+
+		if flagBlockProfile {
+			f, err := os.Create(utils.GetLogFileLocation(utils.Cfg.LogSettings.FileLocation) + ".blk.prof")
+			if err != nil {
+				l4g.Error("Error creating block profile log: " + err.Error())
+			}
+
+			l4g.Info("Writing Block Profiler to: " + utils.GetLogFileLocation(utils.Cfg.LogSettings.FileLocation) + ".blk.prof")
+			pprof.Lookup("block").WriteTo(f, 0)
+			f.Close()
+			runtime.SetBlockProfileRate(0)
+		}
+
+		if flagMemProfile {
+			f, err := os.Create(utils.GetLogFileLocation(utils.Cfg.LogSettings.FileLocation) + ".mem.prof")
+			if err != nil {
+				l4g.Error("Error creating memory profile file: " + err.Error())
+			}
+
+			l4g.Info("Writing Memory Profiler to: " + utils.GetLogFileLocation(utils.Cfg.LogSettings.FileLocation) + ".mem.prof")
+			runtime.GC()
+			if err := pprof.WriteHeapProfile(f); err != nil {
+				l4g.Error("Error creating memory profile: " + err.Error())
+			}
+			f.Close()
+		}
 	}
 }
 
@@ -380,6 +435,10 @@ func parseCmds() {
 	flag.BoolVar(&flagCmdActivateUser, "activate_user", false, "")
 	flag.BoolVar(&flagCmdSlackImport, "slack_import", false, "")
 	flag.BoolVar(&flagUserSetInactive, "inactive", false, "")
+	flag.BoolVar(&flagCpuProfile, "cpuprofile", false, "")
+	flag.BoolVar(&flagMemProfile, "memprofile", false, "")
+	flag.BoolVar(&flagBlockProfile, "blkprofile", false, "")
+	flag.BoolVar(&flagHttpProfiler, "httpprofiler", false, "")
 
 	flag.Parse()