From 9cd5fe4e25cc2aab9b20e18b375efcc6377e5c36 Mon Sep 17 00:00:00 2001
From: Joe Kralicky <joekralicky@gmail.com>
Date: Thu, 7 Nov 2024 14:55:33 -0500
Subject: [PATCH] testenv: Add utility to pause/resume profiling (#5361)

---
 internal/benchmarks/latency_bench_test.go |  8 ++--
 internal/testenv/snippets/pprof.go        | 54 +++++++++++++++++++++++
 2 files changed, 58 insertions(+), 4 deletions(-)
 create mode 100644 internal/testenv/snippets/pprof.go

diff --git a/internal/benchmarks/latency_bench_test.go b/internal/benchmarks/latency_bench_test.go
index 489bec2ff..519732cfd 100644
--- a/internal/benchmarks/latency_bench_test.go
+++ b/internal/benchmarks/latency_bench_test.go
@@ -6,8 +6,8 @@ import (
 	"io"
 	"math/rand/v2"
 	"net/http"
-	"runtime"
 	"testing"
+	"time"
 
 	"github.com/pomerium/pomerium/internal/testenv"
 	"github.com/pomerium/pomerium/internal/testenv/scenarios"
@@ -27,7 +27,7 @@ func init() {
 }
 
 func TestRequestLatency(t *testing.T) {
-	runtime.MemProfileRate = 0
+	resume := snippets.PauseProfiling(t)
 	env := testenv.New(t, testenv.Silent())
 	users := []*scenarios.User{}
 	for i := range numRoutes {
@@ -52,8 +52,8 @@ func TestRequestLatency(t *testing.T) {
 	env.AddUpstream(up)
 
 	env.Start()
-	snippets.WaitStartupComplete(env)
-	runtime.MemProfileRate = 512 * 1024
+	snippets.WaitStartupComplete(env, 1*time.Hour)
+	resume()
 
 	out := testing.Benchmark(func(b *testing.B) {
 		b.ReportAllocs()
diff --git a/internal/testenv/snippets/pprof.go b/internal/testenv/snippets/pprof.go
new file mode 100644
index 000000000..04a020b3e
--- /dev/null
+++ b/internal/testenv/snippets/pprof.go
@@ -0,0 +1,54 @@
+package snippets
+
+import (
+	"flag"
+	"os"
+	"path/filepath"
+	"runtime"
+	"runtime/pprof"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// PauseProfiling will suspend CPU and memory profiling, if started using the
+// -cpuprofile and/or -memprofile test flags. The returned function will restart
+// profiling when called. Existing CPU profile data is overwritten, but
+// existing memory profile data is kept.
+func PauseProfiling(t testing.TB) (resume func()) {
+	resumeFuncs := []func(){}
+
+	outputdir := flag.Lookup("test.outputdir")
+	if f := flag.Lookup("test.cpuprofile"); f != nil {
+		filename := f.Value.String()
+		if outputdir != nil {
+			filename = filepath.Join(outputdir.Value.String(), filename)
+		}
+		if _, err := os.Stat(filename); err == nil {
+			pprof.StopCPUProfile()
+			t.Logf("pausing cpu profiling (%s)", filename)
+			resumeFuncs = append(resumeFuncs, func() {
+				t.Logf("resuming cpu profiling (%s)", filename)
+				f, err := os.Create(filename)
+				require.NoError(t, err)
+				require.NoError(t, pprof.StartCPUProfile(f))
+			})
+		}
+	}
+
+	if f := flag.Lookup("test.memprofile"); f != nil {
+		rate := runtime.MemProfileRate
+		runtime.MemProfileRate = 0
+		t.Log("pausing memory profiling")
+		resumeFuncs = append(resumeFuncs, func() {
+			t.Log("resuming memory profiling")
+			runtime.MemProfileRate = rate
+		})
+	}
+	return sync.OnceFunc(func() {
+		for _, f := range resumeFuncs {
+			f()
+		}
+	})
+}