feat(controller): enable pprof profiling

Signed-off-by: John Wood <[email protected]>

cmd/rollouts-controller/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"net/http"
 	"os"
 	"strings"
 	"time"
@@ -79,6 +80,8 @@ func newCommand() *cobra.Command {
 		printVersion                   bool
 		selfServiceNotificationEnabled bool
 		controllersEnabled             []string
+		pprofEnabled                   bool
+		pprofPort                      int
 	)
 	electOpts := controller.NewLeaderElectionOptions()
 	var command = cobra.Command{
@@ -204,6 +207,11 @@ func newCommand() *cobra.Command {
 			ingressWrapper, err := ingressutil.NewIngressWrapper(mode, kubeClient, kubeInformerFactory)
 			checkError(err)
 
+			if pprofEnabled {
+				mux := controller.NewPProfServer()
+				go func() { log.Println(http.ListenAndServe(fmt.Sprintf("localhost:%d", pprofPort), mux)) }()
+			}
+
 			var cm *controller.Manager
 
 			enabledControllers, err := getEnabledControllers(controllersEnabled)
@@ -283,6 +291,7 @@ func newCommand() *cobra.Command {
 	command.Flags().IntVar(&klogLevel, "kloglevel", 0, "Set the klog logging level")
 	command.Flags().IntVar(&metricsPort, "metricsport", controller.DefaultMetricsPort, "Set the port the metrics endpoint should be exposed over")
 	command.Flags().IntVar(&healthzPort, "healthzPort", controller.DefaultHealthzPort, "Set the port the healthz endpoint should be exposed over")
+	command.Flags().IntVar(&pprofPort, "pprof-port", controller.DefaultPProfPort, "Set the port the pprof endpoint should be exposed over")
 	command.Flags().StringVar(&instanceID, "instance-id", "", "Indicates which argo rollout objects the controller should operate on")
 	command.Flags().Float32Var(&qps, "qps", defaults.DefaultQPS, "Maximum QPS (queries per second) to the K8s API server")
 	command.Flags().IntVar(&burst, "burst", defaults.DefaultBurst, "Maximum burst for throttle.")
@@ -310,6 +319,7 @@ func newCommand() *cobra.Command {
 	command.Flags().DurationVar(&electOpts.LeaderElectionRetryPeriod, "leader-election-retry-period", controller.DefaultLeaderElectionRetryPeriod, "The duration the clients should wait between attempting acquisition and renewal of a leadership. This is only applicable if leader election is enabled.")
 	command.Flags().BoolVar(&selfServiceNotificationEnabled, "self-service-notification-enabled", false, "Allows rollouts controller to pull notification config from the namespace that the rollout resource is in. This is useful for self-service notification.")
 	command.Flags().StringSliceVar(&controllersEnabled, "controllers", nil, "Explicitly specify the list of controllers to run, currently only supports 'analysis', eg. --controller=analysis. Default: all controllers are enabled")
+	command.Flags().BoolVar(&pprofEnabled, "enable-pprof", false, "Enable pprof profiling on controller. Default endpoint exposed via :6060/debug/pprof. Change the port with the --pprof-port flag.")
 	return &command
 }
 

controller/controller.go

@@ -69,6 +69,9 @@ const (
 	// DefaultMetricsPort is the default port to expose the metrics endpoint
 	DefaultMetricsPort = 8090
 
+	// DefaultPProfPort is the default port to expose the pprof endpoint
+	DefaultPProfPort = 6060
+
 	// DefaultRolloutThreads is the default number of rollout worker threads to start with the controller
 	DefaultRolloutThreads = 10
 

controller/profiling.go

@@ -0,0 +1,24 @@
+package controller
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/pprof"
+)
+
+const (
+	ProfilingPath = "/debug/pprof"
+)
+
+// NewPProfServer returns a new pprof server to gather to gather runtime profiling data
+func NewPProfServer() *http.ServeMux {
+	mux := http.NewServeMux()
+
+	mux.HandleFunc(ProfilingPath, pprof.Index)
+	mux.HandleFunc(fmt.Sprintf("%s/cmdline", ProfilingPath), pprof.Cmdline)
+	mux.HandleFunc(fmt.Sprintf("%s/profile", ProfilingPath), pprof.Profile)
+	mux.HandleFunc(fmt.Sprintf("%s/symbol", ProfilingPath), pprof.Symbol)
+	mux.HandleFunc(fmt.Sprintf("%s/trace", ProfilingPath), pprof.Trace)
+
+	return mux
+}