From 57177a2f27e445cb5b8645a05c49f8b494be3945 Mon Sep 17 00:00:00 2001
From: Zohaib Sibte Hassan <zohaib@sibte.pk>
Date: Thu, 28 Dec 2023 08:57:20 -0800
Subject: [PATCH] Enable CLI based PPROF

---
 cfg/config.go |  1 +
 marmot.go     | 20 +++++++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/cfg/config.go b/cfg/config.go
index 20d3fce..5e6addf 100644
--- a/cfg/config.go
+++ b/cfg/config.go
@@ -118,6 +118,7 @@ var SaveSnapshotFlag = flag.Bool("save-snapshot", false, "Only take snapshot and
 var ClusterAddrFlag = flag.String("cluster-addr", "", "Cluster listening address")
 var ClusterPeersFlag = flag.String("cluster-peers", "", "Comma separated list of clusters")
 var LeafServerFlag = flag.String("leaf-servers", "", "Comma separated list of leaf servers")
+var ProfServer = flag.String("pprof", "", "PProf listening address")
 
 var DataRootDir = os.TempDir()
 var Config = &Configuration{
diff --git a/marmot.go b/marmot.go
index c036ef3..41b8a66 100644
--- a/marmot.go
+++ b/marmot.go
@@ -4,6 +4,9 @@ import (
 	"context"
 	"flag"
 	"io"
+	"net/http"
+	"net/http/pprof"
+	_ "net/http/pprof"
 	"os"
 	"time"
 
@@ -22,7 +25,6 @@ import (
 
 func main() {
 	flag.Parse()
-
 	err := cfg.Load(*cfg.ConfigPathFlag)
 	if err != nil {
 		panic(err)
@@ -44,6 +46,22 @@ func main() {
 		log.Logger = gLog.Level(zerolog.InfoLevel)
 	}
 
+	if *cfg.ProfServer != "" {
+		go func() {
+			mux := http.NewServeMux()
+			mux.HandleFunc("/debug/pprof/", pprof.Index)
+			mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+			mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+			mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+			mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+
+			err := http.ListenAndServe(*cfg.ProfServer, mux)
+			if err != nil {
+				log.Error().Err(err).Msg("unable to bind profiler server")
+			}
+		}()
+	}
+
 	log.Debug().Msg("Initializing telemetry")
 	telemetry.InitializeTelemetry()