From e7f4e6722c2b802096c96979ec4d6ef3735f62cb Mon Sep 17 00:00:00 2001
From: Yash Thakkar <yathakka@amazon.com>
Date: Sun, 22 Dec 2024 06:32:26 +0000
Subject: [PATCH] adding flag to enable profiling

---
 cmd/main.go                     | 14 ++++++++++++++
 pkg/config/controller_config.go | 17 ++++++++++++-----
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index add7aea..580eec2 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -18,6 +18,8 @@ package main
 
 import (
 	"context"
+	"net/http"
+	_ "net/http/pprof"
 	"os"
 
 	"github.com/go-logr/logr"
@@ -70,6 +72,18 @@ func main() {
 		infoLogger.Error(err, "unable to load controller config")
 		os.Exit(1)
 	}
+
+	if controllerCFG.EnableProfiling {
+		infoLogger.Info("enable profiling true")
+		go func() {
+			infoLogger.Info("Starting pprof server at :7443/debug/pprof")
+			if err := http.ListenAndServe(":7443", nil); err != nil {
+				infoLogger.Error(err, "failed to start pprof server")
+				os.Exit(1)
+			}
+		}()
+	}
+
 	ctrlLogger := getLoggerWithLogLevel(controllerCFG.LogLevel)
 	ctrl.SetLogger(ctrlLogger)
 
diff --git a/pkg/config/controller_config.go b/pkg/config/controller_config.go
index 8b7347b..8b88dd4 100644
--- a/pkg/config/controller_config.go
+++ b/pkg/config/controller_config.go
@@ -11,12 +11,15 @@ const (
 	flagMaxConcurrentReconciles      = "max-concurrent-reconciles"
 	flagEnableConfigMapCheck         = "enable-configmap-check"
 	flagEndpointChunkSize            = "endpoint-chunk-size"
-	defaultLogLevel                  = "info"
-	defaultMaxConcurrentReconciles   = 3
-	defaultEndpointsChunkSize        = 200
-	defaultEnableConfigMapCheck      = true
+	flagEnableProfiling              = "enable-profiling"
 	flagPodUpdateBatchPeriodDuration = "pod-update-batch-period-duration"
-	defaultBatchPeriodDuration       = 1 * time.Second
+
+	defaultLogLevel                = "info"
+	defaultMaxConcurrentReconciles = 3
+	defaultEndpointsChunkSize      = 200
+	defaultEnableConfigMapCheck    = true
+	defaultBatchPeriodDuration     = 1 * time.Second
+	defaultEnableProfiling         = false
 )
 
 // ControllerConfig contains the controller configuration
@@ -33,6 +36,8 @@ type ControllerConfig struct {
 	PodUpdateBatchPeriodDuration time.Duration
 	// Configurations for the Controller Runtime
 	RuntimeConfig RuntimeConfig
+	// To enable profiling
+	EnableProfiling bool
 }
 
 func (cfg *ControllerConfig) BindFlags(fs *pflag.FlagSet) {
@@ -46,5 +51,7 @@ func (cfg *ControllerConfig) BindFlags(fs *pflag.FlagSet) {
 		"Enable checking the configmap for starting the network policy controller")
 	fs.DurationVar(&cfg.PodUpdateBatchPeriodDuration, flagPodUpdateBatchPeriodDuration, defaultBatchPeriodDuration, ""+
 		"Duration between batch updates of pods")
+	fs.BoolVar(&cfg.EnableProfiling, flagEnableProfiling, defaultEnableProfiling,
+		"Enable checking the configmap for starting the network policy controller")
 	cfg.RuntimeConfig.BindFlags(fs)
 }