From c5f10553735c4fd8eb2feb8f347286936a15ed90 Mon Sep 17 00:00:00 2001
From: Ivan Dlugos <dlugos.ivan@gmail.com>
Date: Tue, 15 Aug 2023 13:08:17 +0200
Subject: [PATCH] wait for profiler to stop

---
 profiler.go      | 13 ++++++++-----
 profiler_test.go | 14 +++++++-------
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/profiler.go b/profiler.go
index 22c3bf123..5517c11d1 100644
--- a/profiler.go
+++ b/profiler.go
@@ -70,7 +70,7 @@ const profilerRuntimeLimit = 30 // seconds
 type profileRecorder struct {
 	startTime         time.Time
 	stopSignal        chan struct{}
-	stopped           sync.WaitGroup
+	stopped           int64
 	mutex             sync.RWMutex
 	testProfilerPanic int64
 
@@ -133,8 +133,10 @@ func (p *profileRecorder) run(started chan struct{}) {
 			Logger.Printf("Profiler panic in run(): %v\n", err)
 		}
 		atomic.StoreInt64(&testProfilerPanic, 0)
-		atomic.StoreInt64(&profilerRunning, 0)
 		close(started)
+		p.stopSignal <- struct{}{}
+		atomic.StoreInt64(&p.stopped, 1)
+		atomic.StoreInt64(&profilerRunning, 0)
 	}()
 
 	p.testProfilerPanic = atomic.LoadInt64(&testProfilerPanic)
@@ -159,17 +161,18 @@ func (p *profileRecorder) run(started chan struct{}) {
 			p.onTick()
 			collectTicker.Ticked()
 		case <-p.stopSignal:
-			p.stopped.Done()
 			return
 		}
 	}
 }
 
 func (p *profileRecorder) Stop(wait bool) {
-	p.stopped.Add(1)
+	if atomic.LoadInt64(&p.stopped) == 1 {
+		return
+	}
 	p.stopSignal <- struct{}{}
 	if wait {
-		p.stopped.Wait()
+		<-p.stopSignal
 	}
 }
 
diff --git a/profiler_test.go b/profiler_test.go
index b5b4fab4a..9bcc8ddce 100644
--- a/profiler_test.go
+++ b/profiler_test.go
@@ -35,7 +35,7 @@ func (t *profilerTestTicker) Ticked() {
 }
 
 func (t *profilerTestTicker) Stop() {
-	t.log("Ticker: stopped.\n")
+	t.log("Ticker: stopping.\n")
 	close(t.ticked)
 }
 
@@ -48,7 +48,7 @@ func (t *profilerTestTicker) Tick() bool {
 	select {
 	case _, ok := <-t.ticked:
 		if ok {
-			t.log("Ticker: tick acknowledged by the profiler.\n") // logged on the test goroutine
+			t.log("Ticker: tick acknowledged (received on the test goroutine).\n") // logged on the test goroutine
 			return true
 		}
 		t.log("Ticker: tick not acknowledged (ticker stopped).\n")
@@ -83,7 +83,7 @@ func TestProfilerCollection(t *testing.T) {
 
 		start := time.Now()
 		profiler := startProfiling(start)
-		defer profiler.Stop(false)
+		defer profiler.Stop(true)
 		if isCI() {
 			doWorkFor(5 * time.Second)
 		} else {
@@ -106,7 +106,7 @@ func TestProfilerCollection(t *testing.T) {
 
 		start := time.Now()
 		profiler := startProfiling(start)
-		defer profiler.Stop(false)
+		defer profiler.Stop(true)
 		require.True(ticker.Tick())
 		end := time.Now()
 		result := profiler.GetSlice(start, end)
@@ -134,7 +134,7 @@ func TestProfilerStackTrace(t *testing.T) {
 
 	start := time.Now()
 	profiler := startProfiling(start)
-	defer profiler.Stop(false)
+	defer profiler.Stop(true)
 	require.True(ticker.Tick())
 	result := profiler.GetSlice(start, time.Now())
 	require.NotNil(result)
@@ -203,7 +203,7 @@ func TestProfilerPanicOnTick(t *testing.T) {
 
 	start := time.Now()
 	profiler := startProfiling(start)
-	defer profiler.Stop(false)
+	defer profiler.Stop(true)
 	assert.True(ticker.Tick())
 	assert.False(ticker.Tick())
 
@@ -307,7 +307,7 @@ func TestProfilerSamplingRate(t *testing.T) {
 
 	start := time.Now()
 	profiler := startProfiling(start)
-	defer profiler.Stop(false)
+	defer profiler.Stop(true)
 	doWorkFor(500 * time.Millisecond)
 	end := time.Now()
 	result := profiler.GetSlice(start, end)