From 21b97c0d3bd4440131c675a832901f4be3ad57f5 Mon Sep 17 00:00:00 2001
From: Baptiste Girard-Carrabin <baptiste.girardcarrabin@datadoghq.com>
Date: Thu, 13 Jul 2023 11:48:03 +0200
Subject: [PATCH] metrics: add pprof support

Create new function `RegisterPprofToServer` to enable pprof profiling on the existing metrics server.
This can be used to gain useful insights on the performance of the application.
---
 metrics/metrics.go      | 21 +++++++++++++++
 metrics/metrics_test.go | 59 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/metrics/metrics.go b/metrics/metrics.go
index b641fffcb..33b64c7cd 100644
--- a/metrics/metrics.go
+++ b/metrics/metrics.go
@@ -20,6 +20,7 @@ import (
 	"bufio"
 	"fmt"
 	"net/http"
+	"net/http/pprof"
 	"sort"
 	"strings"
 	"time"
@@ -99,6 +100,12 @@ type CSIMetricsManager interface {
 	// RegisterToServer registers an HTTP handler for this metrics manager to the
 	// given server at the specified address/path.
 	RegisterToServer(s Server, metricsPath string)
+
+	// RegisterPprofToServer registers the HTTP handlers necessary to enable pprof
+	// for this metrics manager to the given server at the usual path.
+	// This function is not needed when using DefaultServeMux as the Server since
+	// the handlers will automatically be registered when importing pprof.
+	RegisterPprofToServer(s Server)
 }
 
 // Server represents any type that could serve HTTP requests for the metrics
@@ -388,6 +395,20 @@ func (cmm *csiMetricsManager) RegisterToServer(s Server, metricsPath string) {
 			ErrorHandling: metrics.ContinueOnError}))
 }
 
+// RegisterPprofToServer registers the HTTP handlers necessary to enable pprof
+// for this metrics manager to the given server at the usual path.
+// This function is not needed when using DefaultServeMux as the Server since
+// the handlers will automatically be registered when importing pprof.
+func (cmm *csiMetricsManager) RegisterPprofToServer(s Server) {
+	// Needed handlers can be seen here:
+	// https://github.com/golang/go/blob/master/src/net/http/pprof/pprof.go#L27
+	s.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
+	s.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
+	s.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
+	s.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
+	s.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
+}
+
 // VerifyMetricsMatch is a helper function that verifies that the expected and
 // actual metrics are identical excluding metricToIgnore.
 // This method is only used by tests. Ideally it should be in the _test file,
diff --git a/metrics/metrics_test.go b/metrics/metrics_test.go
index d6e6d4b0f..26ca60722 100644
--- a/metrics/metrics_test.go
+++ b/metrics/metrics_test.go
@@ -537,6 +537,65 @@ func TestRegisterToServer_Noop(t *testing.T) {
 	}
 }
 
+func TestRegisterPprofToServer_AllEndpointsAvailable(t *testing.T) {
+	endpoints := []string{
+		"/debug/pprof/",
+		"/debug/pprof/cmdline",
+		"/debug/pprof/profile",
+		"/debug/pprof/symbol",
+		"/debug/pprof/trace",
+	}
+
+	for _, endpoint := range endpoints {
+		t.Run(endpoint, func(t *testing.T) {
+			testRegisterPprofToServer_AllEndpointsAvailable(t, endpoint)
+		})
+	}
+}
+
+func testRegisterPprofToServer_AllEndpointsAvailable(t *testing.T, endpoint string) {
+	// Arrange
+	cmm := NewCSIMetricsManagerForSidecar(
+		"fake.csi.driver.io" /* driverName */)
+	mux := http.NewServeMux()
+
+	// Act
+	cmm.RegisterPprofToServer(mux)
+
+	// Assert
+	request := httptest.NewRequest("GET", endpoint, strings.NewReader(""))
+	rec := httptest.NewRecorder()
+	mux.ServeHTTP(rec, request)
+	resp := rec.Result()
+
+	if resp.StatusCode != 200 {
+		t.Fatalf("%s response status not 200. Response was: %+v", endpoint, resp)
+	}
+
+	contentBytes, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatalf("Failed to parse pprof index response.  Response was: %+v Error: %v", resp, err)
+	}
+
+	// Other endpoints return binary data
+	if endpoint == "/debug/pprof/" {
+		actualPprofIndex := string(contentBytes)
+
+		// This is the exepcted index html page if pprof is running
+		expectedPprofIndexSubstr := `<body>
+/debug/pprof/
+<br>
+<p>Set debug=1 as a query parameter to export in legacy text format</p>
+<br>
+Types of profiles available:
+<table>`
+
+		if ok := strings.Contains(actualPprofIndex, expectedPprofIndexSubstr); !ok {
+			t.Fatalf("Pprof index returned by end point do not match expectation. Expected: %s \nGot: %s", expectedPprofIndexSubstr, actualPprofIndex)
+		}
+	}
+}
+
 func TestProcessStartTimeMetricExist(t *testing.T) {
 	// Arrange
 	cmm := NewCSIMetricsManagerForSidecar(