Merge pull request #171 from aramase/pprof

feat: add an option to enable pprof profiling

README.md

@@ -62,6 +62,8 @@ There are two UNIX domain sockets used by the node-driver-registrar:
 
 * `--mode <mode>` (default: `--mode=registration`): The running mode of node-driver-registrar. `registration` runs node-driver-registrar as a long running process to register the driver with kubelet. `kubelet-registration-probe` runs as a health check and returns a status code of 0 if the driver was registered successfully. In the probe definition make sure that the value of `--kubelet-registration-path` is the same as in the container.
 
+* `--enable-pprof`: Enable pprof profiling on the TCP network address specified by `--http-endpoint`.
+
 ### Required permissions
 
 The node-driver-registrar does not interact with the Kubernetes API, so no RBAC

cmd/csi-node-driver-registrar/main.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	_ "net/http/pprof"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -66,9 +67,10 @@ var (
 	pluginRegistrationPath  = flag.String("plugin-registration-path", "/registration", "Path to Kubernetes plugin registration directory.")
 	kubeletRegistrationPath = flag.String("kubelet-registration-path", "", "Path of the CSI driver socket on the Kubernetes host machine.")
 	healthzPort             = flag.Int("health-port", 0, "(deprecated) TCP port for healthz requests. Set to 0 to disable the healthz server. Only one of `--health-port` and `--http-endpoint` can be set.")
-	httpEndpoint            = flag.String("http-endpoint", "", "The TCP network address where the HTTP server for diagnostics, including the health check indicating whether the registration socket exists, will listen (example: `:8080`). The default is empty string, which means the server is disabled. Only one of `--health-port` and `--http-endpoint` can be set.")
+	httpEndpoint            = flag.String("http-endpoint", "", "The TCP network address where the HTTP server for diagnostics, including pprof and the health check indicating whether the registration socket exists, will listen (example: `:8080`). The default is empty string, which means the server is disabled. Only one of `--health-port` and `--http-endpoint` can be set.")
 	showVersion             = flag.Bool("version", false, "Show version.")
 	mode                    = flag.String("mode", ModeRegistration, `The running mode of node-driver-registrar. "registration" runs node-driver-registrar as a long running process. "kubelet-registration-probe" runs as a health check and returns a status code of 0 if the driver was registered successfully, in the probe definition make sure that the value of --kubelet-registration-path is the same as in the container.`)
+	enableProfile           = flag.Bool("enable-pprof", false, "enable pprof profiling")
 
 	// Set during compilation time
 	version = "unknown"

cmd/csi-node-driver-registrar/node_register.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"net/http/pprof"
 	"os"
 	"os/signal"
 	"runtime"
@@ -68,7 +69,7 @@ func nodeRegister(csiDriverName, httpEndpoint string) {
 	// Registers kubelet plugin watcher api.
 	registerapi.RegisterRegistrationServer(grpcServer, registrar)
 
-	go healthzServer(socketPath, httpEndpoint)
+	go httpServer(socketPath, httpEndpoint)
 	go removeRegSocket(csiDriverName)
 	// Starts service
 	if err := grpcServer.Serve(lis); err != nil {
@@ -86,14 +87,16 @@ func buildSocketPath(csiDriverName string) string {
 	return fmt.Sprintf("%s/%s-reg.sock", *pluginRegistrationPath, csiDriverName)
 }
 
-func healthzServer(socketPath string, httpEndpoint string) {
+func httpServer(socketPath string, httpEndpoint string) {
 	if httpEndpoint == "" {
-		klog.Infof("Skipping healthz server because HTTP endpoint is set to: %q", httpEndpoint)
+		klog.Infof("Skipping HTTP server because endpoint is set to: %q", httpEndpoint)
 		return
 	}
-	klog.Infof("Starting healthz server at HTTP endpoint: %v\n", httpEndpoint)
+	klog.Infof("Starting HTTP server at endpoint: %v\n", httpEndpoint)
 
-	http.HandleFunc("/healthz", func(w http.ResponseWriter, req *http.Request) {
+	// Prepare http endpoint for healthz + profiling (if enabled)
+	mux := http.NewServeMux()
+	mux.HandleFunc("/healthz", func(w http.ResponseWriter, req *http.Request) {
 		socketExists, err := util.DoesSocketExist(socketPath)
 		if err == nil && socketExists {
 			w.WriteHeader(http.StatusOK)
@@ -110,6 +113,16 @@ func healthzServer(socketPath string, httpEndpoint string) {
 		}
 	})
 
+	if *enableProfile {
+		klog.InfoS("Starting profiling", "endpoint", httpEndpoint)
+
+		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)
+	}
+
 	klog.Fatal(http.ListenAndServe(httpEndpoint, nil))
 }