Preserve distributed transactions on pooled connection reset

src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs

@@ -950,7 +950,7 @@ private void ResetConnection()
                 // distributed transaction - otherwise don't reset!
                 // Prepare the parser for the connection reset - the next time a trip
                 // to the server is made.
-                _parser.PrepareResetConnection(IsTransactionRoot && !IsNonPoolableTransactionRoot);
+                _parser.PrepareResetConnection(EnlistedTransaction != null && Pool != null);
 
                 // Reset dictionary values, since calling reset will not send us env_changes.
                 CurrentDatabase = _originalDatabase;

src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlInternalConnectionTds.cs

@@ -1077,7 +1077,7 @@ private void ResetConnection()
                 {
                     // Prepare the parser for the connection reset - the next time a trip
                     // to the server is made.
-                    _parser.PrepareResetConnection(IsTransactionRoot && !IsNonPoolableTransactionRoot);
+                    _parser.PrepareResetConnection(EnlistedTransaction != null && Pool != null);
                 }
                 else if (!IsEnlistedInTransaction)
                 {

src/Microsoft.Data.SqlClient/src/Microsoft/Data/ProviderBase/DbConnectionPool.cs

@@ -1805,6 +1805,8 @@ internal void PutObjectFromTransactedPool(DbConnectionInternal obj)
             Debug.Assert(obj != null, "null pooledObject?");
             Debug.Assert(obj.EnlistedTransaction == null, "pooledObject is still enlisted?");
 
+            obj.DeactivateConnection();
+
             // called by the transacted connection pool , once it's removed the
             // connection from it's list.  We put the connection back in general
             // circulation.

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

@@ -202,6 +202,7 @@
     <Compile Include="SQL\SqlNotificationTest\SqlNotificationTest.cs" />
     <Compile Include="SQL\SqlSchemaInfoTest\SqlSchemaInfoTest.cs" />
     <Compile Include="SQL\SqlStatisticsTest\SqlStatisticsTest.cs" />
+    <Compile Include="SQL\TransactionTest\DistributedTransactionTest.Windows.cs" />
     <Compile Include="SQL\TransactionTest\DistributedTransactionTest.cs" />
     <Compile Include="SQL\TransactionTest\TransactionTest.cs" />
     <Compile Include="SQL\TransactionTest\TransactionEnlistmentTest.cs" />

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/TransactionTest/DistributedTransactionTest.Windows.cs

@@ -0,0 +1,46 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Runtime.InteropServices;
+using System.Threading.Tasks;
+using System.Transactions;
+using Xunit;
+
+#if NET8_0_OR_GREATER
+
+namespace Microsoft.Data.SqlClient.ManualTesting.Tests
+{
+    [PlatformSpecific(TestPlatforms.Windows)]
+    public class DistributedTransactionTestWindows
+    {
+        private static bool s_DelegatedTransactionCondition => DataTestUtility.AreConnStringsSetup() && DataTestUtility.IsNotAzureServer() && DataTestUtility.IsNotX86Architecture;
+
+        [ConditionalFact(nameof(s_DelegatedTransactionCondition), Timeout = 10000)]
+        public async Task Delegated_transaction_deadlock_in_SinglePhaseCommit()
+        {
+            TransactionManager.ImplicitDistributedTransactions = true;
+            using var transaction = new CommittableTransaction();
+
+            // Uncommenting the following makes the deadlock go away as a workaround. If the transaction is promoted before
+            // the first SqlClient enlistment, it never goes into the delegated state.
+            // _ = TransactionInterop.GetTransmitterPropagationToken(transaction);
+            await using var conn = new SqlConnection(DataTestUtility.TCPConnectionString);
+            await conn.OpenAsync();
+            conn.EnlistTransaction(transaction);
+
+            // Enlisting the transaction in second connection causes the transaction to be promoted.
+            // After this, the transaction state will be "delegated" (delegated to SQL Server), and the commit below will
+            // trigger a call to SqlDelegatedTransaction.SinglePhaseCommit.
+            await using var conn2 = new SqlConnection(DataTestUtility.TCPConnectionString);
+            await conn2.OpenAsync();
+            conn2.EnlistTransaction(transaction);
+
+            // Possible deadlock
+            transaction.Commit();
+        }
+    }
+}
+
+#endif

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/TransactionTest/DistributedTransactionTest.cs

@@ -1,46 +1,139 @@
-// Licensed to the .NET Foundation under one or more agreements.
-// The .NET Foundation licenses this file to you under the MIT license.
-// See the LICENSE file in the project root for more information.
-
+using System.Data;
 using System;
-using System.Runtime.InteropServices;
-using System.Threading.Tasks;
 using System.Transactions;
 using Xunit;
-
-#if NET8_0_OR_GREATER
+using Microsoft.Data.SqlClient.TestUtilities;
 
 namespace Microsoft.Data.SqlClient.ManualTesting.Tests
 {
-    [PlatformSpecific(TestPlatforms.Windows)]
     public class DistributedTransactionTest
     {
-        private static bool s_DelegatedTransactionCondition => DataTestUtility.AreConnStringsSetup() && DataTestUtility.IsNotAzureServer() && DataTestUtility.IsNotX86Architecture;
+        [ConditionalFact(nameof(DataTestUtility.AreConnStringsSetup), Timeout = 10000)]
+        public void Test_EnlistedTransactionPreservedWhilePooled()
+        {
+            RunTestSet(EnlistedTransactionPreservedWhilePooled);
+        }
+
+        private void EnlistedTransactionPreservedWhilePooled()
+        {
+            Exception commandException = null;
+            Exception transactionException = null;
+
+            try
+            {
+                using (TransactionScope txScope = new TransactionScope(TransactionScopeOption.Required, TimeSpan.MaxValue))
+                {
+                    SqlConnection rootConnection = new SqlConnection(ConnectionString);
+                    rootConnection.Open();
+                    using (SqlCommand command = rootConnection.CreateCommand())
+                    {
+                        command.CommandText = $"INSERT INTO {TestTableName} VALUES ({InputCol1}, '{InputCol2}')";
+                        command.ExecuteNonQuery();
+                    }
+                    // Leave first connection open so that the transaction is promoted
+
+                    SqlConnection enlistedConnection = new SqlConnection(ConnectionString);
+                    // Closing and reopening cycles the connection through the pool.
+                    // We want to verify that the transaction state is preserved through this cycle.
+                    enlistedConnection.Open();
+                    enlistedConnection.Close();
+                    enlistedConnection.Open();
+
+                    // Forcibly kill the root connection to mimic gateway's behavior when using the proxy connection policy
+                    // https://techcommunity.microsoft.com/blog/azuredbsupport/azure-sql-database-idle-sessions-are-killed-after-about-30-minutes-when-proxy-co/3268601
+                    KillProcess(rootConnection.ServerProcessId);
+
+
+                    using (SqlCommand command = enlistedConnection.CreateCommand())
+                    {
+                        command.CommandText = $"INSERT INTO {TestTableName} VALUES ({InputCol1}, '{InputCol2}')";
+                        try
+                        {
+                            command.ExecuteNonQuery();
+                        }
+                        catch (Exception ex)
+                        {
+                            commandException = ex;
+                        }
+                    }
+
+                    txScope.Complete();
+                }
+            } catch (Exception ex)
+            {
+                transactionException = ex;
+            }
 
-        [ConditionalFact(nameof(s_DelegatedTransactionCondition), Timeout = 10000)]
-        public async Task Delegated_transaction_deadlock_in_SinglePhaseCommit()
+
+            Assert.IsType<SqlException>(commandException);
+            if (Utils.IsAzureSqlServer(new SqlConnectionStringBuilder((ConnectionString)).DataSource))
+            {
+                // See https://learn.microsoft.com/en-us/sql/relational-databases/errors-events/database-engine-events-and-errors-3000-to-3999?view=sql-server-ver16
+                // Error 3971 corresponds to "The server failed to resume the transaction."
+                Assert.Equal(3971, ((SqlException)commandException).Number);
+            } else
+            {
+                // See https://learn.microsoft.com/en-us/sql/relational-databases/errors-events/database-engine-events-and-errors-8000-to-8999?view=sql-server-ver16
+                // The distributed transaction failed
+                Assert.Equal(8525, ((SqlException)commandException).Number);
+            }
+
+            if (Utils.IsAzureSqlServer(new SqlConnectionStringBuilder((ConnectionString)).DataSource))
+            {
+                // Even if an application swallows the command exception, completing the transaction should indicate that it failed.
+                Assert.IsType<TransactionInDoubtException>(transactionException);
+            }
+            else
+            {
+                Assert.IsType<TransactionAbortedException>(transactionException);
+            }
+
+            // Verify that nothing made it into the database
+            DataTable result = DataTestUtility.RunQuery(ConnectionString, $"select col2 from {TestTableName} where col1 = {InputCol1}");
+            Assert.True(result.Rows.Count == 0);
+        }
+
+        private void KillProcess(int serverProcessId)
+        {
+            using (TransactionScope txScope = new TransactionScope(TransactionScopeOption.Suppress))
+            {
+                using (SqlConnection connection = new SqlConnection(ConnectionString))
+                {
+                    connection.Open();
+                    using (SqlCommand command = connection.CreateCommand())
+                    {
+                        command.CommandText = $"KILL {serverProcessId}";
+                        command.ExecuteNonQuery();
+                    }
+                }
+                txScope.Complete();
+            }
+        }
+
+        private static string TestTableName;
+        private static string ConnectionString;
+        private const int InputCol1 = 1;
+        private const string InputCol2 = "One";
+
+        private static void RunTestSet(Action TestCase)
         {
-            TransactionManager.ImplicitDistributedTransactions = true;
-            using var transaction = new CommittableTransaction();
-
-            // Uncommenting the following makes the deadlock go away as a workaround. If the transaction is promoted before
-            // the first SqlClient enlistment, it never goes into the delegated state.
-            // _ = TransactionInterop.GetTransmitterPropagationToken(transaction);
-            await using var conn = new SqlConnection(DataTestUtility.TCPConnectionString);
-            await conn.OpenAsync();
-            conn.EnlistTransaction(transaction);
-
-            // Enlisting the transaction in second connection causes the transaction to be promoted.
-            // After this, the transaction state will be "delegated" (delegated to SQL Server), and the commit below will
-            // trigger a call to SqlDelegatedTransaction.SinglePhaseCommit.
-            await using var conn2 = new SqlConnection(DataTestUtility.TCPConnectionString);
-            await conn2.OpenAsync();
-            conn2.EnlistTransaction(transaction);
-
-            // Possible deadlock
-            transaction.Commit();
+            SqlConnectionStringBuilder builder = new SqlConnectionStringBuilder(DataTestUtility.TCPConnectionString);
+
+            builder.Pooling = true;
+            builder.MaxPoolSize = 5;
+            builder.Enlist = true;
+            ConnectionString = builder.ConnectionString;
+
+            TestTableName = DataTestUtility.GenerateObjectName();
+            DataTestUtility.RunNonQuery(ConnectionString, $"create table {TestTableName} (col1 int, col2 text)");
+            try
+            {
+                TestCase();
+            }
+            finally
+            {
+                DataTestUtility.RunNonQuery(ConnectionString, $"drop table {TestTableName}");
+            }
         }
     }
 }
-
-#endif