From 75303456201686776bdb052adc57387da3da3f32 Mon Sep 17 00:00:00 2001
From: Jeff McCune <jeff@openinfrastructure.co>
Date: Fri, 29 Nov 2024 09:57:32 -0800
Subject: [PATCH] main: add tracing and profiling

Writes files based on parent pid and process pid to avoid collisions.

Analyze with:

export HOLOS_TRACE=trace.%d.%d.out
go tool trace trace.999.1000.out

export HOLOS_CPU_PROFILE=cpu.%d.%d.prof
go tool pprof cpu.999.1000.prof

export HOLOS_MEM_PROFILE=mem.%d.%d.prof
go tool pprof mem.999.1000.prof
---
 internal/cli/main.go | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/internal/cli/main.go b/internal/cli/main.go
index 2b7e8ed8..e55cf874 100644
--- a/internal/cli/main.go
+++ b/internal/cli/main.go
@@ -4,6 +4,9 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"os"
+	"runtime/pprof"
+	"runtime/trace"
 
 	"connectrpc.com/connect"
 	cue "cuelang.org/go/cue/errors"
@@ -12,12 +15,48 @@ import (
 	"google.golang.org/genproto/googleapis/rpc/errdetails"
 )
 
+func memProfile(ctx context.Context, cfg *holos.Config) {
+	if format := os.Getenv("HOLOS_MEM_PROFILE"); format != "" {
+		f, _ := os.Create(fmt.Sprintf(format, os.Getppid(), os.Getpid()))
+		defer f.Close()
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			_ = HandleError(ctx, err, cfg)
+		}
+	}
+}
+
 // MakeMain makes a main function for the cli or tests.
 func MakeMain(options ...holos.Option) func() int {
 	return func() (exitCode int) {
 		cfg := holos.New(options...)
 		slog.SetDefault(cfg.Logger())
 		ctx := context.Background()
+
+		if format := os.Getenv("HOLOS_CPU_PROFILE"); format != "" {
+			f, _ := os.Create(fmt.Sprintf(format, os.Getppid(), os.Getpid()))
+			err := pprof.StartCPUProfile(f)
+			defer func() {
+				pprof.StopCPUProfile()
+				f.Close()
+			}()
+			if err != nil {
+				return HandleError(ctx, err, cfg)
+			}
+		}
+		defer memProfile(ctx, cfg)
+
+		if format := os.Getenv("HOLOS_TRACE"); format != "" {
+			f, _ := os.Create(fmt.Sprintf(format, os.Getppid(), os.Getpid()))
+			err := trace.Start(f)
+			defer func() {
+				trace.Stop()
+				f.Close()
+			}()
+			if err != nil {
+				return HandleError(ctx, err, cfg)
+			}
+		}
+
 		feature := &holos.EnvFlagger{}
 		if err := New(cfg, feature).ExecuteContext(ctx); err != nil {
 			return HandleError(ctx, err, cfg)