From 0ff477a726de3c2556bc9d3a598c03e1e78efcee Mon Sep 17 00:00:00 2001
From: Swapnil Mhamane <swapnil.mhamane@sap.com>
Date: Fri, 12 Oct 2018 10:50:31 +0530
Subject: [PATCH] Add dynamic profiling endpoint

Signed-off-by: Swapnil Mhamane <swapnil.mhamane@sap.com>
---
 cmd/server.go         |  2 ++
 cmd/types.go          |  4 +++-
 pkg/server/httpAPI.go | 42 ++++++++++++++++++++++++++++++++++--------
 3 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/cmd/server.go b/cmd/server.go
index a14559880..73d7a9afa 100644
--- a/cmd/server.go
+++ b/cmd/server.go
@@ -76,6 +76,7 @@ func NewServerCommand(stopCh <-chan struct{}) *cobra.Command {
 				Logger:          logger,
 				Status:          http.StatusServiceUnavailable,
 				StopCh:          make(chan struct{}),
+				EnableProfiling: enableProfiling,
 			}
 			logger.Info("Regsitering the http request handlers...")
 			handler.RegisterHandler()
@@ -187,6 +188,7 @@ func NewServerCommand(stopCh <-chan struct{}) *cobra.Command {
 // initializeServerFlags adds the flags to <cmd>
 func initializeServerFlags(serverCmd *cobra.Command) {
 	serverCmd.Flags().IntVarP(&port, "server-port", "p", defaultServerPort, "port on which server should listen")
+	serverCmd.Flags().BoolVar(&enableProfiling, "enable-profiling", false, "enable profiling")
 }
 
 // ProbeEtcd will make the snapshotter probe for etcd endpoint to be available
diff --git a/cmd/types.go b/cmd/types.go
index b70ede261..6a11a6c2e 100644
--- a/cmd/types.go
+++ b/cmd/types.go
@@ -43,7 +43,9 @@ var (
 	caFile                         string
 
 	//server flags
-	port int
+	port            int
+	enableProfiling bool
+
 	//restore flags
 	restoreCluster      string
 	restoreClusterToken string
diff --git a/pkg/server/httpAPI.go b/pkg/server/httpAPI.go
index e5c0a65a3..8aa2483ca 100644
--- a/pkg/server/httpAPI.go
+++ b/pkg/server/httpAPI.go
@@ -17,6 +17,7 @@ package server
 import (
 	"fmt"
 	"net/http"
+	"net/http/pprof"
 	"sync"
 
 	"github.com/gardener/etcd-backup-restore/pkg/initializer"
@@ -40,25 +41,50 @@ type HTTPHandler struct {
 	initializationStatus      string
 	Status                    int
 	StopCh                    chan struct{}
+	EnableProfiling           bool
 }
 
 // RegisterHandler registers the handler for different requests
 func (h *HTTPHandler) RegisterHandler() {
+	mux := http.NewServeMux()
+	if h.EnableProfiling {
+		registerPProfHandler(mux)
+	}
+
 	h.initializationStatus = "New"
-	http.HandleFunc("/initialization/start", h.serveInitialize)
-	http.HandleFunc("/initialization/status", h.serveInitializationStatus)
-	http.HandleFunc("/healthz", h.serveHealthz)
+	mux.HandleFunc("/initialization/start", h.serveInitialize)
+	mux.HandleFunc("/initialization/status", h.serveInitializationStatus)
+	mux.HandleFunc("/healthz", h.serveHealthz)
+
+	h.server = &http.Server{
+		Addr:    fmt.Sprintf(":%d", h.Port),
+		Handler: mux,
+	}
 	return
 }
 
+// registerPProfHandler registers the PProf handler for profiling.
+func registerPProfHandler(mux *http.ServeMux) {
+	mux.HandleFunc("/debug/pprof/", pprof.Index)
+	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+	mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+	mux.HandleFunc("/debug/pprof/heap", pprof.Handler("heap").ServeHTTP)
+	mux.HandleFunc("/debug/pprof/goroutine", pprof.Handler("goroutine").ServeHTTP)
+	mux.HandleFunc("/debug/pprof/threadcreate", pprof.Handler("threadcreate").ServeHTTP)
+	mux.HandleFunc("/debug/pprof/block", pprof.Handler("block").ServeHTTP)
+	mux.HandleFunc("/debug/pprof/mutex", pprof.Handler("mutex").ServeHTTP)
+}
+
 // Start start the http server to listen for request
 func (h *HTTPHandler) Start() {
-	h.server = &http.Server{
-		Addr:    fmt.Sprintf(":%d", h.Port),
-		Handler: nil,
-	}
+	h.Logger.Infof("Starting Http server at addr: %s", h.server.Addr)
 	err := h.server.ListenAndServe()
-	h.Logger.Fatalf("Failed to start http server: %v", err)
+	if err != nil && err != http.ErrServerClosed {
+		h.Logger.Fatalf("Failed to start http server: %v", err)
+	}
+	h.Logger.Infof("Http server closed gracefully.")
 	return
 }