feat(debug): optional pprof

criteo · Oct 22, 2024 · e93ba62 · e93ba62
1 parent 23228f9
commit e93ba62
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/cmd/data-aggregation-api/main.go b/cmd/data-aggregation-api/main.go
@@ -85,7 +85,7 @@ func run() error {
 	triggerNewBuild := dispatchSingleRequest(newBuildRequest)
 
 	go job.StartBuildLoop(&deviceRepo, &reports, triggerNewBuild)
-	if err := router.NewManager(&deviceRepo, &reports, newBuildRequest).ListenAndServe(ctx, config.Cfg.API.ListenAddress, config.Cfg.API.ListenPort); err != nil {
+	if err := router.NewManager(&deviceRepo, &reports, newBuildRequest).ListenAndServe(ctx, config.Cfg.API.ListenAddress, config.Cfg.API.ListenPort, config.Cfg.Debug.Pprof.Enabled); err != nil {
 		return fmt.Errorf("webserver error: %w", err)
 	}
 

diff --git a/internal/api/router/manager.go b/internal/api/router/manager.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/http/pprof"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -42,7 +43,7 @@ func NewManager(deviceRepo DevicesRepository, reports *report.Repository, restar
 }
 
 // ListenAndServe starts to serve Web API requests.
-func (m *Manager) ListenAndServe(ctx context.Context, address string, port int) error {
+func (m *Manager) ListenAndServe(ctx context.Context, address string, port int, enablepprof bool) error {
 	defer func() {
 		close(m.newBuildRequest)
 		log.Warn().Msg("Shutdown.")
@@ -67,6 +68,16 @@ func (m *Manager) ListenAndServe(ctx context.Context, address string, port int)
 	router.GET("/v1/report/last/successful", withAuth.Wrap(m.getLastSuccessfulReport))
 	router.POST("/v1/build/trigger", withAuth.Wrap(m.triggerBuild))
 
+	if enablepprof {
+		router.HandlerFunc(http.MethodGet, "/debug/pprof/", pprof.Index)
+		router.HandlerFunc(http.MethodGet, "/debug/pprof/allocs", pprof.Index)
+		router.HandlerFunc(http.MethodGet, "/debug/pprof/goroutine", pprof.Index)
+		router.HandlerFunc(http.MethodGet, "/debug/pprof/heap", pprof.Index)
+		router.HandlerFunc(http.MethodGet, "/debug/pprof/profile", pprof.Profile)
+		router.HandlerFunc(http.MethodGet, "/debug/pprof/trace", pprof.Trace)
+		router.HandlerFunc(http.MethodGet, "/debug/pprof/symbol", pprof.Symbol)
+	}
+
 	listenSocket := fmt.Sprint(address, ":", port)
 	log.Info().Msgf("Start webserver - listening on %s", listenSocket)
 

diff --git a/internal/config/settings.go b/internal/config/settings.go
@@ -50,6 +50,11 @@ type Config struct {
 		Interval            time.Duration
 		AllDevicesMustBuild bool
 	}
+	Debug struct {
+		Pprof struct {
+			Enabled bool
+		}
+	}
 }
 
 type AuthConfig struct {
@@ -91,6 +96,8 @@ func setDefaults() {
 	viper.SetDefault("Authentication.LDAP.WorkersCount", defaultLDAPWorkersCount)
 	viper.SetDefault("Authentication.LDAP.Timeout", defaultLDAPTimeout)
 	viper.SetDefault("Authentication.LDAP.MaxConnectionLifetime", time.Minute)
+
+	viper.SetDefault("Debug.Pprof.Enabled", false)
 }
 
 func LoadConfig() error {