From aef34367330af1a5ea89004f4f47ec9ce62783a6 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-91-211.ec2.internal>
Date: Thu, 21 Nov 2024 23:37:31 +0000
Subject: [PATCH] add pprof

---
 disperser/apiserver/server.go          |  6 ++++++
 disperser/cmd/apiserver/config.go      |  6 ++++--
 disperser/cmd/apiserver/flags/flags.go | 15 +++++++++++++++
 disperser/server_config.go             |  3 +++
 go.mod                                 |  2 +-
 5 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/disperser/apiserver/server.go b/disperser/apiserver/server.go
index 87e82b5a1..34f60efd4 100644
--- a/disperser/apiserver/server.go
+++ b/disperser/apiserver/server.go
@@ -819,6 +819,12 @@ func (s *DispersalServer) GetRateConfig() *RateConfig {
 }
 
 func (s *DispersalServer) Start(ctx context.Context) error {
+	pprofProfiler := NewPprofProfiler(s.serverConfig.PprofHttpPort, s.logger)
+	if s.serverConfig.EnablePprof {
+		go pprofProfiler.Start(s.logger)
+		s.logger.Info("Enabled pprof for disperser apiserver", "port", s.serverConfig.PprofHttpPort)
+	}
+
 	go func() {
 		t := time.NewTicker(s.rateConfig.AllowlistRefreshInterval)
 		defer t.Stop()
diff --git a/disperser/cmd/apiserver/config.go b/disperser/cmd/apiserver/config.go
index 6e25b8af6..a27bd0b79 100644
--- a/disperser/cmd/apiserver/config.go
+++ b/disperser/cmd/apiserver/config.go
@@ -95,8 +95,10 @@ func NewConfig(ctx *cli.Context) (Config, error) {
 		DisperserVersion: DisperserVersion(version),
 		AwsClientConfig:  aws.ReadClientConfig(ctx, flags.FlagPrefix),
 		ServerConfig: disperser.ServerConfig{
-			GrpcPort:    ctx.GlobalString(flags.GrpcPortFlag.Name),
-			GrpcTimeout: ctx.GlobalDuration(flags.GrpcTimeoutFlag.Name),
+			GrpcPort:      ctx.GlobalString(flags.GrpcPortFlag.Name),
+			GrpcTimeout:   ctx.GlobalDuration(flags.GrpcTimeoutFlag.Name),
+			PprofHttpPort: ctx.GlobalString(flags.PprofHttpPort.Name),
+			EnablePprof:   ctx.GlobalBool(flags.EnablePprof.Name),
 		},
 		BlobstoreConfig: blobstore.Config{
 			BucketName: ctx.GlobalString(flags.S3BucketNameFlag.Name),
diff --git a/disperser/cmd/apiserver/flags/flags.go b/disperser/cmd/apiserver/flags/flags.go
index 45eb41133..63f8965cf 100644
--- a/disperser/cmd/apiserver/flags/flags.go
+++ b/disperser/cmd/apiserver/flags/flags.go
@@ -154,6 +154,19 @@ var (
 		EnvVar:   common.PrefixEnvVar(envVarPrefix, "MAX_NUM_SYMBOLS_PER_BLOB"),
 		Required: false,
 	}
+	PprofHttpPort = cli.StringFlag{
+		Name:     common.PrefixFlag(FlagPrefix, "pprof-http-port"),
+		Usage:    "the http port which the pprof server is listening",
+		Required: false,
+		Value:    "6060",
+		EnvVar:   common.PrefixEnvVar(envVarPrefix, "PPROF_HTTP_PORT"),
+	}
+	EnablePprof = cli.BoolFlag{
+		Name:     common.PrefixFlag(FlagPrefix, "enable-pprof"),
+		Usage:    "start prrof server",
+		Required: false,
+		EnvVar:   common.PrefixEnvVar(envVarPrefix, "ENABLE_PPROF"),
+	}
 )
 
 var kzgFlags = []cli.Flag{
@@ -246,6 +259,8 @@ var optionalFlags = []cli.Flag{
 	GlobalRateTableName,
 	OnchainStateRefreshInterval,
 	MaxNumSymbolsPerBlob,
+	PprofHttpPort,
+	EnablePprof,
 }
 
 // Flags contains the list of configuration options available to the binary.
diff --git a/disperser/server_config.go b/disperser/server_config.go
index b094fca2a..152ff675c 100644
--- a/disperser/server_config.go
+++ b/disperser/server_config.go
@@ -9,4 +9,7 @@ const (
 type ServerConfig struct {
 	GrpcPort    string
 	GrpcTimeout time.Duration
+
+	PprofHttpPort string
+	EnablePprof   bool
 }
diff --git a/go.mod b/go.mod
index 38c95f915..45e0bbe49 100644
--- a/go.mod
+++ b/go.mod
@@ -37,6 +37,7 @@ require (
 	go.uber.org/mock v0.4.0
 	golang.org/x/exp v0.0.0-20231110203233-9a3e6036ecaa
 	golang.org/x/sync v0.8.0
+	golang.org/x/time v0.5.0
 	google.golang.org/grpc v1.64.1
 )
 
@@ -151,7 +152,6 @@ require (
 	go.uber.org/zap v1.27.0 // indirect
 	golang.org/x/arch v0.4.0 // indirect
 	golang.org/x/oauth2 v0.22.0 // indirect
-	golang.org/x/time v0.5.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240814211410-ddb44dafa142 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240827150818-7e3bb234dfed // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect