Merge pull request #435 from tidwall/pprof

Added pprof flags

cmd/tile38-server/main.go

@@ -7,9 +7,11 @@ import (
 	"io/ioutil"
 	"net"
 	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"os/signal"
 	"runtime"
+	"runtime/pprof"
 	"strconv"
 	"strings"
 	"sync"
@@ -32,17 +34,35 @@ var (
 	devMode     bool
 	quiet       bool
 	pidfile     string
+	cpuprofile  string
+	memprofile  string
+	pprofport   int
 )
 
 // TODO: Set to false in 2.*
 var httpTransport = true
 
+////////////////////////////////////////////////////////////////////////////////
+//
 // Fire up a webhook test server by using the --webhook-http-consumer-port
 // for example
 //   $ ./tile38-server --webhook-http-consumer-port 9999
 //
 // The create hooks like such...
 //   SETHOOK myhook http://localhost:9999/myhook NEARBY mykey FENCE POINT 33.5 -115.5 1000
+//
+////////////////////////////////////////////////////////////////////////////////
+//
+// Memory profiling - start the server with the -pprofport flag
+//
+//   $ ./tile38-server -pprofport 6060
+//
+// Then, at any point, from a different terminal execute:
+//   $ go tool pprof -svg http://localhost:6060/debug/pprof/heap > out.svg
+//
+// Load the SVG into a web browser to visualize the memory usage
+//
+////////////////////////////////////////////////////////////////////////////////
 
 type hserver struct{}
 
@@ -240,6 +260,9 @@ Developer Options:
 	flag.BoolVar(&verbose, "v", false, "Enable verbose logging.")
 	flag.BoolVar(&quiet, "q", false, "Quiet logging. Totally silent.")
 	flag.BoolVar(&veryVerbose, "vv", false, "Enable very verbose logging.")
+	flag.IntVar(&pprofport, "pprofport", 0, "pprofport http at port")
+	flag.StringVar(&cpuprofile, "cpuprofile", "", "write cpu profile to `file`")
+	flag.StringVar(&memprofile, "memprofile", "", "write memory profile to `file`")
 	flag.Parse()
 
 	var logw io.Writer = os.Stderr
@@ -263,24 +286,75 @@ Developer Options:
 	if host != "" {
 		hostd = "Addr: " + host + ", "
 	}
-	var pidferr error
 
-	var cleanedup bool
-	var cleanupMu sync.Mutex
-	cleanup := func() {
-		cleanupMu.Lock()
-		defer cleanupMu.Unlock()
-		if cleanedup {
+	// pprof
+	if cpuprofile != "" {
+		log.Debugf("cpuprofile active")
+		f, err := os.Create(cpuprofile)
+		if err != nil {
+			log.Fatal("could not create CPU profile: ", err)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			log.Fatal("could not start CPU profile: ", err)
+		}
+	}
+	if memprofile != "" {
+		log.Debug("memprofile active")
+	}
+
+	var pprofcleanedup bool
+	var pprofcleanupMu sync.Mutex
+	pprofcleanup := func() {
+		pprofcleanupMu.Lock()
+		defer pprofcleanupMu.Unlock()
+		if pprofcleanedup {
 			return
 		}
 		// cleanup code
-		if pidfile != "" {
-			os.Remove(pidfile)
+		if cpuprofile != "" {
+			pprof.StopCPUProfile()
 		}
-		cleanedup = true
+		if memprofile != "" {
+			f, err := os.Create(memprofile)
+			if err != nil {
+				log.Fatal("could not create memory profile: ", err)
+			}
+			runtime.GC() // get up-to-date statistics
+			if err := pprof.WriteHeapProfile(f); err != nil {
+				log.Fatal("could not write memory profile: ", err)
+			}
+			f.Close()
+		}
+		pprofcleanedup = true
 	}
-	defer cleanup()
+	defer pprofcleanup()
 
+	if pprofport != 0 {
+		log.Debugf("pprof http at port %d", pprofport)
+		go func() {
+			log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", pprofport), nil))
+		}()
+	}
+
+	// pid file
+	var pidferr error
+	var pidcleanedup bool
+	var pidcleanupMu sync.Mutex
+	pidcleanup := func() {
+		if pidfile != "" {
+			pidcleanupMu.Lock()
+			defer pidcleanupMu.Unlock()
+			if pidcleanedup {
+				return
+			}
+			// cleanup code
+			if pidfile != "" {
+				os.Remove(pidfile)
+			}
+			pidcleanedup = true
+		}
+	}
+	defer pidcleanup()
 	if pidfile != "" {
 		pidferr := ioutil.WriteFile(pidfile, []byte(fmt.Sprintf("%d\n", os.Getpid())), 0666)
 		if pidferr == nil {
@@ -296,9 +370,8 @@ Developer Options:
 				continue
 			}
 			log.Warnf("signal: %v", s)
-			if pidfile != "" {
-				cleanup()
-			}
+			pidcleanup()
+			pprofcleanup()
 			switch {
 			default:
 				os.Exit(-1)