guard the connection closed flag with atomic access

transport/connection.go

@@ -156,8 +156,7 @@ type Connection struct {
 	heartbeatCount *atomic.Uint64
 
 	wg          sync.WaitGroup
-	closed      bool
-	mu          sync.RWMutex
+	closed      atomic.Bool
 	writeMu     sync.Mutex
 	doneOnce    sync.Once
 	cleanupOnce sync.Once
@@ -228,7 +227,7 @@ func (c *Connection) Send(data []byte) error {
 	c.writeMu.Lock()
 	defer c.writeMu.Unlock()
 
-	if c.closed {
+	if c.closed.Load() {
 		return ErrConnectionClosed
 	}
 
@@ -425,14 +424,10 @@ func (c *Connection) Close() {
 
 func (c *Connection) shutdownExternal() {
 	// set closed
-	c.mu.Lock()
-	if c.closed {
+	if !c.closed.CompareAndSwap(false, true) {
 		// idempotent shutdown
-		c.mu.Unlock()
 		return
 	}
-	c.closed = true
-	c.mu.Unlock()
 
 	// perform shutdown
 	c.shutdownInner()
@@ -444,14 +439,10 @@ func (c *Connection) shutdownExternal() {
 // must wait for itself to exit.
 func (c *Connection) shutdownInternal() {
 	// set closed
-	c.mu.Lock()
-	if c.closed {
+	if !c.closed.CompareAndSwap(false, true) {
 		// idempotent shutdown
-		c.mu.Unlock()
 		return
 	}
-	c.closed = true
-	c.mu.Unlock()
 
 	// perform shutdown
 	c.shutdownInner()

transport/connection_test.go

@@ -93,7 +93,7 @@ func TestNewConnection(t *testing.T) {
 			assert.NotNil(t, conn.incoming)
 			assert.NotNil(t, conn.errors)
 			assert.NotNil(t, conn.done)
-			assert.False(t, conn.closed)
+			assert.False(t, conn.closed.Load())
 
 			// Verify default config is used if nil is passed.
 			gotCfg := conn.config
@@ -353,6 +353,32 @@ func TestStartPingerSubMinimumInterval(t *testing.T) {
 	}, "pinger should exit cleanly after Close")
 }
 
+func TestConnectionSendCloseRace(t *testing.T) {
+	// Regression for the closed-flag data race; only the race detector
+	// observes the fault, so this is meaningful under `go test -race`.
+	socket := honeybeetest.NewMockSocket()
+	conn, err := NewConnection(context.Background(), socket, nil, nil)
+	assert.NoError(t, err)
+
+	stop := make(chan struct{})
+	finished := make(chan struct{})
+	go func() {
+		defer close(finished)
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				_ = conn.Send([]byte("data"))
+			}
+		}
+	}()
+
+	conn.Close()
+	close(stop)
+	<-finished
+}
+
 // ----------------------------------------------------------------------------
 // Send
 // ----------------------------------------------------------------------------
@@ -676,7 +702,7 @@ func TestDisconnectedConnectionClose(t *testing.T) {
 
 		conn.Close()
 
-		assert.True(t, conn.closed)
+		assert.True(t, conn.closed.Load())
 		_, ok := <-conn.incoming
 		assert.False(t, ok)
 		_, ok = <-conn.errors