fix(pool): data race on Conn.onClose during init vs Close

internal/pool/conn.go

@@ -134,7 +134,12 @@ type Conn struct {
 	// Connection initialization function for reconnections
 	initConnFunc func(context.Context, *Conn) error
 
-	onClose func() error
+	// onClose is read and cleared by Close while initConn (running inside
+	// SetNetConnAndInitConn under the INITIALIZING state) installs it via
+	// SetOnClose; Close transitions to CLOSED from any state, so the two race
+	// when a pool shutdown closes a connection mid-init. Stored as an atomic
+	// pointer so the setter and Close don't need a mutex (keeping Conn slim).
+	onClose atomic.Pointer[func() error]
 }
 
 func NewConn(netConn net.Conn) *Conn {
@@ -635,7 +640,11 @@ func (cn *Conn) getEffectiveWriteTimeout(normalTimeout time.Duration) time.Durat
 // time, and a richer registry here would not even solve the "stale
 // closure" hazard described above.
 func (cn *Conn) SetOnClose(fn func() error) {
-	cn.onClose = fn
+	if fn == nil {
+		cn.onClose.Store(nil)
+		return
+	}
+	cn.onClose.Store(&fn)
 }
 
 // SetInitConnFunc sets the connection initialization function to be called on reconnections.
@@ -967,10 +976,9 @@ func (cn *Conn) Close() error {
 	// Transition to CLOSED state
 	cn.stateMachine.Transition(StateClosed)
 
-	if cn.onClose != nil {
+	if fn := cn.onClose.Swap(nil); fn != nil {
 		// ignore error
-		_ = cn.onClose()
-		cn.onClose = nil
+		_ = (*fn)()
 	}
 
 	// Lock-free netConn access for better performance

internal/pool/conn_onclose_race_test.go

@@ -0,0 +1,53 @@
+package pool
+
+import (
+	"context"
+	"net"
+	"sync"
+	"testing"
+)
+
+// TestConnOnCloseRaceWithInitConn is a regression test for a data race on
+// Conn.onClose between a connection being (re)initialized and a concurrent
+// Close.
+//
+// baseClient.initConn installs the close callback via SetOnClose (e.g. the
+// StreamingCredentialsProvider unsubscribe). initConn runs inside
+// SetNetConnAndInitConn while the connection is in the INITIALIZING state.
+// Conn.Close transitions to CLOSED from any state, then reads and nils
+// onClose without synchronization. A pool shutdown (or connection removal)
+// that closes a connection whose init is still in flight therefore races the
+// setter.
+func TestConnOnCloseRaceWithInitConn(t *testing.T) {
+	iterations := 5000
+	if testing.Short() {
+		iterations = 1000
+	}
+
+	for i := 0; i < iterations; i++ {
+		c1, c2 := net.Pipe()
+
+		cn := NewConn(c1)
+		cn.SetInitConnFunc(func(ctx context.Context, c *Conn) error {
+			// Mirror baseClient.initConn: install the close hook, then mark
+			// the connection idle so it is ready for use.
+			c.SetOnClose(func() error { return nil })
+			c.GetStateMachine().Transition(StateIdle)
+			return nil
+		})
+
+		var wg sync.WaitGroup
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			_ = cn.SetNetConnAndInitConn(context.Background(), c1)
+		}()
+		go func() {
+			defer wg.Done()
+			_ = cn.Close()
+		}()
+		wg.Wait()
+
+		c2.Close()
+	}
+}