From 7331d3a7ee83e1a8c733af2f4c63201cd2f8d01b Mon Sep 17 00:00:00 2001
From: James Shubin <james@shubin.ca>
Date: Fri, 31 Jan 2025 00:00:05 -0500
Subject: [PATCH] util: pprof: Improve pprof for ease of use

This makes it slightly cleaner.
---
 main.go             | 20 ++++++++++----------
 util/pprof/pprof.go |  9 ++++++---
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/main.go b/main.go
index ca8264c9..db92cc09 100644
--- a/main.go
+++ b/main.go
@@ -77,16 +77,6 @@ func main() {
 		return              // for safety
 	}
 
-	// Run profiling if it's activated.
-	// TODO: Should we pass a logger into this?
-	ctx, cancel := context.WithCancel(context.Background())
-	if err := pprof.Run(ctx); err != nil {
-		fmt.Println(err)
-		os.Exit(1)
-		//return // redundant
-	}
-	defer cancel()
-
 	cliUtil.LogSetup(debug)
 	data := &cliUtil.Data{
 		Program: program,
@@ -102,6 +92,16 @@ func main() {
 		Args: os.Args,
 	}
 
+	// Run profiling if it's activated.
+	// TODO: Should we pass a logger into this?
+	ctx, cancel := context.WithCancel(context.Background())
+	if err := pprof.Run(ctx); err != nil {
+		fmt.Println(err)
+		os.Exit(1)
+		//return // redundant
+	}
+	defer cancel()
+
 	name := ""
 	if len(data.Args) > 1 {
 		name = data.Args[1]
diff --git a/util/pprof/pprof.go b/util/pprof/pprof.go
index 8dd68636..5503d529 100644
--- a/util/pprof/pprof.go
+++ b/util/pprof/pprof.go
@@ -46,7 +46,7 @@ import (
 // var, it returns nil. If this is not able to start logging, then it errors. If
 // it starts logging, this waits for an exit signal in a goroutine and returns
 // nil. The magic env name is MGMT_PPROF_PATH. Example usage:
-// MGMT_PPROF_PATH="~/pprof/out.pprof" ./mgmt run lang examples/lang/hello.mcl
+// MGMT_PPROF_PATH=~/pprof/"out.pprof" ./mgmt run lang examples/lang/hello0.mcl
 // go tool pprof -no_browser -http :10000 ~/pprof/out.pprof
 func Run(ctx context.Context) error {
 	s := os.Getenv("MGMT_PPROF_PATH")
@@ -55,10 +55,12 @@ func Run(ctx context.Context) error {
 		log.Printf(format, v...) // XXX: use parent logger when available
 	}
 
-	if s == "" || !strings.HasPrefix(s, "/") {
+	if s == "" {
 		return nil // not activated
 	}
-	logf("pprof logging to: %s", s)
+	if !strings.HasPrefix(s, "/") {
+		return fmt.Errorf("pprof path is not absolute")
+	}
 
 	f, err := os.Create(s)
 	if err != nil {
@@ -67,6 +69,7 @@ func Run(ctx context.Context) error {
 	if err := pprof.StartCPUProfile(f); err != nil {
 		return fmt.Errorf("could not start CPU profile: %v", err)
 	}
+	logf("pprof logging to: %s", s)
 
 	signals := make(chan os.Signal, 1+1) // 1 * ^C + 1 * SIGTERM
 	signal.Notify(signals, os.Interrupt) // catch ^C