Implement pool shutdown for ChannelDbConnectionPool and harden WaitHandleDbConnectionPool shutdown

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs

@@ -5,7 +5,6 @@
 using System.Collections.Concurrent;
 using System.Collections.ObjectModel;
 using System.Data.Common;
-using System.Diagnostics;
 using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Channels;
@@ -92,6 +91,13 @@ internal sealed class ChannelDbConnectionPool : IDbConnectionPool, IDisposable
         /// Must be updated using <see cref="Interlocked"/> operations to ensure thread safety.
         /// </summary>
         private volatile int _isClearing;
+
+        /// <summary>
+        /// Tracks whether <see cref="Shutdown"/> has already initiated the shutdown sequence so that
+        /// repeated calls are observed as no-ops. Updated atomically via
+        /// <see cref="Interlocked.CompareExchange(ref int, int, int)"/>.
+        /// </summary>
+        private int _shutdownInitiated;
         #endregion
 
         /// <summary>
@@ -264,16 +270,56 @@ public void ReturnInternalConnection(DbConnectionInternal connection, DbConnecti
             }
             else
             {
-                var written = _idleChannel.TryWrite(connection);
-                Debug.Assert(written, "Failed to write returning connection to the idle channel.");
+                if (!_idleChannel.TryWrite(connection))
+                {
+                    // The channel has been completed (pool is shutting down). Race window
+                    // between the State check above and TryWrite: destroy instead of pooling.
+                    RemoveConnection(connection);
+                }
             }
         }
 
         /// <inheritdoc />
         public void Shutdown()
         {
+            // idempotent. Compare-and-exchange ensures only one caller performs shutdown work.
+            if (Interlocked.CompareExchange(ref _shutdownInitiated, 1, 0) != 0)
+            {
+                return;
+            }
+
+            SqlClientEventSource.Log.TryPoolerTraceEvent(
+                "<prov.DbConnectionPool.Shutdown|RES|INFO|CPOOL> {0}", Id);
+
+            // Transition to ShuttingDown. After this point, ReturnInternalConnection
+            // routes returning connections to RemoveConnection.
             State = ShuttingDown;
+
+            // Stop the idle-pruning timer before draining so a tick cannot race with
+            // the final drain below. PoolPruner.Dispose is idempotent and non-throwing.
             Pruner?.Dispose();
+
+            // Complete the channel writer so:
+            //  - no further idle connections can be enqueued (TryWrite returns false), and
+            //  - in-flight / future async waiters on ReadAsync fault with ChannelClosedException.
+            _idleChannel.Complete();
+
+            // Reuse Clear() for the drain. Clear bumps _clearGeneration so any active
+            // checked-out connection fails IsLiveConnection on return and is removed, and it
+            // drains the idle channel up to its captured IdleCount.
+            Clear();
+
+            // Clear() may short-circuit if another caller is already draining. Because the
+            // channel is now completed, no new items can be enqueued, so it is safe to do a
+            // final unbounded drain to mop up anything Clear() may have skipped.
+            while (_idleChannel.TryRead(out DbConnectionInternal? connection))
+            {
+                if (connection is not null)
+                {
+                    RemoveConnection(connection);
+                }
+                // null sentinels are wake-up signals only; nothing to destroy.
+            }
         }
 
         /// <summary>
@@ -284,7 +330,12 @@ public void Shutdown()
         /// <inheritdoc />
         public void Startup()
         {
-            // No-op for now, warmup will be implemented later.
+            // This pool has no background timers today (idle timeout is enforced lazily in
+            // IsLiveConnection on retrieval; pruning is not implemented). State is set to Running
+            // in the constructor, so this is currently the symmetrical counterpart of Shutdown.
+            // Background work (warmup, pruning timers) will be added here when introduced.
+            SqlClientEventSource.Log.TryPoolerTraceEvent(
+                "<prov.DbConnectionPool.Startup|RES|INFO|CPOOL> {0}", Id);
         }
 
         /// <inheritdoc />
@@ -300,6 +351,19 @@ public bool TryGetConnection(
             TimeoutTimer timeout,
             out DbConnectionInternal? connection)
         {
+            // Short-circuit when the pool is not Running (i.e., shut down or never started).
+            // Returning (true, null) matches WaitHandleDbConnectionPool.TryGetConnection and tells
+            // the caller "completed; no connection available" without entering the channel path,
+            // which would otherwise reserve a slot, attempt to open a fresh physical connection,
+            // and then immediately destroy it on return because State == ShuttingDown.
+            if (State is not Running)
+            {
+                SqlClientEventSource.Log.TryPoolerTraceEvent(
+                    "<prov.DbConnectionPool.TryGetConnection|RES|CPOOL> {0}, State != Running.", Id);
+                connection = null;
+                return true;
+            }
+
             // If taskCompletionSource is null, we are in a sync context.
             if (taskCompletionSource is null)
             {

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/IdleConnectionChannel.cs

@@ -26,10 +26,19 @@ internal IdleConnectionChannel()
         {
             var channel = Channel.CreateUnbounded<DbConnectionInternal?>();
             _reader = channel.Reader;
-            //TODO: the channel should be completed on pool shutdown
             _writer = channel.Writer;
         }
 
+        /// <summary>
+        /// Marks the channel writer as complete. After completion, <see cref="TryWrite"/>
+        /// returns <see langword="false"/> for any future writes, and any in-flight or future
+        /// <see cref="ReadAsync"/> waiters will fault with <see cref="System.Threading.Channels.ChannelClosedException"/>
+        /// once the channel is drained. Used by the connection pool to signal shutdown.
+        /// </summary>
+        /// <returns><see langword="true"/> if this call completed the channel; otherwise <see langword="false"/>
+        /// (channel was already completed).</returns>
+        internal bool Complete() => _writer.TryComplete();
+
         /// <summary>
         /// The number of non-null connections currently in the channel.
         /// </summary>

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/WaitHandleDbConnectionPool.cs

@@ -191,16 +191,16 @@ public void Dispose()
 
         private readonly WaitCallback _poolCreateRequest;
 
-        private int _waitCount;
+        internal int _waitCount;
         private readonly PoolWaitHandles _waitHandles;
 
         private Exception _resError;
         private volatile bool _errorOccurred;
 
         private int _errorWait;
-        private Timer _errorTimer;
+        internal Timer _errorTimer;
 
-        private Timer _cleanupTimer;
+        internal Timer _cleanupTimer;
 
         private readonly TransactedConnectionPool _transactedConnectionPool;
 
@@ -329,8 +329,16 @@ public bool IsRunning
 
         public TransactedConnectionPool TransactedConnectionPool => _transactedConnectionPool;
 
-        private void CleanupCallback(object state)
+        internal void CleanupCallback(object state)
         {
+            // If the pool is shutting down, skip work. Shutdown disposes the timer, but
+            // a callback may already be in-flight when Shutdown runs; this guard ensures it does
+            // not perform pruning or re-arm pool create requests.
+            if (State == ShuttingDown)
+            {
+                return;
+            }
+
             // Called when the cleanup-timer ticks over.
 
             // This is the automatic pruning method.  Every period, we will
@@ -762,8 +770,15 @@ private void DestroyObject(DbConnectionInternal obj)
             }
         }
 
-        private void ErrorCallback(object state)
+        internal void ErrorCallback(object state)
         {
+            // Skip work if the pool is shutting down. The shutdown path disposes the
+            // timer; this guard handles the in-flight-callback race.
+            if (State == ShuttingDown)
+            {
+                return;
+            }
+
             SqlClientEventSource.Log.TryPoolerTraceEvent("<prov.DbConnectionPool.ErrorCallback|RES|CPOOL> {0}, Resetting Error handling.", Id);
             _errorOccurred = false;
             _waitHandles.ErrorEvent.Reset();
@@ -989,6 +1004,31 @@ private bool TryGetConnection(DbConnection owningObject, uint waitForMultipleObj
                     {
                         waitResult = WaitHandle.WaitAny(_waitHandles.GetHandles(allowCreate), unchecked((int)waitForMultipleObjectsTimeout));
 
+                        // After waking, observe shutdown state and bail out so waiters
+                        // do not spin against a drained pool. If WaitAny consumed a
+                        // PoolSemaphore slot, release it back so the accounting stays
+                        // balanced; otherwise the slot would leak and other waiters
+                        // (or callers that arrive after Shutdown completes its own
+                        // Release loop) would starve.
+                        if (State != Running)
+                        {
+                            SqlClientEventSource.Log.TryPoolerTraceEvent("<prov.DbConnectionPool.GetConnection|RES|CPOOL> {0}, Pool is shutting down; abandoning wait.", Id);
+                            if (waitResult == SEMAPHORE_HANDLE || waitResult == WAIT_ABANDONED + SEMAPHORE_HANDLE)
+                            {
+                                try
+                                {
+                                    _waitHandles.PoolSemaphore.Release(1);
+                                }
+                                catch (SemaphoreFullException)
+                                {
+                                    // Pool semaphore was already saturated by Shutdown's bulk release; safe to ignore.
+                                }
+                            }
+                            Interlocked.Decrement(ref _waitCount);
+                            connection = null;
+                            return false;
+                        }
+
                         // From the WaitAny docs: "If more than one object became signaled during
                         // the call, this is the array index of the signaled object with the
                         // smallest index value of all the signaled objects."  This is important
@@ -1504,6 +1544,16 @@ private bool ReclaimEmancipatedObjects()
         public void Startup()
         {
             SqlClientEventSource.Log.TryPoolerTraceEvent("<prov.DbConnectionPool.Startup|RES|INFO|CPOOL> {0}, CleanupWait={1}", Id, _cleanupWait);
+
+            // Do not resurrect a shut-down pool. Without this guard a Startup() call after
+            // Shutdown() would create a fresh _cleanupTimer and queue a PoolCreateRequest
+            // against a pool that will never accept connections back, leaking the timer
+            // and scheduling background work that immediately short-circuits.
+            if (State is ShuttingDown)
+            {
+                return;
+            }
+
             _cleanupTimer = CreateCleanupTimer();
 
             if (NeedToReplenish)
@@ -1515,15 +1565,45 @@ public void Startup()
         public void Shutdown()
         {
             SqlClientEventSource.Log.TryPoolerTraceEvent("<prov.DbConnectionPool.Shutdown|RES|INFO|CPOOL> {0}", Id);
+
+            // Idempotent: subsequent calls observe ShuttingDown and bail.
+            if (State == ShuttingDown)
+            {
+                return;
+            }
             State = ShuttingDown;
 
-            // deactivate timer callbacks
-            Timer t = _cleanupTimer;
-            _cleanupTimer = null;
-            if (t != null)
+            // Dispose all background timers so they no longer schedule new work.
+            // Note that any timer callback already in flight may still observe State == ShuttingDown
+            // and short-circuit (see CleanupCallback / ErrorCallback).
+            Timer cleanup = Interlocked.Exchange(ref _cleanupTimer, null);
+            cleanup?.Dispose();
+            Timer error = Interlocked.Exchange(ref _errorTimer, null);
+            error?.Dispose();
+
+            // Wake any threads parked in WaitHandle.WaitAny by releasing as many semaphore
+            // slots as there are recorded waiters. Using _waitCount (rather than MaxPoolSize)
+            // avoids ArgumentOutOfRangeException when MaxPoolSize == 0 (unlimited) and ensures
+            // we wake every parked waiter even when _waitCount exceeds MaxPoolSize. Waiters
+            // observe State != Running after wake-up and bail.
+            int waitersToWake = Volatile.Read(ref _waitCount);
+            if (waitersToWake > 0)
             {
-                t.Dispose();
+                try
+                {
+                    _waitHandles.PoolSemaphore.Release(waitersToWake);
+                }
+                catch (SemaphoreFullException)
+                {
+                    // Semaphore already saturated; nothing to do.
+                }
             }
+
+            // Reuse Clear() to doom every connection (including active checked-out ones), drain
+            // both idle stacks, and reclaim emancipated objects. Active connections destroy
+            // themselves on return either via the doom flag or via DeactivateObject's
+            // State == ShuttingDown branch.
+            Clear();
         }
 
         // TransactionEnded merely provides the plumbing for DbConnectionInternal to access the transacted pool