From feeeef71f1d88aca28de1925285e136a8047aea4 Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Mon, 22 Jan 2024 09:39:06 -0800
Subject: [PATCH] Add extra protection for grid reconnects (#18840)

Race checks would occasionally show race on handleMsgWg WaitGroup by debug messages (used in test only).

Use the `connMu` mutex to protect this against concurrent Wait/Add.

Fixes #18827
---
 internal/grid/connection.go | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/internal/grid/connection.go b/internal/grid/connection.go
index 40dc000ca..970ec45bf 100644
--- a/internal/grid/connection.go
+++ b/internal/grid/connection.go
@@ -831,7 +831,9 @@ func (c *Connection) reconnected() {
 	c.outgoing.Clear()
 
 	// Wait for existing to exit
+	c.connMu.Lock()
 	c.handleMsgWg.Wait()
+	c.connMu.Unlock()
 }
 
 func (c *Connection) updateState(s State) {
@@ -855,7 +857,9 @@ func (c *Connection) updateState(s State) {
 
 func (c *Connection) handleMessages(ctx context.Context, conn net.Conn) {
 	// Read goroutine
+	c.connMu.Lock()
 	c.handleMsgWg.Add(2)
+	c.connMu.Unlock()
 	ctx, cancel := context.WithCancelCause(ctx)
 	go func() {
 		defer func() {
@@ -1534,7 +1538,9 @@ func (c *Connection) debugMsg(d debugMsg, args ...any) {
 			c.debugInConn.Close()
 		}
 	case debugWaitForExit:
+		c.connMu.Lock()
 		c.handleMsgWg.Wait()
+		c.connMu.Unlock()
 	case debugSetConnPingDuration:
 		c.connMu.Lock()
 		defer c.connMu.Unlock()