Merge pull request #217 from heroku/vmc-pprof

Refactor and Start PProf debug server

Author: Vamsi-Mundra

cmdutil/debug/config.go

@@ -3,4 +3,24 @@ package debug
 // Config describes the configurable parameters for debugging.
 type Config struct {
 	Port int `env:"DEBUG_PORT,default=9999"`
+	PProf
+}
+
+type PProf struct {
+	Port int `env:"DEBUG_PPROF_PORT,default=9998"`
+	// This controls how much of fraction of mutexes we need to consider for profiling.
+	MutexProfileFraction int `env:"DBEUG_PPROF_MUTEX_PROFILE_FRACTION,default=2"`
+	// This controls how much blocking time in nano seconds we need to consider.
+	// Default 10 indicates, we consider all 10 ns blocking events .
+	BlockProfileRate int `env:"DBEUG_PPROF_BLOCK_PROFILE_RATE,default=10"`
+	// Heap Profiling is enabled by default.
+	// This controls the frequency of heap allocation sampling.
+	// It defines the number of bytes allocated between samples.
+	// Default is actual default value of MemProfileRate
+	MemProfileRate int  `env:"DBEUG_PPROF_MEM_PROFILE_RATE,default=524288"`
+	Enabled        bool `env:"DEBUG_PPROF_ENABLE,default=false"`
+	// Mutex Profiling is not enabled by default
+	EnableMutexProfiling bool `env:"DEBUG_PPROF_MUTEX_PROFILE_ENABLE,default=false"`
+	// Block Profiling is not enabled default
+	EnableBlockProfiling bool `env:"DEBUG_PPROF_BLOCK_PROFILE_ENABLE,default=false"`
 }

cmdutil/debug/debug.go

@@ -16,10 +16,12 @@ package debug
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net/http"
 	"net/http/pprof"
 	"runtime"
+	"sync"
 	"time"
 
 	"github.com/google/gops/agent"
@@ -31,25 +33,75 @@ import (
 // Connect to the debug server with gops:
 //
 //	gops stack localhost:PORT
-func New(l logrus.FieldLogger, port int) *Server {
-	return &Server{
+//
+// Connect to the pprof server with pprof.port when pprof enabled
+func New(l logrus.FieldLogger, config Config) *Server {
+	server := &Server{
 		logger: l,
-		addr:   fmt.Sprintf("127.0.0.1:%d", port),
+		addr:   fmt.Sprintf("127.0.0.1:%d", config.Port),
 		done:   make(chan struct{}),
 	}
+	if config.Enabled {
+		server.pprof = NewPProfServer(l, &config.PProf)
+	}
+	return server
 }
 
 // Server wraps a gops server for easy use with oklog/group.
 type Server struct {
 	logger logrus.FieldLogger
 	addr   string
 	done   chan struct{}
+	pprof  *PProfServer
 }
 
 // Run starts the debug server.
 //
-// It implements oklog group's runFn.
+// It implements oklog group's runFn and pprof.
 func (s *Server) Run() error {
+
+	var wg sync.WaitGroup
+	var gopsErr error
+	var pprofErr error
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		gopsErr = s.RunGOPS()
+		if gopsErr != nil {
+			s.logger.WithError(gopsErr).Error("gops server failed")
+		}
+	}()
+
+	if s.pprof != nil {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			pprofErr = s.pprof.Run()
+			if pprofErr != nil {
+				s.pprof.logger.WithError(pprofErr).Error("pprof server failed")
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	var err error
+	if gopsErr != nil {
+		err = fmt.Errorf("gops error: %w", gopsErr)
+	}
+	if pprofErr != nil {
+		errPProf := fmt.Errorf("pprof error: %w", pprofErr)
+		err = errors.Join(err, errPProf)
+	}
+
+	return err
+}
+
+// Run starts the debug server.
+//
+// It implements oklog group's runFn.
+func (s *Server) RunGOPS() error {
 	s.logger.WithFields(logrus.Fields{
 		"at":      "binding",
 		"service": "debug",
@@ -70,11 +122,15 @@ func (s *Server) Run() error {
 
 // Stop shuts down the debug server.
 //
-// It implements oklog group's interruptFn.
+// It implements oklog group's interruptFn and pprof stop.
 func (s *Server) Stop(_ error) {
 	agent.Close()
 
 	close(s.done)
+
+	if s.pprof != nil {
+		s.pprof.Stop(nil)
+	}
 }
 
 // PProfServer wraps a pprof server.
@@ -85,45 +141,34 @@ type PProfServer struct {
 	pprofServer *http.Server
 }
 
-// ProfileConfig holds the configuration for the pprof server.
-type PProfServerConfig struct {
-	Addr                 string
-	MutexProfileFraction int
-}
+// NewPProfServer sets up a pprof server with configurable profiling types and returns a PProfServer instance.
+func NewPProfServer(l logrus.FieldLogger, pprofConfig *PProf) *PProfServer {
 
-// defaultMutexProfileFraction is the default value for MutexProfileFraction
-const defaultMutexProfileFraction = 2
+	runtime.MemProfileRate = pprofConfig.MemProfileRate
 
-// NewPProfServer sets up a pprof server with configurable profiling types and returns a PProfServer instance.
-func NewPProfServer(config PProfServerConfig, l logrus.FieldLogger) *PProfServer {
-	if config.Addr == "" {
-		config.Addr = "127.0.0.1:9998" // Default port
+	if pprofConfig.EnableMutexProfiling {
+		runtime.SetMutexProfileFraction(pprofConfig.MutexProfileFraction)
 	}
 
-	// Use a local variable for the mutex profile fraction
-	mpf := defaultMutexProfileFraction
-	if config.MutexProfileFraction != 0 {
-		mpf = config.MutexProfileFraction
+	if pprofConfig.EnableBlockProfiling {
+		runtime.SetBlockProfileRate(pprofConfig.BlockProfileRate)
 	}
-	runtime.SetMutexProfileFraction(mpf)
 
 	httpServer := &http.Server{
-		Addr:              config.Addr,
+		Addr:              fmt.Sprintf("127.0.0.1:%d", pprofConfig.Port),
 		Handler:           http.HandlerFunc(pprof.Index),
 		ReadHeaderTimeout: 5 * time.Second,
 	}
 
 	return &PProfServer{
 		logger:      l,
-		addr:        config.Addr,
+		addr:        httpServer.Addr,
 		done:        make(chan struct{}),
 		pprofServer: httpServer,
 	}
 }
 
 // Run starts the pprof server.
-//
-// It implements oklog group's runFn.
 func (s *PProfServer) Run() error {
 	if s.pprofServer == nil {
 		return fmt.Errorf("pprofServer is nil")
@@ -144,8 +189,6 @@ func (s *PProfServer) Run() error {
 }
 
 // Stop shuts down the pprof server.
-//
-// It implements oklog group's interruptFn.
 func (s *PProfServer) Stop(_ error) {
 	if s.pprofServer != nil {
 		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

cmdutil/debug/debug_test.go

@@ -13,40 +13,46 @@ func TestNewPProfServer(t *testing.T) {
 	logger := logrus.New()
 
 	tests := []struct {
-		name                  string
-		config                PProfServerConfig
-		expectedAddr          string
-		expectedMutexFraction int
+		name                   string
+		expectedAddr           string
+		pprofConfig            *PProf
+		expectedMemProfileRate int
 	}{
 		{
-			name:                  "DefaultAddr",
-			config:                PProfServerConfig{},
-			expectedAddr:          "127.0.0.1:9998",
-			expectedMutexFraction: defaultMutexProfileFraction,
+			name:         "test port as 9998 and mpf as 2",
+			expectedAddr: "127.0.0.1:9998",
+			pprofConfig: &PProf{
+				Port:           9998,
+				Enabled:        true,
+				MemProfileRate: 524288,
+			},
+			expectedMemProfileRate: 524288,
 		},
 		{
-			name:                  "CustomAddr",
-			config:                PProfServerConfig{Addr: "127.0.0.1:9090"},
-			expectedAddr:          "127.0.0.1:9090",
-			expectedMutexFraction: defaultMutexProfileFraction,
-		},
-		{
-			name:                  "CustomMutexProfileFraction",
-			config:                PProfServerConfig{MutexProfileFraction: 5},
-			expectedAddr:          "127.0.0.1:9998",
-			expectedMutexFraction: 5,
+			name:         "test port as 9997 and mpf as 4",
+			expectedAddr: "127.0.0.1:9997",
+			pprofConfig: &PProf{
+				Port:           9997,
+				Enabled:        true,
+				MemProfileRate: 524287,
+			},
+			expectedMemProfileRate: 524287,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			server := NewPProfServer(tt.config, logger)
+			server := NewPProfServer(logger, tt.pprofConfig)
 
 			// Check server address
 			if server.addr != tt.expectedAddr {
 				t.Errorf("NewPProfServer() addr = %v, want %v", server.addr, tt.expectedAddr)
 			}
 
+			if runtime.MemProfileRate != tt.expectedMemProfileRate {
+				t.Errorf("MemProfileRate expected  %v, got  %v", tt.expectedMemProfileRate, runtime.MemProfileRate)
+			}
+
 			// Start the server
 			go func() {
 				if err := server.Run(); err != nil {
@@ -57,12 +63,6 @@ func TestNewPProfServer(t *testing.T) {
 			// Give the server a moment to start
 			time.Sleep(100 * time.Millisecond)
 
-			// Check mutex profile fraction
-			if got := runtime.SetMutexProfileFraction(0); got != tt.expectedMutexFraction {
-				t.Errorf("runtime.SetMutexProfileFraction() = %v, want %v", got, tt.expectedMutexFraction)
-			}
-			runtime.SetMutexProfileFraction(tt.expectedMutexFraction) // Reset to the expected value
-
 			// Perform HTTP GET request to the root path
 			url := "http://" + server.addr + "/debug/pprof/"
 			client := &http.Client{}

cmdutil/service/standard.go

@@ -73,7 +73,7 @@ func New(appConfig interface{}, ofs ...OptionFunc) *Standard {
 		s.Add(cmdutil.NewContextServer(l2met.Run))
 	}
 
-	s.Add(debug.New(logger, sc.Debug.Port))
+	s.Add(debug.New(logger, sc.Debug))
 	s.Add(signals.NewServer(logger, syscall.SIGINT, syscall.SIGTERM))
 
 	// only setup an exporter if indicated && the AgentAddress is set