From 40c5ba6c0adabed95a1513443c8caf8afb4f37b3 Mon Sep 17 00:00:00 2001
From: Jackson Owens <jackson@cockroachlabs.com>
Date: Mon, 6 Feb 2023 16:16:10 -0500
Subject: [PATCH] storage: fix TestPebbleMetricEventListener

In #96096 handling of disk slow events was made asynchronous. This introduced a
race during Close where the goroutine logging could still be inflight.

Fix #96638.
Epic: None
Release note: None
---
 pkg/storage/pebble.go | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/pkg/storage/pebble.go b/pkg/storage/pebble.go
index 3004fd44174c..9ecbf23ba47a 100644
--- a/pkg/storage/pebble.go
+++ b/pkg/storage/pebble.go
@@ -693,6 +693,7 @@ type Pebble struct {
 		syncutil.Mutex
 		flushCompletedCallback func()
 	}
+	asyncDone sync.WaitGroup
 
 	// supportsRangeKeys is 1 if the database supports range keys. It must
 	// be accessed atomically.
@@ -943,7 +944,7 @@ func NewPebble(ctx context.Context, cfg PebbleConfig) (p *Pebble, err error) {
 	oldDiskSlow := lel.DiskSlow
 	lel.DiskSlow = func(info pebble.DiskSlowInfo) {
 		// Run oldDiskSlow asynchronously.
-		go oldDiskSlow(info)
+		p.async(func() { oldDiskSlow(info) })
 	}
 	cfg.Opts.EventListener = pebble.TeeEventListener(
 		p.makeMetricEtcEventListener(ctx),
@@ -992,6 +993,17 @@ func NewPebble(ctx context.Context, cfg PebbleConfig) (p *Pebble, err error) {
 	return p, nil
 }
 
+// async launches the provided function in a new goroutine. It uses a wait group
+// to synchronize with (*Pebble).Close to ensure all launched goroutines have
+// exited before Close returns.
+func (p *Pebble) async(fn func()) {
+	p.asyncDone.Add(1)
+	go func() {
+		defer p.asyncDone.Done()
+		fn()
+	}()
+}
+
 func (p *Pebble) makeMetricEtcEventListener(ctx context.Context) pebble.EventListener {
 	return pebble.EventListener{
 		WriteStallBegin: func(info pebble.WriteStallBeginInfo) {
@@ -1042,8 +1054,10 @@ func (p *Pebble) makeMetricEtcEventListener(ctx context.Context) pebble.EventLis
 					log.Fatalf(ctx, "disk stall detected: pebble unable to write to %s in %.2f seconds",
 						info.Path, redact.Safe(info.Duration.Seconds()))
 				} else {
-					go log.Errorf(ctx, "disk stall detected: pebble unable to write to %s in %.2f seconds",
-						info.Path, redact.Safe(info.Duration.Seconds()))
+					p.async(func() {
+						log.Errorf(ctx, "disk stall detected: pebble unable to write to %s in %.2f seconds",
+							info.Path, redact.Safe(info.Duration.Seconds()))
+					})
 				}
 				return
 			}
@@ -1083,6 +1097,9 @@ func (p *Pebble) Close() {
 	}
 	p.closed = true
 
+	// Wait for any asynchronous goroutines to exit.
+	p.asyncDone.Wait()
+
 	handleErr := func(err error) {
 		if err == nil {
 			return