From e3289023d36968f6c106cbb03d8c1fa306678230 Mon Sep 17 00:00:00 2001
From: Lyndon-Li <lyonghui@vmware.com>
Date: Wed, 4 Sep 2024 18:46:01 +0800
Subject: [PATCH] add profile to node-agent

Signed-off-by: Lyndon-Li <lyonghui@vmware.com>
---
 pkg/cmd/cli/nodeagent/server.go | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/pkg/cmd/cli/nodeagent/server.go b/pkg/cmd/cli/nodeagent/server.go
index 16543794d3..074a765ef7 100644
--- a/pkg/cmd/cli/nodeagent/server.go
+++ b/pkg/cmd/cli/nodeagent/server.go
@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"math"
 	"net/http"
+	"net/http/pprof"
 	"os"
 	"strings"
 	"time"
@@ -84,6 +85,7 @@ type nodeAgentServerConfig struct {
 	metricsAddress          string
 	resourceTimeout         time.Duration
 	dataMoverPrepareTimeout time.Duration
+	profilerAddress         string
 }
 
 func NewServerCommand(f client.Factory) *cobra.Command {
@@ -120,6 +122,7 @@ func NewServerCommand(f client.Factory) *cobra.Command {
 	command.Flags().DurationVar(&config.resourceTimeout, "resource-timeout", config.resourceTimeout, "How long to wait for resource processes which are not covered by other specific timeout parameters. Default is 10 minutes.")
 	command.Flags().DurationVar(&config.dataMoverPrepareTimeout, "data-mover-prepare-timeout", config.dataMoverPrepareTimeout, "How long to wait for preparing a DataUpload/DataDownload. Default is 30 minutes.")
 	command.Flags().StringVar(&config.metricsAddress, "metrics-address", config.metricsAddress, "The address to expose prometheus metrics")
+	command.Flags().StringVar(&config.profilerAddress, "profiler-address", config.profilerAddress, "The address to expose the pprof profiler.")
 
 	return command
 }
@@ -236,6 +239,10 @@ func newNodeAgentServer(logger logrus.FieldLogger, factory client.Factory, confi
 func (s *nodeAgentServer) run() {
 	signals.CancelOnShutdown(s.cancelFunc, s.logger)
 
+	if s.config.profilerAddress != "" {
+		go s.runProfiler()
+	}
+
 	go func() {
 		metricsMux := http.NewServeMux()
 		metricsMux.Handle("/metrics", promhttp.Handler())
@@ -308,6 +315,24 @@ func (s *nodeAgentServer) run() {
 	}
 }
 
+func (s *nodeAgentServer) runProfiler() {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/debug/pprof/", pprof.Index)
+	mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+
+	server := &http.Server{
+		Addr:              s.config.profilerAddress,
+		Handler:           mux,
+		ReadHeaderTimeout: 3 * time.Second,
+	}
+	if err := server.ListenAndServe(); err != nil {
+		s.logger.WithError(errors.WithStack(err)).Error("error running profiler http server")
+	}
+}
+
 // validatePodVolumesHostPath validates that the pod volumes path contains a
 // directory for each Pod running on this node
 func (s *nodeAgentServer) validatePodVolumesHostPath(client kubernetes.Interface) error {