Support building with cpu, mem, and trace profiles (#641)

'mr build --pprof` will generate a build
that will add three new flags: --cpuprofile, --memprofile, and --trace.
These can be used to generate profiles and debug performance issues.

These are not included by default to avoid bloating the binary size.

main.go

@@ -211,14 +211,24 @@ func main() {
 		logger.Fatalf("%v: %v", cmdName, err)
 	}
 
+	if err := cmd.Profile.Start(); err != nil {
+		logger.Error("Error creating trace file", "error", err)
+	}
+
 	if err := kctx.Run(builtinShorthands); err != nil {
 		logger.Fatalf("%v: %v", cmdName, err)
 	}
+
+	if err := cmd.Profile.Stop(); err != nil {
+		logger.Error("Error closing trace file", "error", err)
+	}
 }
 
 type mainCmd struct {
 	kong.Plugins
 
+	Profile ProfileFlags `embed:""`
+
 	// Global options that are never accessed directly by subcommands.
 	Globals struct {
 		// Flags with built-in side effects.

mise.toml

@@ -79,7 +79,16 @@ run = "go mod tidy -diff"
 description = "Ensure go.mod and go.sum are up to date"
 
 [tasks.build]
-run = "go install go.abhg.dev/gs"
+usage = """
+flag "--pprof" help="Build with pofiling flags"
+"""
+run = """
+go install -tags="
+  {%- if flag(name='pprof') == 'true' -%}
+    profile
+  {%- endif -%}
+" go.abhg.dev/gs
+"""
 wait_for = ["generate"]
 description = "Build the gs binary"
 

profile_disable.go

@@ -0,0 +1,17 @@
+//go:build !profile
+
+package main
+
+// ProfileFlags is a placeholder for the profile flags.
+// The corresponding _enable Go file contains the actual implementation.
+//
+// To build gs with profiling enabled, use:
+//
+//	go build -tags profile
+type ProfileFlags struct{}
+
+// Start is a no-op for the profile disabled build.
+func (*ProfileFlags) Start() error { return nil }
+
+// Stop is a no-op for the profile disabled build.
+func (*ProfileFlags) Stop() error { return nil }

profile_enable.go

@@ -0,0 +1,138 @@
+//go:build profile
+
+package main
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"os"
+	"runtime/pprof"
+	"runtime/trace"
+)
+
+type ProfileFlags struct {
+	CPUProfile string `name:"cpuprofile" placeholder:"FILE" help:"Write CPU profile to file"`
+	MemProfile string `name:"memprofile" placeholder:"FILE" help:"Write memory profile to file"`
+	Trace      string `name:"trace" placeholder:"FILE" help:"Write trace to file"`
+
+	cpuProfileFile io.WriteCloser
+	memProfileFile io.WriteCloser
+	traceFile      io.WriteCloser
+}
+
+func (pf *ProfileFlags) Start() error {
+	return errors.Join(
+		pf.startCPUProfile(),
+		pf.startMemProfile(),
+		pf.startTrace(),
+	)
+}
+
+func (pf *ProfileFlags) Stop() error {
+	return errors.Join(
+		pf.stopCPUProfile(),
+		pf.stopMemProfile(),
+		pf.stopTrace(),
+	)
+}
+
+func (pf *ProfileFlags) startCPUProfile() error {
+	path := pf.CPUProfile
+	if path == "" {
+		return nil
+	}
+
+	f, err := os.Create(path)
+	if err != nil {
+		return fmt.Errorf("create CPU profile file: %w", err)
+	}
+
+	if err := pprof.StartCPUProfile(f); err != nil {
+		return fmt.Errorf("start CPU profile: %w", err)
+	}
+
+	pf.cpuProfileFile = f
+	return nil
+}
+
+func (pf *ProfileFlags) stopCPUProfile() error {
+	f := pf.cpuProfileFile
+	if f == nil {
+		return nil
+	}
+
+	pprof.StopCPUProfile()
+
+	if err := f.Close(); err != nil {
+		return fmt.Errorf("close CPU profile file: %w", err)
+	}
+
+	return nil
+}
+
+func (pf *ProfileFlags) startMemProfile() error {
+	path := pf.MemProfile
+	if path == "" {
+		return nil
+	}
+
+	f, err := os.Create(path)
+	if err != nil {
+		return fmt.Errorf("create memory profile file: %w", err)
+	}
+
+	pf.memProfileFile = f
+	return nil
+}
+
+func (pf *ProfileFlags) stopMemProfile() error {
+	f := pf.memProfileFile
+	if f == nil {
+		return nil
+	}
+
+	if err := pprof.Lookup("heap").WriteTo(f, 0); err != nil {
+		return fmt.Errorf("write memory profile: %w", err)
+	}
+
+	if err := f.Close(); err != nil {
+		return fmt.Errorf("close memory profile file: %w", err)
+	}
+
+	return nil
+}
+
+func (pf *ProfileFlags) startTrace() error {
+	path := pf.Trace
+	if path == "" {
+		return nil
+	}
+
+	f, err := os.Create(path)
+	if err != nil {
+		return fmt.Errorf("create trace file: %w", err)
+	}
+
+	if err := trace.Start(f); err != nil {
+		return fmt.Errorf("start trace: %w", err)
+	}
+
+	pf.traceFile = f
+	return nil
+}
+
+func (pf *ProfileFlags) stopTrace() error {
+	f := pf.traceFile
+	if f == nil {
+		return nil
+	}
+
+	trace.Stop()
+
+	if err := f.Close(); err != nil {
+		return fmt.Errorf("close trace file: %w", err)
+	}
+
+	return nil
+}