From e93ba62ba7952b0a78d00f5de1d01c25a07f28ec Mon Sep 17 00:00:00 2001
From: Kevin Petremann <k.petremann@criteo.com>
Date: Thu, 17 Oct 2024 19:49:26 +0200
Subject: [PATCH] feat(debug): optional pprof

---
 cmd/data-aggregation-api/main.go |  2 +-
 internal/api/router/manager.go   | 13 ++++++++++++-
 internal/config/settings.go      |  7 +++++++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/cmd/data-aggregation-api/main.go b/cmd/data-aggregation-api/main.go
index f46aaee..d6dbb2f 100644
--- 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
index 3d3e617..70c4777 100644
--- 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
index 3f6482b..b2c6657 100644
--- 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 {