add --pprof option

cmd/ccat/main.go

@@ -51,6 +51,7 @@ var (
 	argInsecure     = flag.BoolP("insecure", "k", false, "get files insecurely (globally)")
 	argCompletion   = flag.StringP("completion", "C", "", "print shell completion script")
 	argLess         = flag.BoolP("ui", "T", false, "display with a minimal ui")
+	argPprof        = flag.Bool("pprof", false, "enable cpu and mem profiling")
 
 	argSetKey = new(bool) // set by init()
 
@@ -84,10 +85,15 @@ func init() {
 func main() {
 	log.Debugln("STARTING ccat")
 	log.Debugf(buildLine())
+	if *argPprof {
+		log.Debugln("pprof enabled")
+		enablePprof()
+		defer endPprof()
+	}
 
 	if *argSetKey {
 		interactivelySetKey()
-		os.Exit(0)
+		os.Exit(0) //nolint
 	}
 
 	if *argVersion {

cmd/ccat/pprof.go

@@ -0,0 +1,46 @@
+package main
+
+import (
+	"os"
+	"runtime"
+	"runtime/pprof"
+
+	"github.com/batmac/ccat/pkg/log"
+	// _ "net/http/pprof"
+)
+
+var (
+	cpuPprofFile *os.File
+	memPprofFile *os.File
+	err          error
+)
+
+func enablePprof() {
+	log.Debugln("enabling pprof")
+
+	cpuPprofFile, err = os.Create("cpu.pprof")
+	if err != nil {
+		log.Fatal("could not create CPU profile: ", err)
+	}
+	if err := pprof.StartCPUProfile(cpuPprofFile); err != nil {
+		log.Fatal("could not start CPU profile: ", err)
+	}
+
+	memPprofFile, err = os.Create("mem.pprof")
+	if err != nil {
+		log.Fatal("could not create memory profile: ", err)
+	}
+}
+
+func endPprof() {
+	log.Debugln("ending pprof")
+
+	pprof.StopCPUProfile()
+	_ = cpuPprofFile.Close()
+
+	runtime.GC() // get up-to-date statistics
+	if err := pprof.WriteHeapProfile(memPprofFile); err != nil {
+		log.Fatal("could not write memory profile: ", err)
+	}
+	_ = memPprofFile.Close()
+}