Move pprof paths to HELM_PPROF env variable

Signed-off-by: Evans Mungai <mbuevans@gmail.com>

CONTRIBUTING.md

@@ -280,11 +280,11 @@ Small, ad-hoc changes/PRs to Helm which introduce user facing changes, which wou
 
 ### Profiling PRs
 
-If your contribution requires profiling to check memory and/or CPU usage, you can use `--cpuprofile` and/or `--memprofile` global cli flags to collect runtime profiling data for analysis. You can use Golang's [pprof](https://github.com/google/pprof/blob/main/doc/README.md) tool to inspect the results.
+If your contribution requires profiling to check memory and/or CPU usage, you can set `HELM_PPROF=cpu=/path/to/cpu.prof,mem=/path/to/mem.prof` environment variable to collect runtime profiling data for analysis. You can use Golang's [pprof](https://github.com/google/pprof/blob/main/doc/README.md) tool to inspect the results.
 
 Example analysing collected profiling data
 ```
-helm show all bitnami/nginx --memprofile mem.prof --cpuprofile cpu.prof
+HELM_PPROF=cpu=/path/to/cpu.prof,mem=/path/to/mem.prof helm show all bitnami/nginx
 
 # Visualize graphs. You need to have installed graphviz package in your system
 go tool pprof -http=":8000" cpu.prof

cmd/helm/profiling.go

@@ -19,6 +19,8 @@ package main
 import (
 	"fmt"
 	"os"
+	"path"
+	"path/filepath"
 	"runtime"
 	"runtime/pprof"
 	"strings"
@@ -28,11 +30,17 @@ import (
 
 var (
 	cpuProfileFile *os.File
+	pprofPaths     map[string]string
 )
 
+func init() {
+	pprofPaths = parsePProfPaths(os.Getenv("HELM_PPROF"))
+}
+
 // startProfiling starts profiling CPU usage
-func startProfiling(cpuprofile string) error {
-	if cpuprofile != "" {
+func startProfiling() error {
+	cpuprofile, ok := pprofPaths["cpu"]
+	if ok && cpuprofile != "" {
 		var err error
 		cpuProfileFile, err = os.Create(cpuprofile)
 		if err != nil {
@@ -48,8 +56,8 @@ func startProfiling(cpuprofile string) error {
 }
 
 // stopProfiling stops profiling CPU and memory usage and writes the results to
-// the files specified by --cpuprofile and --memprofile flags respectively.
-func stopProfiling(memprofile string) error {
+// the files specified by HELM_PPROF=cpu=/path/to/cpu.prof,mem=/path/to/mem.prof
+func stopProfiling() error {
 	errs := []string{}
 
 	// Stop CPU profiling if it was started
@@ -62,7 +70,8 @@ func stopProfiling(memprofile string) error {
 		cpuProfileFile = nil
 	}
 
-	if memprofile != "" {
+	memprofile, ok := pprofPaths["mem"]
+	if ok && memprofile != "" {
 		f, err := os.Create(memprofile)
 		if err != nil {
 			errs = append(errs, err.Error())
@@ -88,3 +97,28 @@ func addProfilingFlags(cmd *cobra.Command) {
 	cmd.PersistentFlags().String("cpuprofile", "", "File path to write cpu profiling data")
 	cmd.PersistentFlags().String("memprofile", "", "File path to write memory profiling data")
 }
+
+func parsePProfPaths(env string) map[string]string {
+	// Initial empty paths
+	m := map[string]string{}
+	for _, pprofs := range strings.Split(env, ",") {
+		// Is of the format mem=/path/to/memprof
+		tuple := strings.Split(pprofs, "=")
+		if len(tuple) != 2 {
+			continue
+		}
+		if tuple[0] != "cpu" && tuple[0] != "mem" {
+			continue
+		}
+
+		s, err := filepath.Abs(path.Clean(tuple[1]))
+		if err != nil {
+			continue
+		}
+		if !strings.HasSuffix(s, string(filepath.Separator)) {
+			// Ensure its not a directory
+			m[tuple[0]] = s
+		}
+	}
+	return m
+}

cmd/helm/profiling_test.go

@@ -0,0 +1,79 @@
+/*
+Copyright The Helm Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package main
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_parsePProfPaths(t *testing.T) {
+	cwd, err := os.Getwd()
+	require.NoError(t, err)
+
+	tests := []struct {
+		name string
+		env  string
+		want map[string]string
+	}{
+		{
+			name: "no env",
+			env:  "",
+			want: map[string]string{},
+		},
+		{
+			name: "single path",
+			env:  "cpu=cpu.pprof",
+			want: map[string]string{
+				"cpu": cwd + "/cpu.pprof",
+			},
+		},
+		{
+			name: "mem and cpu paths",
+			env:  "cpu=cpu.pprof,mem=mem.pprof",
+			want: map[string]string{
+				"cpu": cwd + "/cpu.pprof",
+				"mem": cwd + "/mem.pprof",
+			},
+		},
+		{
+			name: "extra commas",
+			env:  "cpu=cpu.pprof,mem=mem.pprof,",
+			want: map[string]string{
+				"cpu": cwd + "/cpu.pprof",
+				"mem": cwd + "/mem.pprof",
+			},
+		},
+		{
+			name: "unknown keys",
+			env:  "cpu=cpu.pprof,mem=mem.pprof,foo=bar",
+			want: map[string]string{
+				"cpu": cwd + "/cpu.pprof",
+				"mem": cwd + "/mem.pprof",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := parsePProfPaths(tt.env)
+			assert.Equalf(t, tt.want, got, "parsePProfPaths() = %v, want %v", got, tt.want)
+		})
+	}
+}

cmd/helm/root.go

@@ -96,22 +96,12 @@ func newRootCmd(actionConfig *action.Configuration, out io.Writer, args []string
 		Long:         globalUsage,
 		SilenceUsage: true,
 		PersistentPreRun: func(cmd *cobra.Command, _ []string) {
-			cpuprof, err := cmd.Flags().GetString("cpuprofile")
-			if err != nil {
-				log.Printf("Warning: Failed to get cpuprofile flag: %v", err)
-			}
-
-			if err := startProfiling(cpuprof); err != nil {
+			if err := startProfiling(); err != nil {
 				log.Printf("Warning: Failed to start profiling: %v", err)
 			}
 		},
 		PersistentPostRun: func(cmd *cobra.Command, _ []string) {
-			memprof, err := cmd.Flags().GetString("memprofile")
-			if err != nil {
-				log.Printf("Warning: Failed to get memprofile flag: %v", err)
-			}
-
-			if err := stopProfiling(memprof); err != nil {
+			if err := stopProfiling(); err != nil {
 				log.Printf("Warning: Failed to stop profiling: %v", err)
 			}
 		},