Port connection pool tests from devdiv (#1155)

src/Microsoft.Data.SqlClient/tests/ManualTests/Microsoft.Data.SqlClient.ManualTesting.Tests.csproj

@@ -264,6 +264,10 @@
   </ItemGroup>
   <ItemGroup Condition="'$(TargetGroup)' == 'netfx'">
     <Reference Include="System.Transactions" />
+    <Compile Include="SQL\ConnectionPoolTest\TransactionPoolTest.cs" />
+  </ItemGroup>
+  <ItemGroup Condition="$(Configuration.Contains('Debug'))">
+    <Compile Include="SQL\ConnectionPoolTest\ConnectionPoolTest.Debug.cs" />
   </ItemGroup>
   <ItemGroup>
     <ProjectReference Include="$(TestsPath)tools\Microsoft.DotNet.XUnitExtensions\Microsoft.DotNet.XUnitExtensions.csproj">

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/AsyncTest/BeginExecAsyncTest.cs

@@ -34,38 +34,18 @@ private static string GenerateCommandText()
         [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsNotAzureSynapse))]
         public static void ExecuteTest()
         {
-            using (SqlConnection connection = new SqlConnection(DataTestUtility.TCPConnectionString))
+            using SqlConnection connection = new(DataTestUtility.TCPConnectionString);
+
+            using SqlCommand command = new(GenerateCommandText(), connection);
+            connection.Open();
+
+            IAsyncResult result = command.BeginExecuteNonQuery();
+            while (!result.IsCompleted)
             {
-                try
-                {
-                    SqlCommand command = new SqlCommand(GenerateCommandText(), connection);
-                    connection.Open();
-
-                    IAsyncResult result = command.BeginExecuteNonQuery();
-                    while (!result.IsCompleted)
-                    {
-                        System.Threading.Thread.Sleep(100);
-                    }
-                    Assert.True(command.EndExecuteNonQuery(result) > 0, "FAILED: BeginExecuteNonQuery did not complete successfully.");
-                }
-                catch (SqlException ex)
-                {
-                    Console.WriteLine("Error ({0}): {1}", ex.Number, ex.Message);
-                    Assert.Null(ex);
-                }
-                catch (InvalidOperationException ex)
-                {
-                    Console.WriteLine("Error: {0}", ex.Message);
-                    Assert.Null(ex);
-                }
-                catch (Exception ex)
-                {
-                    // You might want to pass these errors
-                    // back out to the caller.
-                    Console.WriteLine("Error: {0}", ex.Message);
-                    Assert.Null(ex);
-                }
+                System.Threading.Thread.Sleep(100);
             }
+
+            Assert.True(command.EndExecuteNonQuery(result) > 0, "FAILED: BeginExecuteNonQuery did not complete successfully.");
         }
 
         // Synapse: Parse error at line: 1, column: 201: Incorrect syntax near ';'.
@@ -74,24 +54,12 @@ public static void FailureTest()
         {
             using (SqlConnection connection = new SqlConnection(DataTestUtility.TCPConnectionString))
             {
-                bool caughtException = false;
                 SqlCommand command = new SqlCommand(GenerateCommandText(), connection);
                 connection.Open();
 
                 //Try to execute a synchronous query on same command
                 IAsyncResult result = command.BeginExecuteNonQuery();
-                try
-                {
-                    command.ExecuteNonQuery();
-                }
-                catch (Exception ex)
-                {
-                    Assert.True(ex is InvalidOperationException, "FAILED: Thrown exception for BeginExecuteNonQuery was not an InvalidOperationException");
-                    caughtException = true;
-                }
-
-                Assert.True(caughtException, "FAILED: No exception thrown after trying second BeginExecuteNonQuery.");
-                caughtException = false;
+                InvalidOperationException ex = Assert.Throws<InvalidOperationException>(() => command.ExecuteNonQuery());
 
                 while (!result.IsCompleted)
                 {

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/Common/InternalConnectionWrapper.cs

@@ -18,6 +18,22 @@ public class InternalConnectionWrapper
         private object _internalConnection = null;
         private object _spid = null;
 
+        /// <summary>
+        /// Is this internal connection enlisted in a distributed transaction?
+        /// </summary>
+        public bool IsEnlistedInTransaction => ConnectionHelper.IsEnlistedInTransaction(_internalConnection);
+
+        /// <summary>
+        /// Is this internal connection the root of a distributed transaction?
+        /// </summary>
+        public bool IsTransactionRoot => ConnectionHelper.IsTransactionRoot(_internalConnection);
+
+        /// <summary>
+        /// True if this connection is the root of a transaction AND it is waiting for the transaction 
+        /// to complete (i.e. it has been 'aged' or 'put into stasis'), otherwise false
+        /// </summary>
+        public bool IsTxRootWaitingForTxEnd => ConnectionHelper.IsTxRootWaitingForTxEnd(_internalConnection);
+
         /// <summary>
         /// Gets the internal connection associated with the given SqlConnection
         /// </summary>

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/Common/SystemDataInternals/ConnectionHelper.cs

@@ -32,6 +32,9 @@ internal static class ConnectionHelper
         private static PropertyInfo s_pendingSQLDNS_AddrIPv4 = s_SQLDNSInfo.GetProperty("AddrIPv4", BindingFlags.Instance | BindingFlags.Public);
         private static PropertyInfo s_pendingSQLDNS_AddrIPv6 = s_SQLDNSInfo.GetProperty("AddrIPv6", BindingFlags.Instance | BindingFlags.Public);
         private static PropertyInfo s_pendingSQLDNS_Port = s_SQLDNSInfo.GetProperty("Port", BindingFlags.Instance | BindingFlags.Public);
+        private static PropertyInfo dbConnectionInternalIsTransRoot = s_dbConnectionInternal.GetProperty("IsTransactionRoot", BindingFlags.Instance | BindingFlags.NonPublic);
+        private static PropertyInfo dbConnectionInternalEnlistedTrans = s_sqlInternalConnection.GetProperty("EnlistedTransaction", BindingFlags.Instance | BindingFlags.NonPublic);
+        private static PropertyInfo dbConnectionInternalIsTxRootWaitingForTxEnd = s_dbConnectionInternal.GetProperty("IsTxRootWaitingForTxEnd", BindingFlags.Instance | BindingFlags.NonPublic);
 
         public static object GetConnectionPool(object internalConnection)
         {
@@ -69,6 +72,24 @@ private static void VerifyObjectIsConnection(object connection)
                 throw new ArgumentException("Object provided was not a SqlConnection", nameof(connection));
         }
 
+        public static bool IsEnlistedInTransaction(object internalConnection)
+        {
+            VerifyObjectIsInternalConnection(internalConnection);
+            return (dbConnectionInternalEnlistedTrans.GetValue(internalConnection, null) != null);
+        }
+
+        public static bool IsTransactionRoot(object internalConnection)
+        {
+            VerifyObjectIsInternalConnection(internalConnection);
+            return (bool)dbConnectionInternalIsTransRoot.GetValue(internalConnection, null);
+        }
+
+        public static bool IsTxRootWaitingForTxEnd(object internalConnection)
+        {
+            VerifyObjectIsInternalConnection(internalConnection);
+            return (bool)dbConnectionInternalIsTxRootWaitingForTxEnd.GetValue(internalConnection, null);
+        }
+
         public static object GetParser(object internalConnection)
         {
             VerifyObjectIsInternalConnection(internalConnection);

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/Common/SystemDataInternals/TdsParserHelper.cs

@@ -18,7 +18,7 @@ private static void VerifyObjectIsTdsParser(object parser)
             if (parser == null)
                 throw new ArgumentNullException("stateObject");
             if (!s_tdsParser.IsInstanceOfType(parser))
-                throw new ArgumentException("Object provided was not a DbConnectionInternal", "internalConnection");
+                throw new ArgumentException("Object provided was not a TdsParser", nameof(parser));
         }
 
         internal static object GetStateObject(object parser)

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/Common/SystemDataInternals/TdsParserStateObjectHelper.cs

@@ -56,7 +56,7 @@ private static void VerifyObjectIsTdsParserStateObject(object stateObject)
             if (stateObject == null)
                 throw new ArgumentNullException(nameof(stateObject));
             if (!s_tdsParserStateObjectManaged.IsInstanceOfType(stateObject))
-                throw new ArgumentException("Object provided was not a DbConnectionInternal", "internalConnection");
+                throw new ArgumentException("Object provided was not a TdsParserStateObjectManaged", nameof(stateObject));
         }
 
         internal static object GetSessionHandle(object stateObject)

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/ConnectionPoolTest/ConnectionPoolTest.Debug.cs

@@ -0,0 +1,193 @@
+using System;
+using System.Runtime.ExceptionServices;
+using System.Threading;
+using System.Threading.Tasks;
+using Xunit;
+using static Microsoft.Data.SqlClient.ManualTesting.Tests.ConnectionPoolTest;
+
+namespace Microsoft.Data.SqlClient.ManualTesting.Tests
+{
+    public static class ConnectionPoolTestDebug
+    {
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsUsingManagedSNI))]
+        [ClassData(typeof(ConnectionPoolConnectionStringProvider))]
+        public static void ReplacementConnectionUsesSemaphoreTest(string connectionString)
+        {
+            string newConnectionString = (new SqlConnectionStringBuilder(connectionString) { MaxPoolSize = 2, ConnectTimeout = 5 }).ConnectionString;
+            SqlConnection.ClearAllPools();
+
+            using SqlConnection liveConnection = new(newConnectionString);
+            using SqlConnection deadConnection = new(newConnectionString);
+            liveConnection.Open();
+            deadConnection.Open();
+            InternalConnectionWrapper deadConnectionInternal = new(deadConnection);
+            InternalConnectionWrapper liveConnectionInternal = new(liveConnection);
+            deadConnectionInternal.KillConnection();
+            deadConnection.Close();
+            liveConnection.Close();
+
+            Task<InternalConnectionWrapper>[] tasks = new Task<InternalConnectionWrapper>[3];
+            Barrier syncBarrier = new(tasks.Length);
+            Func<InternalConnectionWrapper> taskFunction = (() => ReplacementConnectionUsesSemaphoreTask(newConnectionString, syncBarrier));
+            for (int i = 0; i < tasks.Length; i++)
+            {
+                tasks[i] = Task.Factory.StartNew(taskFunction);
+            }
+
+            bool taskWithLiveConnection = false;
+            bool taskWithNewConnection = false;
+            bool taskWithCorrectException = false;
+
+            Task waitAllTask = Task.Factory.ContinueWhenAll(tasks, (completedTasks) =>
+            {
+                foreach (var item in completedTasks)
+                {
+                    if (item.Status == TaskStatus.Faulted)
+                    {
+                        // One task should have a timeout exception
+                        if ((!taskWithCorrectException) && (item.Exception.InnerException is InvalidOperationException) && (item.Exception.InnerException.Message.StartsWith(SystemDataResourceManager.Instance.ADP_PooledOpenTimeout)))
+                            taskWithCorrectException = true;
+                        else if (!taskWithCorrectException)
+                        {
+                            // Rethrow the unknown exception
+                            ExceptionDispatchInfo exceptionInfo = ExceptionDispatchInfo.Capture(item.Exception);
+                            exceptionInfo.Throw();
+                        }
+                    }
+                    else if (item.Status == TaskStatus.RanToCompletion)
+                    {
+                        // One task should get the live connection
+                        if (item.Result.Equals(liveConnectionInternal))
+                        {
+                            if (!taskWithLiveConnection)
+                                taskWithLiveConnection = true;
+                        }
+                        else if (!item.Result.Equals(deadConnectionInternal) && !taskWithNewConnection)
+                            taskWithNewConnection = true;
+                    }
+                    else
+                        Console.WriteLine("ERROR: Task in unknown state: {0}", item.Status);
+                }
+            });
+
+            waitAllTask.Wait();
+            Assert.True(taskWithLiveConnection && taskWithNewConnection && taskWithCorrectException,
+                $"Tasks didn't finish as expected.\n" +
+                $"Task with live connection: {taskWithLiveConnection}\n" +
+                $"Task with new connection: {taskWithNewConnection}\n" +
+                $"Task with correct exception: {taskWithCorrectException}\n");
+        }
+
+        /// <summary>
+        /// Tests if killing the connection using the InternalConnectionWrapper is working
+        /// </summary>
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsUsingManagedSNI))]
+        [ClassData(typeof(ConnectionPoolConnectionStringProvider))]
+        public static void KillConnectionTest(string connectionString)
+        {
+            InternalConnectionWrapper wrapper = null;
+
+            using (SqlConnection connection = new(connectionString))
+            {
+                connection.Open();
+                wrapper = new InternalConnectionWrapper(connection);
+
+                using SqlCommand command = new("SELECT 5;", connection);
+
+                DataTestUtility.AssertEqualsWithDescription(5, command.ExecuteScalar(), "Incorrect scalar result.");
+
+                wrapper.KillConnection();
+            }
+
+            using (SqlConnection connection2 = new(connectionString))
+            {
+                connection2.Open();
+                Assert.False(wrapper.IsInternalConnectionOf(connection2), "New connection has internal connection that was just killed");
+                using SqlCommand command = new("SELECT 5;", connection2);
+
+                DataTestUtility.AssertEqualsWithDescription(5, command.ExecuteScalar(), "Incorrect scalar result.");
+            }
+        }
+
+        /// <summary>
+        /// Tests that cleanup removes connections that are unused for two cleanups
+        /// </summary>
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsUsingManagedSNI))]
+        [ClassData(typeof(ConnectionPoolConnectionStringProvider))]
+        public static void CleanupTest(string connectionString)
+        {
+            SqlConnection.ClearAllPools();
+
+            using SqlConnection conn1 = new(connectionString);
+            using SqlConnection conn2 = new(connectionString);
+            conn1.Open();
+            conn2.Open();
+            ConnectionPoolWrapper connectionPool = new(conn1);
+            Assert.Equal(2, connectionPool.ConnectionCount);
+
+            connectionPool.Cleanup();
+            Assert.Equal(2, connectionPool.ConnectionCount);
+
+            conn1.Close();
+            connectionPool.Cleanup();
+            Assert.Equal(2, connectionPool.ConnectionCount);
+
+            conn2.Close();
+            connectionPool.Cleanup();
+            Assert.Equal(1, connectionPool.ConnectionCount);
+
+            connectionPool.Cleanup();
+            Assert.Equal(0, connectionPool.ConnectionCount);
+
+            using SqlConnection conn3 = new(connectionString);
+            conn3.Open();
+            InternalConnectionWrapper internalConnection3 = new(conn3);
+
+            conn3.Close();
+            internalConnection3.KillConnection();
+            Assert.Equal(1, connectionPool.ConnectionCount);
+            Assert.False(internalConnection3.IsConnectionAlive(), "Connection should not be alive");
+
+            connectionPool.Cleanup();
+            Assert.Equal(1, connectionPool.ConnectionCount);
+        }
+
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsUsingManagedSNI))]
+        [ClassData(typeof(ConnectionPoolConnectionStringProvider))]
+        public static void ReplacementConnectionObeys0TimeoutTest(string connectionString)
+        {
+            string newConnectionString = (new SqlConnectionStringBuilder(connectionString) { ConnectTimeout = 0 }).ConnectionString;
+            SqlConnection.ClearAllPools();
+
+            // Kick off proxy
+            using (ProxyServer proxy = ProxyServer.CreateAndStartProxy(newConnectionString, out newConnectionString))
+            {
+                // Create one dead connection
+                using SqlConnection deadConnection = new(newConnectionString);
+                deadConnection.Open();
+                InternalConnectionWrapper deadConnectionInternal = new(deadConnection);
+                deadConnectionInternal.KillConnection();
+
+                // Block one live connection
+                proxy.PauseCopying();
+                Task<SqlConnection> blockedConnectionTask = Task.Run(() => ReplacementConnectionObeys0TimeoutTask(newConnectionString));
+                Thread.Sleep(100);
+                Assert.Equal(TaskStatus.Running, blockedConnectionTask.Status);
+
+                // Close and re-open the dead connection
+                deadConnection.Close();
+                Task<SqlConnection> newConnectionTask = Task.Run(() => ReplacementConnectionObeys0TimeoutTask(newConnectionString));
+                Thread.Sleep(100);
+                Assert.Equal(TaskStatus.Running, blockedConnectionTask.Status);
+                Assert.Equal(TaskStatus.Running, newConnectionTask.Status);
+
+                // restart the proxy
+                proxy.ResumeCopying();
+
+                Task.WaitAll(blockedConnectionTask, newConnectionTask);
+                blockedConnectionTask.Result.Close();
+                newConnectionTask.Result.Close();
+            }
+        }
+    }
+}

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/ConnectionPoolTest/TransactionPoolTest.cs

@@ -0,0 +1,99 @@
+using System.Transactions;
+using Xunit;
+
+namespace Microsoft.Data.SqlClient.ManualTesting.Tests
+{
+    public static class TransactionPoolTest
+    {
+        /// <summary>
+        /// Tests if connections in a distributed transaction are put into a transaction pool. Also checks that clearallpools 
+        /// does not clear transaction connections and that the transaction root is put into "stasis" when closed
+        /// </summary>
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
+        [ClassData(typeof(ConnectionPoolConnectionStringProvider))]
+        public static void BasicTransactionPoolTest(string connectionString)
+        {
+            SqlConnection.ClearAllPools();
+            ConnectionPoolWrapper connectionPool = null;
+
+            using (TransactionScope transScope = new())
+            {
+                using SqlConnection connection1 = new(connectionString);
+                using SqlConnection connection2 = new(connectionString);
+                connection1.Open();
+                connection2.Open();
+                connectionPool = new ConnectionPoolWrapper(connection1);
+
+                InternalConnectionWrapper internalConnection1 = new(connection1);
+                InternalConnectionWrapper internalConnection2 = new(connection2);
+
+                Assert.True(internalConnection1.IsEnlistedInTransaction, "First connection not in transaction");
+                Assert.True(internalConnection1.IsTransactionRoot, "First connection not transaction root");
+                Assert.True(internalConnection2.IsEnlistedInTransaction, "Second connection not in transaction");
+                Assert.False(internalConnection2.IsTransactionRoot, "Second connection is transaction root");
+
+                // Attempt to re-use root connection
+                connection1.Close();
+                using SqlConnection connection3 = new(connectionString);
+                connection3.Open();
+
+                Assert.True(connectionPool.ContainsConnection(connection3), "New connection in wrong pool");
+                Assert.True(internalConnection1.IsInternalConnectionOf(connection3), "Root connection was not re-used");
+
+                // Attempt to re-use non-root connection
+                connection2.Close();
+                using SqlConnection connection4 = new(connectionString);
+                connection4.Open();
+                Assert.True(internalConnection2.IsInternalConnectionOf(connection4), "Connection did not re-use expected internal connection");
+                Assert.True(connectionPool.ContainsConnection(connection4), "New connection is in the wrong pool");
+                connection4.Close();
+
+                // Use a different connection string
+                using SqlConnection connection5 = new(connectionString + ";App=SqlConnectionPoolUnitTest;");
+                connection5.Open();
+                Assert.False(internalConnection2.IsInternalConnectionOf(connection5), "Connection with different connection string re-used internal connection");
+                Assert.False(connectionPool.ContainsConnection(connection5), "Connection with different connection string is in same pool");
+                connection5.Close();
+
+                transScope.Complete();
+            }
+
+            Assert.Equal(2, connectionPool.ConnectionCount);
+        }
+
+        /// <summary>
+        /// Checks that connections in the transaction pool are not cleaned out, and the root transaction is put into "stasis" when it ages
+        /// </summary>
+        /// <param name="connectionString"></param>
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
+        [ClassData(typeof(ConnectionPoolConnectionStringProvider))]
+        public static void TransactionCleanupTest(string connectionString)
+        {
+            SqlConnection.ClearAllPools();
+            ConnectionPoolWrapper connectionPool = null;
+
+            using (TransactionScope transScope = new())
+            {
+                using SqlConnection connection1 = new(connectionString);
+                using SqlConnection connection2 = new(connectionString);
+                connection1.Open();
+                connection2.Open();
+                InternalConnectionWrapper internalConnection1 = new(connection1);
+                connectionPool = new ConnectionPoolWrapper(connection1);
+
+                connectionPool.Cleanup();
+                Assert.Equal(2, connectionPool.ConnectionCount);
+
+                connection1.Close();
+                connection2.Close();
+                connectionPool.Cleanup();
+                Assert.Equal(2, connectionPool.ConnectionCount);
+
+                connectionPool.Cleanup();
+                Assert.Equal(2, connectionPool.ConnectionCount);
+
+                transScope.Complete();
+            }
+        }
+    }
+}