From 5d634e7db6d8002c6546033441cb4bcd9b73aba5 Mon Sep 17 00:00:00 2001
From: Damien Lespiau <damien@weave.works>
Date: Fri, 14 Sep 2018 15:20:04 +0100
Subject: [PATCH] Add go profile instrumentation to kubectl

This commit adds two new global options to kubectl: --profile and
--profile-output, writing out go profiles to disk to debug interesting and
unexpected kubectl behaviour.

As an example, here is how to capture a block file, eg. for how long are we
blocked on I/O and where?

$ kubectl get nodes --profile=mutex -v6
$ go tool pprof -png ./profile.pprof > out.png
$ google-chrome out.png

Fixes: #68679
---
 pkg/kubectl/cmd/BUILD        |  1 +
 pkg/kubectl/cmd/cmd.go       | 10 ++++
 pkg/kubectl/cmd/profiling.go | 88 ++++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+)
 create mode 100644 pkg/kubectl/cmd/profiling.go

diff --git a/pkg/kubectl/cmd/BUILD b/pkg/kubectl/cmd/BUILD
index 885485458d..b2195fbee8 100644
--- a/pkg/kubectl/cmd/BUILD
+++ b/pkg/kubectl/cmd/BUILD
@@ -40,6 +40,7 @@ go_library(
         "patch.go",
         "plugin.go",
         "portforward.go",
+        "profiling.go",
         "proxy.go",
         "replace.go",
         "rollingupdate.go",
diff --git a/pkg/kubectl/cmd/cmd.go b/pkg/kubectl/cmd/cmd.go
index b15dcc53b1..158e3eccbe 100644
--- a/pkg/kubectl/cmd/cmd.go
+++ b/pkg/kubectl/cmd/cmd.go
@@ -370,6 +370,14 @@ func NewKubectlCommand(in io.Reader, out, err io.Writer) *cobra.Command {
       Find more information at:
             https://kubernetes.io/docs/reference/kubectl/overview/`),
 		Run: runHelp,
+		// Hook before and after Run initialize and write profiles to disk,
+		// respectively.
+		PersistentPreRunE: func(*cobra.Command, []string) error {
+			return initProfiling()
+		},
+		PersistentPostRunE: func(*cobra.Command, []string) error {
+			return flushProfiling()
+		},
 		BashCompletionFunction: bashCompletionFunc,
 	}
 
@@ -380,6 +388,8 @@ func NewKubectlCommand(in io.Reader, out, err io.Writer) *cobra.Command {
 	// a.k.a. change all "_" to "-". e.g. glog package
 	flags.SetNormalizeFunc(utilflag.WordSepNormalizeFunc)
 
+	addProfilingFlags(flags)
+
 	kubeConfigFlags := genericclioptions.NewConfigFlags()
 	kubeConfigFlags.AddFlags(flags)
 	matchVersionKubeConfigFlags := cmdutil.NewMatchVersionFlags(kubeConfigFlags)
diff --git a/pkg/kubectl/cmd/profiling.go b/pkg/kubectl/cmd/profiling.go
new file mode 100644
index 0000000000..2a1c1ce3c1
--- /dev/null
+++ b/pkg/kubectl/cmd/profiling.go
@@ -0,0 +1,88 @@
+/*
+Copyright 2017 The Kubernetes 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 cmd
+
+import (
+	"fmt"
+	"os"
+	"runtime"
+	"runtime/pprof"
+
+	"github.com/spf13/pflag"
+)
+
+var (
+	profileName   string
+	profileOutput string
+)
+
+func addProfilingFlags(flags *pflag.FlagSet) {
+	flags.StringVar(&profileName, "profile", "none", "Name of profile to capture. One of (none|cpu|heap|goroutine|threadcreate|block|mutex)")
+	flags.StringVar(&profileOutput, "profile-output", "profile.pprof", "Name of the file to write the profile to")
+}
+
+func initProfiling() error {
+	switch profileName {
+	case "none":
+		return nil
+	case "cpu":
+		f, err := os.Create(profileOutput)
+		if err != nil {
+			return err
+		}
+		return pprof.StartCPUProfile(f)
+	// Block and mutex profiles need a call to Set{Block,Mutex}ProfileRate to
+	// output anything. We choose to sample all events.
+	case "block":
+		runtime.SetBlockProfileRate(1)
+		return nil
+	case "mutex":
+		runtime.SetMutexProfileFraction(1)
+		return nil
+	default:
+		// Check the profile name is valid.
+		if profile := pprof.Lookup(profileName); profile == nil {
+			return fmt.Errorf("unknown profile '%s'", profileName)
+		}
+	}
+
+	return nil
+}
+
+func flushProfiling() error {
+	switch profileName {
+	case "none":
+		return nil
+	case "cpu":
+		pprof.StopCPUProfile()
+	case "heap":
+		runtime.GC()
+		fallthrough
+	default:
+		profile := pprof.Lookup(profileName)
+		if profile == nil {
+			return nil
+		}
+		f, err := os.Create(profileOutput)
+		if err != nil {
+			return err
+		}
+		profile.WriteTo(f, 0)
+	}
+
+	return nil
+}