smartconnpool: avoid close deadlock with refresh reopen (#20157)

Signed-off-by: Arthur Schreiber <arthur@planetscale.com>

Author: arthurschreiber

go/pools/smartconnpool/pool.go

@@ -295,10 +295,10 @@ func (pool *ConnPool[C]) Close() {
 // pool closing operation
 func (pool *ConnPool[C]) CloseWithContext(ctx context.Context) error {
 	pool.capacityMu.Lock()
-	defer pool.capacityMu.Unlock()
 
-	if pool.close.Load() == nil || pool.capacity.Load() == 0 {
+	if pool.close.Load() == nil {
 		// already closed
+		pool.capacityMu.Unlock()
 		return nil
 	}
 
@@ -310,6 +310,7 @@ func (pool *ConnPool[C]) CloseWithContext(ctx context.Context) error {
 	closeChan := *pool.close.Swap(nil)
 	close(closeChan)
 
+	pool.capacityMu.Unlock()
 	pool.workers.Wait()
 	return err
 }
@@ -318,6 +319,10 @@ func (pool *ConnPool[C]) reopen() {
 	pool.capacityMu.Lock()
 	defer pool.capacityMu.Unlock()
 
+	if pool.close.Load() == nil {
+		return
+	}
+
 	capacity := pool.capacity.Load()
 	if capacity == 0 {
 		return
@@ -411,6 +416,9 @@ func (pool *ConnPool[C]) Get(ctx context.Context, setting *Setting) (*Pooled[C],
 	if ctx.Err() != nil {
 		return nil, ErrCtxTimeout
 	}
+	if pool.close.Load() == nil {
+		return nil, ErrConnPoolClosed
+	}
 	if pool.capacity.Load() == 0 {
 		return nil, ErrConnPoolClosed
 	}
@@ -427,6 +435,10 @@ func (pool *ConnPool[C]) put(conn *Pooled[C]) {
 
 	if conn == nil {
 		var err error
+		if pool.close.Load() == nil {
+			pool.closedConn()
+			return
+		}
 		// Using context.Background() is fine since MySQL connection already enforces
 		// a connect timeout via the `db-connect-timeout-ms` config param.
 		conn, err = pool.connNew(context.Background())
@@ -453,6 +465,12 @@ func (pool *ConnPool[C]) put(conn *Pooled[C]) {
 }
 
 func (pool *ConnPool[C]) tryReturnConn(conn *Pooled[C], updateIdleTime bool) bool {
+	if pool.close.Load() == nil {
+		conn.Close()
+		pool.closedConn()
+		return false
+	}
+
 	if pool.wait.tryReturnConn(conn) {
 		return true
 	}
@@ -588,6 +606,10 @@ func (pool *ConnPool[C]) closedConn() {
 
 func (pool *ConnPool[C]) getNew(ctx context.Context) (*Pooled[C], error) {
 	for {
+		if pool.close.Load() == nil {
+			return nil, ErrConnPoolClosed
+		}
+
 		open := pool.active.Load()
 		if open >= pool.capacity.Load() {
 			return nil, nil
@@ -599,6 +621,11 @@ func (pool *ConnPool[C]) getNew(ctx context.Context) (*Pooled[C], error) {
 				pool.closedConn()
 				return nil, err
 			}
+			if pool.close.Load() == nil {
+				conn.Close()
+				pool.closedConn()
+				return nil, ErrConnPoolClosed
+			}
 			return conn, nil
 		}
 	}
@@ -763,7 +790,12 @@ func (pool *ConnPool[C]) setCapacity(ctx context.Context, newcap int64) error {
 	}
 
 	oldcap := pool.capacity.Swap(newcap)
-	if oldcap == newcap {
+	// Skip the drain only when capacity is unchanged AND we're already at or
+	// below the target. Otherwise we may have been left with active > newcap
+	// by a prior call that timed out (e.g. SetCapacity(0) racing with held
+	// conns), and CloseWithContext relies on a follow-up call here to finish
+	// draining.
+	if oldcap == newcap && pool.active.Load() <= newcap {
 		return nil
 	}
 	// update the idle count to match the new capacity if necessary

go/pools/smartconnpool/pool_test.go

@@ -659,6 +659,78 @@ func TestUserClosing(t *testing.T) {
 	}
 }
 
+func TestCloseWithContextAfterSetCapacityZeroClosesPool(t *testing.T) {
+	var state TestState
+
+	p := NewPool(&Config[*TestConn]{
+		Capacity: 1,
+		LogWait:  state.LogWait,
+	}).Open(newConnector(&state), nil)
+
+	t.Cleanup(func() {
+		if closeChan := p.close.Swap(nil); closeChan != nil {
+			close(*closeChan)
+			p.workers.Wait()
+		}
+	})
+
+	ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second)
+	defer cancel()
+
+	require.NoError(t, p.SetCapacity(ctx, 0))
+	require.True(t, p.IsOpen())
+
+	require.NoError(t, p.CloseWithContext(ctx))
+	require.False(t, p.IsOpen())
+}
+
+// TestCloseWithContextDrainsAfterTimedOutSetCapacityZero guards against an
+// early return in setCapacity that previously caused CloseWithContext to
+// return nil despite active conns being out, when a prior SetCapacity(0)
+// had timed out with conns still borrowed. The path is:
+//
+//   - SetCapacity(0) times out: capacity is swapped to 0 first, then the
+//     drain loop hits the deadline and returns an error.
+//   - The next CloseWithContext calls setCapacity(0) again. With the bug,
+//     oldcap == newcap == 0 short-circuited the drain loop and Close
+//     returned nil even though pool.active was still > 0.
+//
+// CloseWithContext must instead attempt to drain and surface a timeout
+// when borrowed conns aren't returned.
+func TestCloseWithContextDrainsAfterTimedOutSetCapacityZero(t *testing.T) {
+	var state TestState
+
+	p := NewPool(&Config[*TestConn]{
+		Capacity: 2,
+		LogWait:  state.LogWait,
+	}).Open(newConnector(&state), nil)
+
+	t.Cleanup(func() {
+		if closeChan := p.close.Swap(nil); closeChan != nil {
+			close(*closeChan)
+			p.workers.Wait()
+		}
+	})
+
+	conn, err := p.Get(t.Context(), nil)
+	require.NoError(t, err)
+
+	// SetCapacity(0) must time out: a conn is still borrowed.
+	cancelledCtx, cancel := context.WithCancel(t.Context())
+	cancel()
+	require.Error(t, p.SetCapacity(cancelledCtx, 0))
+	require.EqualValues(t, 0, p.Capacity())
+	require.GreaterOrEqual(t, p.Active(), int64(1))
+
+	// CloseWithContext must still attempt to drain — without the fix it
+	// returns nil immediately because setCapacity short-circuits on
+	// oldcap == newcap.
+	require.Error(t, p.CloseWithContext(cancelledCtx))
+
+	// Release the held conn so it is properly closed during teardown.
+	conn.Recycle()
+}
+
 func TestConnReopen(t *testing.T) {
 	var state TestState