Add tests. Clean up diff.

mdaigle · mdaigle · commit 4233f943e863 · 2025-06-09T13:44:36.000-07:00
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs
@@ -110,9 +110,6 @@ public DbConnectionPoolState State
 
         // Counts the total number of open connections tracked by the pool.
         private volatile int _numConnections;
-
-        // Counts the number of connections currently sitting idle in the pool.
-        private volatile int _idleCount;
         #endregion
 
 
@@ -140,7 +137,7 @@ internal ChannelDbConnectionPool(
 
             _connections = new DbConnectionInternal[MaxPoolSize];
 
-            // We enforce Max Pool Size, so no need to to create a bounded channel (which is less efficient)
+            // We enforce Max Pool Size, so no need to create a bounded channel (which is less efficient)
             // On the consuming side, we have the multiplexing write loop but also non-multiplexing Rents
             // On the producing side, we have connections being released back into the pool (both multiplexing and not)
             var idleChannel = Channel.CreateUnbounded<DbConnectionInternal?>();
@@ -198,6 +195,11 @@ public bool TryGetConnection(DbConnection owningObject, TaskCompletionSource<DbC
         {
             if (taskCompletionSource is not null)
             {
+                // This is ugly, but async anti-patterns further up the stack necessitate a fresh task to be created.
+                // Ideally we would just return a Task<DbConnectionInternal> and let the caller await it as needed,
+                // but we need to signal to the provided TaskCompletionSource when the connection is established.
+                // This pattern has implications for connection open retry logic that are intricate enough to merit
+                // dedicated work.
                 Task.Run(async () =>
                 {
                     //TODO: use same timespan everywhere and tick down for queueuing and actual connection opening work
@@ -224,7 +226,7 @@ public bool TryGetConnection(DbConnection owningObject, TaskCompletionSource<DbC
             }
         }
 
-        internal async Task<DbConnectionInternal> GetInternalConnection(DbConnection owningConnection, DbConnectionOptions userOptions, TimeSpan timeout, bool async, CancellationToken cancellationToken)
+        private async Task<DbConnectionInternal> GetInternalConnection(DbConnection owningConnection, DbConnectionOptions userOptions, TimeSpan timeout, bool async, CancellationToken cancellationToken)
         {
             DbConnectionInternal? connection = null;
 
@@ -282,7 +284,8 @@ internal async Task<DbConnectionInternal> GetInternalConnection(DbConnection own
                             }
                         }
 
-                        if (connection != null && CheckIdleConnection(connection))
+                        // TODO: check if connection is still valid
+                        if (connection != null)
                         {
                             //TODO: set connection internal state
                             //PrepareConnection(owningConnection, connection, transaction);
@@ -332,42 +335,16 @@ internal async Task<DbConnectionInternal> GetInternalConnection(DbConnection own
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private DbConnectionInternal? GetIdleConnection()
         {
-
             while (_idleConnectionReader.TryRead(out DbConnectionInternal? connection))
             {
-                if (CheckIdleConnection(connection))
-                {
-                    return connection;
-                }
+                // TODO: check if connection is still valid
+                // if (CheckIdleConnection(connection))
+                return connection;
             }
 
             return null;
         }
 
-        /// <summary>
-        /// Checks that the provided connection is live and unexpired and closes it if needed.
-        /// Decrements the idle count as long as the connection is not null.
-        /// </summary>
-        /// <param name="connection">The connection to be checked.</param>
-        /// <returns>Returns true if the connection is live and unexpired, otherwise returns false.</returns>
-        /// TODO: profile the inlining to see if it's necessary
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private bool CheckIdleConnection(DbConnectionInternal? connection)
-        {
-            if (connection is null)
-            {
-                return false;
-            }
-
-            // Only decrement when the connection has a value.
-            Interlocked.Decrement(ref _idleCount);
-
-            //TODO: check if the connection is expired
-            //return CheckConnection(connection);
-
-            return true;
-        }
-
         /// <inheritdoc/>
         internal Task<DbConnectionInternal?> OpenNewInternalConnection(DbConnection? owningConnection, DbConnectionOptions userOptions, TimeSpan timeout, bool async, CancellationToken cancellationToken)
         {
diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPoolTest.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPoolTest.cs
@@ -26,7 +26,7 @@ public class ChannelDbConnectionPoolTest
         public ChannelDbConnectionPoolTest()
         {
             // Use a stubbed connection factory to avoid network code
-            connectionFactory = new ConnectionFactoryStub();
+            connectionFactory = new SuccessfulConnectionFactory();
             identity = DbConnectionPoolIdentity.NoIdentity;
             connectionPoolProviderInfo = new DbConnectionPoolProviderInfo();
             poolGroupOptions = new DbConnectionPoolGroupOptions(
@@ -54,20 +54,21 @@ public ChannelDbConnectionPoolTest()
         [InlineData(1)]
         [InlineData(5)]
         [InlineData(10)]
-        public async Task TestGetConnectionFromEmptyPoolSync_ShouldCreateNewConnection(int numConnections)
+        public void TestGetConnectionFromEmptyPoolSync_ShouldCreateNewConnection(int numConnections)
         {
             // Act
             for (int i = 0; i < numConnections; i++)
             {
-                var internalConnection = await pool.GetInternalConnection(
+                DbConnectionInternal internalConnection = null;
+                var completed = pool.TryGetConnection(
                     new SqlConnection(),
+                    taskCompletionSource: null,
                     new DbConnectionOptions("", null),
-                    TimeSpan.Zero,
-                    async: false,
-                    CancellationToken.None
+                    out internalConnection
                 );
 
                 // Assert
+                Assert.True(completed);
                 Assert.NotNull(internalConnection);
             }
 
@@ -85,23 +86,29 @@ public async Task TestGetConnectionFromEmptyPoolAsync_ShouldCreateNewConnection(
             // Act
             for (int i = 0; i < numConnections; i++)
             {
-                var internalConnection = await pool.GetInternalConnection(
+                var tcs = new TaskCompletionSource<DbConnectionInternal>();
+                DbConnectionInternal internalConnection = null;
+                var completed = pool.TryGetConnection(
                     new SqlConnection(),
+                    tcs,
                     new DbConnectionOptions("", null),
-                    TimeSpan.Zero,
-                    async: false,
-                    CancellationToken.None
+                    out internalConnection
                 );
 
                 // Assert
-                Assert.NotNull(internalConnection);
+                Assert.False(completed);
+                Assert.Null(internalConnection);
+                Assert.NotNull(await tcs.Task);
             }
 
 
             // Assert
             Assert.Equal(numConnections, pool.Count);
         }
 
+        // Test that requests to get connection from the pool fails when max pool size is reached
+       
+
         #region Property Tests
 
         [Fact]
@@ -236,15 +243,53 @@ public void TestTransactionEnded()
         {
             Assert.Throws<NotImplementedException>(() => pool.TransactionEnded(null!, null!));
         }
+        [Fact]
+        public void TestGetConnectionFailsWhenMaxPoolSizeReached()
+        {
+            // Arrange
+            for (int i = 0; i < poolGroupOptions.MaxPoolSize; i++)
+            {
+                DbConnectionInternal internalConnection = null;
+                var completed = pool.TryGetConnection(
+                    new SqlConnection(),
+                    taskCompletionSource: null,
+                    new DbConnectionOptions("", null),
+                    out internalConnection
+                );
+
+                Assert.True(completed);
+                Assert.NotNull(internalConnection);
+            }
 
+            try
+            {
+                // Act
+                DbConnectionInternal extraConnection = null;
+                var exceeded = pool.TryGetConnection(
+                    //TODO: set timeout to make this faster
+                    new SqlConnection(),
+                    taskCompletionSource: null,
+                    new DbConnectionOptions("", null),
+                    out extraConnection
+                );
+            } catch (Exception ex)
+            {
+                // Assert
+                Assert.IsType<InvalidOperationException>(ex);
+                Assert.Equal("Timeout expired.  The timeout period elapsed prior to obtaining a connection from the pool.  This may have occurred because all pooled connections were in use and max pool size was reached.", ex.Message);
+            }
+
+            // Assert
+            Assert.Equal(poolGroupOptions.MaxPoolSize, pool.Count);
+        }
         #endregion
 
         #region Test classes
-        internal class ConnectionFactoryStub : DbConnectionFactory
+        internal class SuccessfulConnectionFactory : DbConnectionFactory
         {
             protected override DbConnectionInternal CreateConnection(DbConnectionOptions options, DbConnectionPoolKey poolKey, object poolGroupProviderInfo, IDbConnectionPool pool, DbConnection owningConnection)
             {
-                return new TestDbConnectionInternal();
+                return new SuccessfulDbConnectionInternal();
             }
 
             #region Not Implemented Members
@@ -302,7 +347,7 @@ internal override void SetInnerConnectionTo(DbConnection owningObject, DbConnect
             #endregion
         }
 
-        internal class TestDbConnectionInternal : DbConnectionInternal
+        internal class SuccessfulDbConnectionInternal : DbConnectionInternal
         {
             #region Not Implemented Members
             public override string ServerVersion => throw new NotImplementedException();
@@ -328,6 +373,68 @@ protected override void Deactivate()
             }
             #endregion
         }
+
+        internal class TimeoutConnectionFactory : DbConnectionFactory
+        {
+            protected override DbConnectionInternal CreateConnection(DbConnectionOptions options, DbConnectionPoolKey poolKey, object poolGroupProviderInfo, IDbConnectionPool pool, DbConnection owningConnection)
+            {
+                return new SuccessfulDbConnectionInternal();
+            }
+
+            #region Not Implemented Members
+            public override DbProviderFactory ProviderFactory => throw new NotImplementedException();
+
+            protected override DbConnectionOptions CreateConnectionOptions(string connectionString, DbConnectionOptions previous)
+            {
+                throw new NotImplementedException();
+            }
+
+            protected override DbConnectionPoolGroupOptions CreateConnectionPoolGroupOptions(DbConnectionOptions options)
+            {
+                throw new NotImplementedException();
+            }
+
+            protected override int GetObjectId(DbConnection connection)
+            {
+                throw new NotImplementedException();
+            }
+
+            internal override DbConnectionPoolGroup GetConnectionPoolGroup(DbConnection connection)
+            {
+                throw new NotImplementedException();
+            }
+
+            internal override DbConnectionInternal GetInnerConnection(DbConnection connection)
+            {
+                throw new NotImplementedException();
+            }
+
+            internal override void PermissionDemand(DbConnection outerConnection)
+            {
+                throw new NotImplementedException();
+            }
+
+            internal override void SetConnectionPoolGroup(DbConnection outerConnection, DbConnectionPoolGroup poolGroup)
+            {
+                throw new NotImplementedException();
+            }
+
+            internal override void SetInnerConnectionEvent(DbConnection owningObject, DbConnectionInternal to)
+            {
+                throw new NotImplementedException();
+            }
+
+            internal override bool SetInnerConnectionFrom(DbConnection owningObject, DbConnectionInternal to, DbConnectionInternal from)
+            {
+                throw new NotImplementedException();
+            }
+
+            internal override void SetInnerConnectionTo(DbConnection owningObject, DbConnectionInternal to)
+            {
+                throw new NotImplementedException();
+            }
+            #endregion
+        }
         #endregion
     }
 }
