From cbecb737b3c118b50175a068fe6129eb303d92f7 Mon Sep 17 00:00:00 2001
From: Ben Perry <bhperry94@gmail.com>
Date: Sat, 17 Aug 2024 23:09:33 -0500
Subject: [PATCH] Add profiling support

---
 tools/cmd/bench.go | 36 +++++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/tools/cmd/bench.go b/tools/cmd/bench.go
index 2d29f8f..e95e685 100644
--- a/tools/cmd/bench.go
+++ b/tools/cmd/bench.go
@@ -4,6 +4,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"runtime"
+	"runtime/pprof"
 	"text/tabwriter"
 	"time"
 
@@ -12,7 +14,9 @@ import (
 )
 
 var (
-	benchRunOutput   string
+	benchRunOutput,
+	benchRunCpuprofile,
+	benchRunMemprofile string
 	benchRunDuration time.Duration
 
 	benchStatsInput string
@@ -37,6 +41,20 @@ func NewBenchRunCmd() *cobra.Command {
 		Args:         cobra.MinimumNArgs(1),
 		SilenceUsage: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
+			// Start CPU profile
+			if benchRunCpuprofile != "" {
+				f, err := os.Create(benchRunCpuprofile)
+				if err != nil {
+					return fmt.Errorf("could not create CPU profile: %v", err)
+				}
+				defer f.Close()
+				if err := pprof.StartCPUProfile(f); err != nil {
+					return fmt.Errorf("could not start CPU profile: %v", err)
+				}
+				defer pprof.StopCPUProfile()
+			}
+
+			// Run benchmark(s)
 			benchStats := make(benchmark.StatsCollection, len(args))
 			for i, name := range args {
 				config, err := benchmark.Benchmarks.Get(name)
@@ -55,6 +73,20 @@ func NewBenchRunCmd() *cobra.Command {
 				benchStats[i] = stats
 			}
 
+			// Dump memory profile
+			if benchRunMemprofile != "" {
+				f, err := os.Create(benchRunMemprofile)
+				if err != nil {
+					return fmt.Errorf("could not create memory profile: %v", err)
+				}
+				defer f.Close()
+				runtime.GC() // get up-to-date statistics
+				if err := pprof.WriteHeapProfile(f); err != nil {
+					return fmt.Errorf("could not write memory profile: %v", err)
+				}
+			}
+
+			// Dump stats
 			if benchRunOutput != "" {
 				err := benchStats.Dump(benchRunOutput)
 				if err != nil {
@@ -67,6 +99,8 @@ func NewBenchRunCmd() *cobra.Command {
 
 	run.Flags().StringVarP(&benchRunOutput, "output", "o", "", "Output path for statistics file")
 	run.Flags().DurationVarP(&benchRunDuration, "duration", "d", 0, "Override duration for benchmark runs")
+	run.Flags().StringVarP(&benchRunCpuprofile, "cpuprofile", "c", "", "CPU profiling file")
+	run.Flags().StringVarP(&benchRunMemprofile, "memprofile", "m", "", "Memory profiling file")
 	return run
 }