From c5153bee741e788a14398ba74e09bcfc79f011b8 Mon Sep 17 00:00:00 2001
From: buckhx <buckhx>
Date: Fri, 25 Mar 2016 16:46:41 -0400
Subject: [PATCH] added profile flag

---
 cli.go           | 10 +++++++++-
 geofence/http.go | 18 ++++++++++++++++--
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/cli.go b/cli.go
index 23da892..7b0a298 100644
--- a/cli.go
+++ b/cli.go
@@ -34,6 +34,10 @@ func client(args []string) {
 			Value: "8080",
 			Usage: "Port to bind to",
 		},
+		cli.BoolFlag{
+			Name:  "profile",
+			Usage: "Mount profiling endpoints",
+		},
 	}
 	app.Action = func(c *cli.Context) {
 		args := c.Args()
@@ -51,8 +55,12 @@ func client(args []string) {
 		if err != nil {
 			die(c, err.Error())
 		}
+		prof := c.Bool("profile")
+		if prof {
+			log.Println("Profiling available at /debug/pprof/")
+		}
 		port := fmt.Sprintf(":%s", c.String("port"))
-		err = geofence.ListenAndServe(port, fences)
+		err = geofence.ListenAndServe(port, fences, prof)
 		die(c, err.Error())
 	}
 	app.Run(args)
diff --git a/geofence/http.go b/geofence/http.go
index 402a7ff..b5de26e 100644
--- a/geofence/http.go
+++ b/geofence/http.go
@@ -6,7 +6,7 @@ import (
 	"io/ioutil"
 	"log"
 	"net/http"
-	_ "net/http/pprof"
+	"net/http/pprof"
 	"strconv"
 
 	"github.com/buckhx/diglet/geo"
@@ -15,7 +15,7 @@ import (
 
 var fences FenceIndex
 
-func ListenAndServe(addr string, idx FenceIndex) error {
+func ListenAndServe(addr string, idx FenceIndex, profile bool) error {
 	log.Printf("Fencing on address %s\n", addr)
 	defer log.Printf("Done Fencing\n")
 	fences = idx
@@ -26,6 +26,9 @@ func ListenAndServe(addr string, idx FenceIndex) error {
 	router.POST("/fence/:name/add", postAdd)
 	router.POST("/fence/:name/search", postSearch)
 	router.GET("/fence/:name/search", getSearch)
+	if profile {
+		attachProfiler(router)
+	}
 	return http.ListenAndServe(addr, router)
 }
 
@@ -131,3 +134,14 @@ func getEngarde(w http.ResponseWriter, r *http.Request, params httprouter.Params
 	w.Header().Set("Content-Length", fmt.Sprint(len(response)))
 	fmt.Fprint(w, response)
 }
+
+func attachProfiler(router *httprouter.Router) {
+	router.HandlerFunc("GET", "/debug/pprof/", pprof.Index)
+	router.HandlerFunc("GET", "/debug/pprof/cmdline", pprof.Cmdline)
+	router.HandlerFunc("GET", "/debug/pprof/profile", pprof.Profile)
+	router.HandlerFunc("GET", "/debug/pprof/symbol", pprof.Symbol)
+	router.Handler("GET", "/debug/pprof/heap", pprof.Handler("heap"))
+	router.Handler("GET", "/debug/pprof/block", pprof.Handler("block"))
+	router.Handler("GET", "/debug/pprof/goroutine", pprof.Handler("goroutine"))
+	router.Handler("GET", "/debug/pprof/threadcreate", pprof.Handler("threadcreate"))
+}