Preserve distributed transactions on pooled connection reset

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

@@ -946,11 +946,11 @@ private void ResetConnection()
 
             if (_fResetConnection)
             {
-                // Ensure we are either going against 2000, or we are not enlisted in a
-                // 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);
+                // Pooled connections that are enlisted in a transaction must have their transaction
+                // preserved when reseting the connection state. Otherwise, future uses of the connection
+                // from the pool will execute outside of the transaction, in auto-commit mode.
+                // https://github.com/dotnet/SqlClient/issues/2970
+                _parser.PrepareResetConnection(EnlistedTransaction is not null && Pool is not 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

@@ -1075,9 +1075,11 @@ private void ResetConnection()
                 // distributed transaction - otherwise don't reset!
                 if (Is2000)
                 {
-                    // Prepare the parser for the connection reset - the next time a trip
-                    // to the server is made.
-                    _parser.PrepareResetConnection(IsTransactionRoot && !IsNonPoolableTransactionRoot);
+                    // Pooled connections that are enlisted in a transaction must have their transaction
+                    // preserved when reseting the connection state. Otherwise, future uses of the connection
+                    // from the pool will execute outside of the transaction, in auto-commit mode.
+                    // https://github.com/dotnet/SqlClient/issues/2970
+                    _parser.PrepareResetConnection(EnlistedTransaction is not null && Pool is not 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

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

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

@@ -0,0 +1,181 @@
+// 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.Data;
+using System.Runtime.InteropServices;
+using System.Threading.Tasks;
+using System.Transactions;
+using Microsoft.Data.SqlClient.TestUtilities;
+using Xunit;
+
+namespace Microsoft.Data.SqlClient.ManualTesting.Tests
+{
+
+    [PlatformSpecific(TestPlatforms.Windows)]
+    public class DistributedTransactionTestWindows
+    {
+
+#if NET
+        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
+
+        private static bool s_EnlistedTransactionPreservedWhilePooledCondition => DataTestUtility.AreConnStringsSetup() && DataTestUtility.IsNotX86Architecture;
+
+        [ConditionalFact(nameof(s_EnlistedTransactionPreservedWhilePooledCondition), Timeout = 10000)]
+        public void Test_EnlistedTransactionPreservedWhilePooled()
+        {
+#if NET
+            TransactionManager.ImplicitDistributedTransactions = true;
+#endif
+            RunTestSet(EnlistedTransactionPreservedWhilePooled);
+        }
+
+        private void EnlistedTransactionPreservedWhilePooled()
+        {
+            Exception commandException = null;
+            Exception transactionException = null;
+
+            try
+            {
+                using (TransactionScope txScope = new TransactionScope(TransactionScopeOption.Required, TimeSpan.MaxValue))
+                {
+                    // Leave first connection open so that the transaction is promoted
+                    SqlConnection rootConnection = new SqlConnection(ConnectionString);
+                    rootConnection.Open();
+                    using (SqlCommand command = rootConnection.CreateCommand())
+                    {
+                        command.CommandText = $"INSERT INTO {TestTableName} VALUES ({InputCol1}, '{InputCol2}')";
+                        command.ExecuteNonQuery();
+                    }
+
+                    // Closing and reopening cycles the connection through the pool.
+                    // We want to verify that the transaction state is preserved through this cycle.
+                    SqlConnection enlistedConnection = new SqlConnection(ConnectionString);
+                    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
+                    // Can also represent a general server-side, process failure
+                    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;
+            }
+
+            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);
+                // 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
+            {
+                Assert.IsType<TransactionAbortedException>(transactionException);
+
+#if NETFRAMEWORK
+                // 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);
+#else
+                Assert.IsType<InvalidOperationException>(commandException);
+#endif
+            }
+
+            // 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)
+        {
+            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}");
+            }
+        }
+    }
+}
+