cmd/receivepack9418: Profile

2 files changed, 98 insertions, 2 deletions

diff --git a/cmd/receivepack9418/main.go b/cmd/receivepack9418/main.go
index 377bf7d4..6884f326 100644
--- a/cmd/receivepack9418/main.go
+++ b/cmd/receivepack9418/main.go
@@ -4,20 +4,58 @@ package main
 import (
 	"flag"
 	"log"
+	"os"
 )
 
 func main() {
+	os.Exit(runMain())
+}
+
+func runMain() int {
 	listenAddr := flag.String("listen", ":9418", "listen address")
 	repoPath := flag.String("repo", "", "path to git dir (.git or bare repo root)")
+	cpuProfilePath := flag.String("cpuprofile", "", "write CPU profile to file")
+	memProfilePath := flag.String("memprofile", "", "write heap profile to file at exit")
 
 	flag.Parse()
 
 	if *repoPath == "" {
-		log.Fatal("must provide -repo <path-to-git-dir>")
+		log.Print("must provide -repo <path-to-git-dir>")
+
+		return 2
+	}
+
+	if *cpuProfilePath != "" {
+		stopCPUProfile, err := startCPUProfile(*cpuProfilePath)
+		if err != nil {
+			log.Printf("cpuprofile: %v", err)
+
+			return 1
+		}
+
+		defer func() {
+			stopErr := stopCPUProfile()
+			if stopErr != nil {
+				log.Printf("cpuprofile: %v", stopErr)
+			}
+		}()
+	}
+
+	if *memProfilePath != "" {
+		defer func() {
+			memErr := writeMemProfile(*memProfilePath)
+			if memErr != nil {
+				log.Printf("memprofile: %v", memErr)
+			}
+		}()
 	}
 
 	err := run(*listenAddr, *repoPath)
 	if err != nil {
-		log.Fatalf("run: %v", err)
+		log.Printf("run: %v", err)
+
+		return 1
 	}
+
+	return 0
 }
diff --git a/cmd/receivepack9418/profile.go b/cmd/receivepack9418/profile.go
new file mode 100644
index 00000000..40ba1d56
--- /dev/null
+++ b/cmd/receivepack9418/profile.go
@@ -0,0 +1,58 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"runtime"
+	"runtime/pprof"
+)
+
+func startCPUProfile(path string) (func() error, error) {
+	//#nosec G304
+	file, err := os.Create(path)
+	if err != nil {
+		return nil, fmt.Errorf("create %q: %w", path, err)
+	}
+
+	err = pprof.StartCPUProfile(file)
+	if err != nil {
+		_ = file.Close()
+
+		return nil, fmt.Errorf("start cpu profile %q: %w", path, err)
+	}
+
+	return func() error {
+		pprof.StopCPUProfile()
+
+		err := file.Close()
+		if err != nil {
+			return fmt.Errorf("close cpu profile %q: %w", path, err)
+		}
+
+		return nil
+	}, nil
+}
+
+func writeMemProfile(path string) error {
+	//#nosec G304
+	file, err := os.Create(path)
+	if err != nil {
+		return fmt.Errorf("create %q: %w", path, err)
+	}
+
+	runtime.GC()
+
+	err = pprof.WriteHeapProfile(file)
+	if err != nil {
+		_ = file.Close()
+
+		return fmt.Errorf("write heap profile %q: %w", path, err)
+	}
+
+	err = file.Close()
+	if err != nil {
+		return fmt.Errorf("close heap profile %q: %w", path, err)
+	}
+
+	return nil
+}