From ae507c88fbd171603b3d07024ec9559405884cf6 Mon Sep 17 00:00:00 2001
From: Praveen Rewar <8457124+praveenrewar@users.noreply.github.com>
Date: Sat, 17 Jun 2023 14:34:07 +0530
Subject: [PATCH] Add profiling flags (based on kubectl)

Signed-off-by: Praveen Rewar <8457124+praveenrewar@users.noreply.github.com>
---
 pkg/kapp/cmd/kapp.go      |  7 +++
 pkg/kapp/cmd/profiling.go | 94 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+)
 create mode 100644 pkg/kapp/cmd/profiling.go

diff --git a/pkg/kapp/cmd/kapp.go b/pkg/kapp/cmd/kapp.go
index 2d9b39c13..ced445651 100644
--- a/pkg/kapp/cmd/kapp.go
+++ b/pkg/kapp/cmd/kapp.go
@@ -31,6 +31,7 @@ type KappOptions struct {
 	KubeAPIFlags    cmdcore.KubeAPIFlags
 	KubeconfigFlags cmdcore.KubeconfigFlags
 	WarningFlags    WarningFlags
+	ProfilingFlags  ProfilingFlags
 }
 
 func NewKappOptions(ui *ui.ConfUI, configFactory cmdcore.ConfigFactory,
@@ -63,6 +64,10 @@ func NewKappCmd(o *KappOptions, flagsFactory cmdcore.FlagsFactory) *cobra.Comman
 		DisableAutoGenTag: true,
 		Version:           version.Version,
 
+		PersistentPostRunE: func(*cobra.Command, []string) error {
+			return o.ProfilingFlags.flushProfiling()
+		},
+
 		// TODO bash completion
 	}
 
@@ -80,6 +85,7 @@ func NewKappCmd(o *KappOptions, flagsFactory cmdcore.FlagsFactory) *cobra.Comman
 	o.KubeAPIFlags.Set(cmd, flagsFactory)
 	o.KubeconfigFlags.Set(cmd, flagsFactory)
 	o.WarningFlags.Set(cmd, flagsFactory)
+	o.ProfilingFlags.Set(cmd, flagsFactory)
 
 	o.configFactory.ConfigurePathResolver(o.KubeconfigFlags.Path.Value)
 	o.configFactory.ConfigureContextResolver(o.KubeconfigFlags.Context.Value)
@@ -135,6 +141,7 @@ func NewKappCmd(o *KappOptions, flagsFactory cmdcore.FlagsFactory) *cobra.Comman
 		o.LoggerFlags.Configure(o.logger)
 		o.KubeAPIFlags.Configure(o.configFactory)
 		o.WarningFlags.Configure(o.depsFactory)
+		o.ProfilingFlags.initProfiling()
 		return nil
 	})
 
diff --git a/pkg/kapp/cmd/profiling.go b/pkg/kapp/cmd/profiling.go
new file mode 100644
index 000000000..1b14da1f3
--- /dev/null
+++ b/pkg/kapp/cmd/profiling.go
@@ -0,0 +1,94 @@
+// Copyright 2023 VMware, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package cmd
+
+import (
+	"fmt"
+	"os"
+	"os/signal"
+	"runtime"
+	"runtime/pprof"
+
+	"github.com/spf13/cobra"
+	cmdcore "github.com/vmware-tanzu/carvel-kapp/pkg/kapp/cmd/core"
+)
+
+type ProfilingFlags struct {
+	profileName   string
+	profileOutput string
+}
+
+func (p *ProfilingFlags) Set(cmd *cobra.Command, _ cmdcore.FlagsFactory) {
+	cmd.PersistentFlags().StringVar(&p.profileName, "profile", "none", "Name of profile to capture. One of (none|cpu|heap|goroutine|threadcreate|block|mutex)")
+	cmd.PersistentFlags().StringVar(&p.profileOutput, "profile-output", "profile.pprof", "Name of the file to write the profile to")
+}
+
+func (p *ProfilingFlags) initProfiling() error {
+	var (
+		f   *os.File
+		err error
+	)
+	switch p.profileName {
+	case "none":
+		return nil
+	case "cpu":
+		f, err = os.Create(p.profileOutput)
+		if err != nil {
+			return err
+		}
+		err = pprof.StartCPUProfile(f)
+		if err != nil {
+			return err
+		}
+	// 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)
+	case "mutex":
+		runtime.SetMutexProfileFraction(1)
+	default:
+		// Check the profile name is valid.
+		if profile := pprof.Lookup(p.profileName); profile == nil {
+			return fmt.Errorf("unknown profile '%s'", p.profileName)
+		}
+	}
+
+	// If the command is interrupted before the end (ctrl-c), flush the
+	// profiling files
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt)
+	go func() {
+		<-c
+		f.Close()
+		p.flushProfiling()
+		os.Exit(0)
+	}()
+
+	return nil
+}
+
+func (p *ProfilingFlags) flushProfiling() error {
+	switch p.profileName {
+	case "none":
+		return nil
+	case "cpu":
+		pprof.StopCPUProfile()
+	case "heap":
+		runtime.GC()
+		fallthrough
+	default:
+		profile := pprof.Lookup(p.profileName)
+		if profile == nil {
+			return nil
+		}
+		f, err := os.Create(p.profileOutput)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+		return profile.WriteTo(f, 0)
+	}
+
+	return nil
+}