From 1fbb8723dc414dc9b84bbaf84224b6f5829d34de Mon Sep 17 00:00:00 2001
From: Chris Lu <chris.lu@gmail.com>
Date: Thu, 22 Jun 2017 01:33:58 -0700
Subject: [PATCH] add cpu and mem pprof

---
 weed/command/master.go                        | 15 ++-------
 weed/command/mount_std.go                     |  2 +-
 weed/command/server.go                        |  2 +-
 weed/command/volume.go                        |  7 +++-
 weed/util/pprof.go                            | 33 +++++++++++++++++++
 weed/{command => util}/signal_handling.go     |  2 +-
 .../signal_handling_notsupported.go           |  2 +-
 7 files changed, 46 insertions(+), 17 deletions(-)
 create mode 100644 weed/util/pprof.go
 rename weed/{command => util}/signal_handling.go (96%)
 rename weed/{command => util}/signal_handling_notsupported.go (75%)

diff --git a/weed/command/master.go b/weed/command/master.go
index 09089a4f9..8c5efbb9f 100644
--- a/weed/command/master.go
+++ b/weed/command/master.go
@@ -4,7 +4,6 @@ import (
 	"net/http"
 	"os"
 	"runtime"
-	"runtime/pprof"
 	"strconv"
 	"strings"
 	"time"
@@ -48,6 +47,7 @@ var (
 	masterWhiteListOption = cmdMaster.Flag.String("whiteList", "", "comma separated Ip addresses having write permission. No limit if empty.")
 	masterSecureKey       = cmdMaster.Flag.String("secure.secret", "", "secret to encrypt Json Web Token(JWT)")
 	masterCpuProfile      = cmdMaster.Flag.String("cpuprofile", "", "cpu profile output file")
+	masterMemProfile      = cmdMaster.Flag.String("memprofile", "", "memory profile output file")
 
 	masterWhiteList []string
 )
@@ -57,17 +57,8 @@ func runMaster(cmd *Command, args []string) bool {
 		*mMaxCpu = runtime.NumCPU()
 	}
 	runtime.GOMAXPROCS(*mMaxCpu)
-	if *masterCpuProfile != "" {
-		f, err := os.Create(*masterCpuProfile)
-		if err != nil {
-			glog.Fatal(err)
-		}
-		pprof.StartCPUProfile(f)
-		defer pprof.StopCPUProfile()
-		OnInterrupt(func() {
-			pprof.StopCPUProfile()
-		})
-	}
+	util.SetupProfiling(*masterCpuProfile, *masterMemProfile)
+
 	if err := util.TestFolderWritable(*metaFolder); err != nil {
 		glog.Fatalf("Check Meta Folder (-mdir) Writable %s : %s", *metaFolder, err)
 	}
diff --git a/weed/command/mount_std.go b/weed/command/mount_std.go
index a87068675..0c36e488f 100644
--- a/weed/command/mount_std.go
+++ b/weed/command/mount_std.go
@@ -31,7 +31,7 @@ func runMount(cmd *Command, args []string) bool {
 		return false
 	}
 
-	OnInterrupt(func() {
+	util.OnInterrupt(func() {
 		fuse.Unmount(*mountOptions.dir)
 		c.Close()
 	})
diff --git a/weed/command/server.go b/weed/command/server.go
index c8878e9eb..cd8ddf66a 100644
--- a/weed/command/server.go
+++ b/weed/command/server.go
@@ -290,7 +290,7 @@ func runServer(cmd *Command, args []string) bool {
 		}()
 	}
 
-	OnInterrupt(func() {
+	util.OnInterrupt(func() {
 		volumeServer.Shutdown()
 		pprof.StopCPUProfile()
 	})
diff --git a/weed/command/volume.go b/weed/command/volume.go
index ad9803974..a54ffd1fd 100644
--- a/weed/command/volume.go
+++ b/weed/command/volume.go
@@ -36,6 +36,8 @@ type VolumeServerOptions struct {
 	indexType             *string
 	fixJpgOrientation     *bool
 	readRedirect          *bool
+	cpuProfile            *string
+	memProfile            *string
 }
 
 func init() {
@@ -54,6 +56,8 @@ func init() {
 	v.indexType = cmdVolume.Flag.String("index", "memory", "Choose [memory|leveldb|boltdb|btree] mode for memory~performance balance.")
 	v.fixJpgOrientation = cmdVolume.Flag.Bool("images.fix.orientation", true, "Adjust jpg orientation when uploading.")
 	v.readRedirect = cmdVolume.Flag.Bool("read.redirect", true, "Redirect moved or non-local volumes.")
+	v.cpuProfile = cmdVolume.Flag.String("cpuprofile", "", "cpu profile output file")
+	v.memProfile = cmdVolume.Flag.String("memprofile", "", "memory profile output file")
 }
 
 var cmdVolume = &Command{
@@ -75,6 +79,7 @@ func runVolume(cmd *Command, args []string) bool {
 		*v.maxCpu = runtime.NumCPU()
 	}
 	runtime.GOMAXPROCS(*v.maxCpu)
+	util.SetupProfiling(*v.cpuProfile, *v.memProfile)
 
 	//Set multiple folders and each folder's max volume count limit'
 	v.folders = strings.Split(*volumeFolders, ",")
@@ -156,7 +161,7 @@ func runVolume(cmd *Command, args []string) bool {
 		}()
 	}
 
-	OnInterrupt(func() {
+	util.OnInterrupt(func() {
 		volumeServer.Shutdown()
 	})
 
diff --git a/weed/util/pprof.go b/weed/util/pprof.go
new file mode 100644
index 000000000..94bcdb8b3
--- /dev/null
+++ b/weed/util/pprof.go
@@ -0,0 +1,33 @@
+package util
+
+import (
+	"os"
+	"runtime/pprof"
+
+	"github.com/chrislusf/seaweedfs/weed/glog"
+)
+
+func SetupProfiling(cpuProfile, memProfile string) {
+	if cpuProfile != "" {
+		f, err := os.Create(cpuProfile)
+		if err != nil {
+			glog.Fatal(err)
+		}
+		pprof.StartCPUProfile(f)
+		defer pprof.StopCPUProfile()
+		OnInterrupt(func() {
+			pprof.StopCPUProfile()
+		})
+	}
+	if memProfile != "" {
+		f, err := os.Create(memProfile)
+		if err != nil {
+			glog.Fatal(err)
+		}
+		OnInterrupt(func() {
+			pprof.WriteHeapProfile(f)
+			f.Close()
+		})
+	}
+
+}
diff --git a/weed/command/signal_handling.go b/weed/util/signal_handling.go
similarity index 96%
rename from weed/command/signal_handling.go
rename to weed/util/signal_handling.go
index 182e2754d..7da898738 100644
--- a/weed/command/signal_handling.go
+++ b/weed/util/signal_handling.go
@@ -1,6 +1,6 @@
 // +build !plan9
 
-package command
+package util
 
 import (
 	"os"
diff --git a/weed/command/signal_handling_notsupported.go b/weed/util/signal_handling_notsupported.go
similarity index 75%
rename from weed/command/signal_handling_notsupported.go
rename to weed/util/signal_handling_notsupported.go
index dfcc24a3e..c389cfb7e 100644
--- a/weed/command/signal_handling_notsupported.go
+++ b/weed/util/signal_handling_notsupported.go
@@ -1,6 +1,6 @@
 // +build plan9
 
-package command
+package util
 
 func OnInterrupt(fn func()) {
 }