Add go profile instrumentation to kubectl

pkg/kubectl/cmd/BUILD

@@ -40,6 +40,7 @@ go_library(
         "patch.go",
         "plugin.go",
         "portforward.go",
+        "profiling.go",
         "proxy.go",
         "replace.go",
         "rollingupdate.go",

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)

pkg/kubectl/cmd/profiling.go

+/*
+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
+}