add pprof to disperser (Layr-Labs#927)

Co-authored-by: Ubuntu <[email protected]>

disperser/apiserver/pprof.go

@@ -0,0 +1,31 @@
+package apiserver
+
+import (
+	"fmt"
+	"net/http"
+
+	_ "net/http/pprof"
+
+	"github.com/Layr-Labs/eigensdk-go/logging"
+)
+
+type PprofProfiler struct {
+	logger   logging.Logger
+	httpPort string
+}
+
+func NewPprofProfiler(httpPort string, logger logging.Logger) *PprofProfiler {
+	return &PprofProfiler{
+		logger:   logger.With("component", "PprofProfiler"),
+		httpPort: httpPort,
+	}
+}
+
+// Start the pprof server
+func (p *PprofProfiler) Start(logger logging.Logger) {
+	pprofAddr := fmt.Sprintf("%s:%s", "0.0.0.0", p.httpPort)
+
+	if err := http.ListenAndServe(pprofAddr, nil); err != nil {
+		p.logger.Error("pprof server failed", "error", err, "pprofAddr", pprofAddr)
+	}
+}

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()

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),

disperser/cmd/apiserver/flags/flags.go

@@ -155,6 +155,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{
@@ -247,6 +260,8 @@ var optionalFlags = []cli.Flag{
 	GlobalRateTableName,
 	OnchainStateRefreshInterval,
 	MaxNumSymbolsPerBlob,
+	PprofHttpPort,
+	EnablePprof,
 }
 
 // Flags contains the list of configuration options available to the binary.

disperser/server_config.go

@@ -9,4 +9,7 @@ const (
 type ServerConfig struct {
 	GrpcPort    string
 	GrpcTimeout time.Duration
+
+	PprofHttpPort string
+	EnablePprof   bool
 }

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