From 76665bed7327eeea2e9075a3a24842879fd9d300 Mon Sep 17 00:00:00 2001
From: Hau Nguyen <hauvipapro@gmail.com>
Date: Tue, 17 Jan 2023 18:10:07 +0700
Subject: [PATCH] feat: add go profiler

---
 internal/cli/action.go | 33 ++++++++++++++++++++++++++++++++-
 internal/cli/app.go    | 17 ++++++++++++-----
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/internal/cli/action.go b/internal/cli/action.go
index 1b47d9a..5b91a6d 100644
--- a/internal/cli/action.go
+++ b/internal/cli/action.go
@@ -2,6 +2,9 @@ package cli
 
 import (
 	"fmt"
+	"os"
+	"runtime"
+	"runtime/pprof"
 
 	"github.com/urfave/cli/v2"
 
@@ -15,6 +18,7 @@ type action struct {
 		write         bool
 		diff          bool
 		verbose       bool
+		profiler      bool
 	}
 }
 
@@ -23,11 +27,12 @@ func (a *action) RunHelp(c *cli.Context) error {
 }
 
 func (a *action) getFlags(c *cli.Context) {
+	a.flags.companyPrefix = c.String(flagCompanyPrefixName)
 	a.flags.list = c.Bool(flagListName)
 	a.flags.write = c.Bool(flagWriteName)
 	a.flags.diff = c.Bool(flagDiffName)
 	a.flags.verbose = c.Bool(flagVerboseName)
-	a.flags.companyPrefix = c.String(flagCompanyPrefixName)
+	a.flags.profiler = c.Bool(flagProfilerName)
 
 	if a.flags.verbose {
 		fmt.Printf("flags: %+v\n", a.flags)
@@ -49,6 +54,19 @@ func (a *action) Run(c *cli.Context) error {
 		return a.RunHelp(c)
 	}
 
+	if a.flags.profiler {
+		f, err := os.Create("cpuprofile")
+		if err != nil {
+			return fmt.Errorf("os: failed to create: %w", err)
+		}
+		defer f.Close()
+
+		if err := pprof.StartCPUProfile(f); err != nil {
+			return fmt.Errorf("pprof: failed to start cpu profile: %w", err)
+		}
+		defer pprof.StopCPUProfile()
+	}
+
 	ft, err := imports.NewFormmater(
 		imports.FormatterWithList(a.flags.list),
 		imports.FormatterWithWrite(a.flags.write),
@@ -65,5 +83,18 @@ func (a *action) Run(c *cli.Context) error {
 		return fmt.Errorf("imports formatter: failed to format %v: %w", args, err)
 	}
 
+	if a.flags.profiler {
+		f, err := os.Create("memprofile")
+		if err != nil {
+			return fmt.Errorf("os: failed to create: %w", err)
+		}
+		defer f.Close()
+
+		runtime.GC()
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			return fmt.Errorf("pprof: failed to write heap profile: %w", err)
+		}
+	}
+
 	return nil
 }
diff --git a/internal/cli/app.go b/internal/cli/app.go
index 3490444..427436f 100644
--- a/internal/cli/app.go
+++ b/internal/cli/app.go
@@ -13,6 +13,9 @@ const (
 	usage = "goimports with my opinionated preferences"
 
 	// Inspiration from gofmt flags
+	flagCompanyPrefixName  = "company"
+	flagCompanyPrefixUsage = "company prefix, for example github.com/haunt98"
+
 	flagListName  = "list"
 	flagListUsage = "list files will be changed"
 
@@ -25,8 +28,8 @@ const (
 	flagVerboseName  = "verbose"
 	flagVerboseUsage = "show verbose output, for debug only"
 
-	flagCompanyPrefixName  = "company"
-	flagCompanyPrefixUsage = "company prefix, for example github.com/haunt98"
+	flagProfilerName  = "profiler"
+	flagProfilerUsage = "go profiler, for debug only"
 )
 
 var (
@@ -47,6 +50,10 @@ func NewApp() *App {
 		Name:  name,
 		Usage: usage,
 		Flags: []cli.Flag{
+			&cli.StringFlag{
+				Name:  flagCompanyPrefixName,
+				Usage: flagCompanyPrefixUsage,
+			},
 			&cli.BoolFlag{
 				Name:    flagListName,
 				Usage:   flagListUsage,
@@ -66,9 +73,9 @@ func NewApp() *App {
 				Name:  flagVerboseName,
 				Usage: flagVerboseUsage,
 			},
-			&cli.StringFlag{
-				Name:  flagCompanyPrefixName,
-				Usage: flagCompanyPrefixUsage,
+			&cli.BoolFlag{
+				Name:  flagProfilerName,
+				Usage: flagProfilerUsage,
 			},
 		},
 		Action: a.Run,