From bc70a0a0daa2a088c410a78d5ca8e27fc5efe722 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-91-211.ec2.internal>
Date: Wed, 13 Nov 2024 03:52:09 +0000
Subject: [PATCH] add pprof to disperser

---
 disperser/apiserver/config.go | 24 ++++++++++++++++++++++++
 disperser/apiserver/server.go |  7 +++++++
 go.mod                        |  2 +-
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/disperser/apiserver/config.go b/disperser/apiserver/config.go
index 947e45ac95..6d2237af36 100644
--- a/disperser/apiserver/config.go
+++ b/disperser/apiserver/config.go
@@ -26,6 +26,9 @@ const (
 	AllowlistFileFlagName            = "auth.allowlist-file"
 	AllowlistRefreshIntervalFlagName = "auth.allowlist-refresh-interval"
 
+	PprofHttpPortName = "auth.pprof-http-port"
+	EnablePprofName   = "auth.enable-pprof"
+
 	RetrievalBlobRateFlagName   = "auth.retrieval-blob-rate"
 	RetrievalThroughputFlagName = "auth.retrieval-throughput"
 
@@ -67,6 +70,8 @@ type RateConfig struct {
 
 	AllowlistFile            string
 	AllowlistRefreshInterval time.Duration
+
+	PprofConfig PprofConfig
 }
 
 func AllowlistFileFlag(envPrefix string) cli.Flag {
@@ -137,6 +142,19 @@ func CLIFlags(envPrefix string) []cli.Flag {
 			EnvVar:   common.PrefixEnvVar(envPrefix, "RETRIEVAL_BYTE_RATE"),
 			Required: true,
 		},
+		cli.StringFlag{
+			Name:     PprofHttpPortName,
+			Usage:    "the http port which the pprof server is listening",
+			Required: false,
+			Value:    "6060",
+			EnvVar:   common.PrefixEnvVar(envPrefix, "PPROF_HTTP_PORT"),
+		},
+		cli.BoolFlag{
+			Name:     EnablePprofName,
+			Usage:    "start prrof server",
+			Required: false,
+			EnvVar:   common.PrefixEnvVar(envPrefix, "ENABLE_PPROF"),
+		},
 	}
 }
 
@@ -234,6 +252,11 @@ func ReadCLIConfig(c *cli.Context) (RateConfig, error) {
 		}
 	}
 
+	pprofConfig := PprofConfig{
+		HTTPPort:    c.String(PprofHttpPortName),
+		EnablePprof: c.Bool(EnablePprofName),
+	}
+
 	return RateConfig{
 		QuorumRateInfos:          quorumRateInfos,
 		ClientIPHeader:           c.String(ClientIPHeaderFlagName),
@@ -242,5 +265,6 @@ func ReadCLIConfig(c *cli.Context) (RateConfig, error) {
 		RetrievalThroughput:      common.RateParam(c.Int(RetrievalThroughputFlagName)),
 		AllowlistFile:            c.String(AllowlistFileFlagName),
 		AllowlistRefreshInterval: c.Duration(AllowlistRefreshIntervalFlagName),
+		PprofConfig:              pprofConfig,
 	}, nil
 }
diff --git a/disperser/apiserver/server.go b/disperser/apiserver/server.go
index 001d663ed1..1f92eada83 100644
--- a/disperser/apiserver/server.go
+++ b/disperser/apiserver/server.go
@@ -812,6 +812,13 @@ func (s *DispersalServer) GetRateConfig() *RateConfig {
 }
 
 func (s *DispersalServer) Start(ctx context.Context) error {
+	pprofProfiler := NewPprofProfiler(s.rateConfig.PprofConfig.HTTPPort, s.logger)
+	if s.rateConfig.PprofConfig.EnablePprof {
+		port := s.rateConfig.PprofConfig.HTTPPort
+		go pprofProfiler.Start(port, s.logger)
+		s.logger.Info("Enabled pprof for disperser apiserver", "port", port)
+	}
+
 	go func() {
 		t := time.NewTicker(s.rateConfig.AllowlistRefreshInterval)
 		defer t.Stop()
diff --git a/go.mod b/go.mod
index ab0174e66d..df8d5f9aa5 100644
--- a/go.mod
+++ b/go.mod
@@ -15,7 +15,6 @@ require (
 	github.com/fxamacker/cbor/v2 v2.5.0
 	github.com/gin-contrib/logger v0.2.6
 	github.com/gin-gonic/gin v1.9.1
-	github.com/golang/protobuf v1.5.4
 	github.com/hashicorp/go-multierror v1.1.1
 	github.com/jedib0t/go-pretty/v6 v6.5.9
 	github.com/joho/godotenv v1.5.1
@@ -102,6 +101,7 @@ require (
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang-jwt/jwt v3.2.2+incompatible // indirect
 	github.com/golang-jwt/jwt/v4 v4.5.0 // indirect
+	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-bexpr v0.1.10 // indirect