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 index b2d4427b47..2e52feb9a4 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs @@ -5,6 +5,7 @@ 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; @@ -250,6 +251,14 @@ public ConcurrentDictionary< private int MinPoolSize => PoolGroupOptions.MinPoolSize; + /// + /// Indicates whether connections may be vended from (and parked in) the + /// . This mirrors automatic transaction enlistment: + /// when enlistment is disabled a connection is never bound to an ambient transaction, so + /// the transacted store must not be consulted. + /// + private bool HasTransactionAffinity => PoolGroupOptions.HasTransactionAffinity; + /// /// The most recently launched warmup/replenishment loop task, exposed so tests can await a /// warmup pass to a deterministic completion instead of polling pool counters. May be null @@ -313,7 +322,35 @@ public void Clear() /// public void PutObjectFromTransactedPool(DbConnectionInternal connection) { - throw new NotImplementedException(); + Debug.Assert(connection is not null, + "PutObjectFromTransactedPool was called with a null connection."); + Debug.Assert(connection.EnlistedTransaction is null, + "PutObjectFromTransactedPool was called with a connection that is still enlisted. " + + "The transaction must have ended and been detached before the connection returns to " + + "general circulation, otherwise it could be vended to a caller in a different transaction."); + + // Called by the transacted connection pool once it has removed the connection from its + // list. We put the connection back into general circulation. + // + // NOTE: no locking is required here because if we're in this method we can safely + // presume that the caller is the only one using the connection, that all pre-push logic + // has been done, and that all transactions have ended. + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Connection {1}, Transaction has ended.", + Id, + connection.ObjectID); + + if (State is Running && connection.CanBePooled) + { + connection.ResetConnection(); + PutConnectionInIdleChannel(connection); + } + else + { + // RemoveConnection triggers replenishment, which is the channel pool's equivalent + // of the wait handle pool's QueuePoolCreateRequest. + RemoveConnection(connection); + } } /// @@ -331,7 +368,8 @@ public DbConnectionInternal ReplaceConnection( if (newConnection is not null) { - // TODO: Full transaction enlistment support (Story 2). + // Carry the old connection's enlistment over to the replacement so that a connection + // replaced mid-transaction stays bound to the same transaction. PrepareConnection(owningObject, newConnection, oldConnection.EnlistedTransaction); oldConnection.DeactivateConnection(); RemoveConnection(oldConnection); @@ -373,7 +411,8 @@ public DbConnectionInternal ReplaceConnection( newConnection.PostPop(owningObject); } - // TODO: Full transaction enlistment support (Story 2). + // Carry the old connection's enlistment over to the replacement so that a + // connection replaced mid-transaction stays bound to the same transaction. newConnection.ActivateConnection(oldConnection.EnlistedTransaction); // Place new into old's slot @@ -414,12 +453,149 @@ public void ReturnInternalConnection(DbConnectionInternal connection, DbConnecti { ValidateOwnershipAndSetPoolingState(connection, owningObject); + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Connection {1}, Deactivating.", + Id, + connection.ObjectID); + + // Deactivate before inspecting the connection's transaction state. Deactivation is what + // detaches a completed transaction, so reading EnlistedTransaction beforehand could park + // a connection in the transacted pool under a transaction that has already ended. + connection.DeactivateConnection(); + + if (connection.IsConnectionDoomed) + { + // The connection is not fit for reuse -- just dispose of it. + RemoveConnection(connection); + return; + } + + // The disposition is decided while holding the connection's lock so that a transaction + // completing asynchronously on another thread cannot race with us, but the resulting + // I/O is performed outside of it. + ReturnDisposition disposition; + lock (connection) + { + disposition = DecideReturnDisposition(connection); + } + + switch (disposition) + { + case ReturnDisposition.Reuse: + PutConnectionInIdleChannel(connection); + break; + + case ReturnDisposition.Destroy: + RemoveConnection(connection); + break; + + case ReturnDisposition.HeldByTransaction: + // Nothing further to do. The connection is parked in the transacted store or + // in stasis, and comes back through PutObjectFromTransactedPool once its + // transaction ends. + break; + } + } + + /// + /// Decides what should happen to a returning connection that has already been deactivated + /// and is not doomed. Must be called while holding the connection's lock: parking a + /// connection in the transacted store has to be atomic with respect to a transaction + /// completing on another thread. + /// + /// The connection being returned. + /// The action the caller must take for this connection. + private ReturnDisposition DecideReturnDisposition(DbConnectionInternal connection) + { + // A connection is only fit to go back into circulation if the pool is still running, + // the connection itself is still poolable, and it isn't a transaction root that has + // already been detached from its pool. + bool isReusable = + State is Running && + connection.CanBePooled && + !(connection.IsTransactionRoot && connection.Pool is null); + + if (!isReusable) + { + // A transaction root that cannot be pooled must be put in stasis rather than + // closed. Closing it would orphan the root transaction with no means to promote + // itself to a full delegated transaction, or to commit or roll back. + // System.Transactions keeps the connection owned (not lost) and is certain to + // hand it back to us when the transaction ends. + if (connection.IsTransactionRoot) + { + connection.SetInStasis(); + return ReturnDisposition.HeldByTransaction; + } + + return ReturnDisposition.Destroy; + } + + // A connection that is still enlisted cannot be handed to a different customer until + // its transaction actually completes, so it is parked in the transacted store keyed by + // that transaction and comes back via PutObjectFromTransactedPool when it ends. + // + // NOTE: we do not hold a lock on State, so its value could have changed since the + // check above. That is acceptable because the DelegatedTransactionEnded event cleans + // the connection up appropriately regardless of the pool state. + // + // Transacted connections are deliberately not stamped with a returned time: they are + // never proactively closed (doing so would abort a possibly distributed transaction), + // so idle-timeout enforcement does not apply while they are parked. They are stamped + // when they rejoin the idle channel in PutObjectFromTransactedPool. + Transaction? transaction = connection.EnlistedTransaction; + if (transaction is not null) + { + TransactedConnectionPool.PutTransactedObject(transaction, connection); + return ReturnDisposition.HeldByTransaction; + } + + return ReturnDisposition.Reuse; + } + + /// + /// The outcome of evaluating a connection that is being returned to the pool. + /// + private enum ReturnDisposition + { + /// + /// The connection is fit for general reuse and belongs in the idle channel. + /// + Reuse, + + /// + /// The connection cannot be reused and must be closed. + /// + Destroy, + + /// + /// The connection is owned by a live transaction, either parked in the transacted + /// store or held in stasis, and must not be touched by the pool until that + /// transaction ends. + /// + HeldByTransaction, + } + + /// + /// Places a connection that is fit for general reuse into the idle channel, stamping its + /// idle-return time and dropping it if it is no longer live. + /// + /// The connection to make available to other callers. + private void PutConnectionInIdleChannel(DbConnectionInternal connection) + { // Stamp the return time before IsLiveConnection runs so the idle-expiry gate inside it // measures time-in-pool, not time-since-last-return. Without this, a connection whose // checkout exceeded IdleTimeout (e.g. a long-running query) would be wrongly evicted on // return even though it was actively in use on the wire. The same gating conditions are // applied here as in IsLiveConnection so we avoid the per-return timestamp read when // idle expiry is disabled or the legacy idle-timeout behavior is in effect. + // + // A connection parked in the transacted store does not pass through here, so it is not + // subject to idle timeout for as long as its transaction is live. That is intentional + // and matches WaitHandleDbConnectionPool: closing it would abort a possibly distributed + // transaction. It is stamped here when the transaction ends and + // PutObjectFromTransactedPool returns it to general circulation, so the idle clock + // starts from the moment it actually becomes available to other callers. if (!LocalAppContextSwitches.UseLegacyIdleTimeoutBehavior && PoolGroupOptions.IdleTimeout != TimeSpan.Zero) { @@ -432,27 +608,12 @@ public void ReturnInternalConnection(DbConnectionInternal connection, DbConnecti return; } - SqlClientEventSource.Log.TryPoolerTraceEvent( - " {0}, Connection {1}, Deactivating.", - Id, - connection.ObjectID); - connection.DeactivateConnection(); - - if (connection.IsConnectionDoomed || - !connection.CanBePooled || - State == ShuttingDown) + if (!_idleChannel.TryWrite(connection)) { + // The channel has been completed (pool is shutting down). Race window + // between the State check by the caller and TryWrite: destroy instead of pooling. RemoveConnection(connection); } - else - { - 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); - } - } } /// @@ -606,7 +767,21 @@ public void Startup() /// public void TransactionEnded(Transaction transaction, DbConnectionInternal transactedObject) { - throw new NotImplementedException(); + Debug.Assert(transaction is not null, "null transaction?"); + Debug.Assert(transactedObject is not null, "null transactedObject?"); + + // Note: the connection may still be associated with the transaction due to the explicit + // unbinding requirement. + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Transaction {1}, Connection {2}, Transaction Completed", + Id, + transaction.GetHashCode(), + transactedObject.ObjectID); + + // If the connection is in the transacted pool, remove it there and return it to general + // circulation. TransactedConnectionPool.TransactionEnded calls back into + // PutObjectFromTransactedPool for us once it has removed the connection from its list. + TransactedConnectionPool.TransactionEnded(transaction, transactedObject); } /// @@ -632,10 +807,12 @@ public bool TryGetConnection( // If taskCompletionSource is null, we are in a sync context. if (taskCompletionSource is null) { + // We're on the caller's thread, so the ambient transaction is directly observable. var task = GetInternalConnection( owningObject, async: false, - timeout); + timeout, + ADP.GetCurrentTransaction()); // When running synchronously, we are guaranteed that the task is already completed. // We don't need to guard the managed threadpool at this spot because we pass the async flag as false @@ -664,6 +841,22 @@ public bool TryGetConnection( // OpenAsync call. This means that we cannot cancel the connection open operation if the caller's token // is cancelled. We can only cancel based on our own timeout, which is set to the owningObject's // ConnectionTimeout. + // The ambient transaction is captured here, on the caller's thread, because + // Transaction.Current does not flow into the Task.Run below: a TransactionScope keeps + // the ambient transaction in thread-static storage unless it was created with + // TransactionScopeAsyncFlowOption.Enabled. We rely on the caller to capture the ambient + // transaction in the TaskCompletionSource's AsyncState, and then hand it to + // GetInternalConnection explicitly. + // + // Note that we deliberately do not assign Transaction.Current on the thread pool + // thread. That assignment writes to thread-static storage which is *not* unwound when + // the ExecutionContext is restored, so it would outlive this open and be observed by + // unrelated work later scheduled onto the same thread pool thread -- including the + // login-time auto-enlistment that non-pooled connections perform against + // Transaction.Current. The WaitHandle pool can get away with assigning it because it + // processes pending opens on a dedicated non-thread-pool thread. + Transaction? ambientTransaction = taskCompletionSource.Task.AsyncState as Transaction; + Task.Run(async () => { if (taskCompletionSource.Task.IsCompleted) @@ -671,10 +864,6 @@ public bool TryGetConnection( return; } - // We're potentially on a new thread, so we need to properly set the ambient transaction. - // We rely on the caller to capture the ambient transaction in the TaskCompletionSource's AsyncState - // so that we can access it here. Read: area for improvement. - // TODO: ADP.SetCurrentTransaction(taskCompletionSource.Task.AsyncState as Transaction); DbConnectionInternal? connection = null; try @@ -682,7 +871,8 @@ public bool TryGetConnection( connection = await GetInternalConnection( owningObject, async: true, - timeout + timeout, + ambientTransaction ).ConfigureAwait(false); if (!taskCompletionSource.TrySetResult(connection)) @@ -941,6 +1131,19 @@ private bool IsLiveConnection(DbConnectionInternal connection) /// The connection to be closed. private void RemoveConnection(DbConnectionInternal connection) { + // A connection with a delegated transaction cannot be disposed of until the delegated + // transaction has actually completed; disposing it would abort the (possibly + // distributed) transaction. Leave it alone: when the transaction completes it comes + // back through PutObjectFromTransactedPool, which calls us again. + if (connection.IsTxRootWaitingForTxEnd) + { + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Connection {1}, Has Delegated Transaction, waiting to Dispose.", + Id, + connection.ObjectID); + return; + } + _connectionSlots.TryRemove(connection); // Removing a connection from the pool opens a free slot. @@ -995,6 +1198,10 @@ private void RemoveConnection(DbConnectionInternal connection) /// A boolean indicating whether the operation should be asynchronous. /// The overall timeout budget for this connection request. Time spent waiting /// in the pool is deducted from the budget available for physical connection creation. + /// The ambient transaction captured on the caller's thread, or + /// null when the caller is not inside a transaction. It is passed explicitly rather than read + /// from because this method may run on a thread pool thread + /// that the ambient transaction does not flow to. /// Returns a DbConnectionInternal that is retrieved from the pool. /// /// Thrown when an OperationCanceledException is caught, indicating that the timeout period @@ -1007,20 +1214,46 @@ private void RemoveConnection(DbConnectionInternal connection) private async Task GetInternalConnection( DbConnection owningConnection, bool async, - TimeoutTimer timeout) + TimeoutTimer timeout, + Transaction? ambientTransaction) { DbConnectionInternal? connection = null; + // When automatic enlistment is disabled, the connection must never be bound to the + // ambient transaction, so we neither consult the transacted store nor hand the + // transaction to activation. HasTransactionAffinity is derived from the connection + // string's Enlist keyword. + Transaction? transaction = HasTransactionAffinity ? ambientTransaction : null; + // Derive a CancellationTokenSource from the TimeoutTimer so pool-internal wait operations // (channel reads, semaphore waits) are cancelled when the overall budget expires. using CancellationTokenSource cancellationTokenSource = timeout.CreateCancellationTokenSource(); CancellationToken cancellationToken = cancellationTokenSource.Token; - // Continue looping until we create or retrieve a connection - do + // Continue looping until we create or retrieve a connection. + while (connection is null) { try { + // A connection already enlisted in our transaction is always preferred, since + // reusing it avoids promoting the transaction to a distributed one. This is + // re-checked on every iteration so that a connection returned to the transacted + // store while we were looping is picked up rather than being passed over in + // favor of a fresh connection. + if (transaction is not null) + { + connection = GetFromTransactedPool(transaction); + if (connection is not null) + { + // Skip the liveness/idle/generation gate at the bottom of the loop: + // GetFromTransactedPool has already probed liveness, and a transacted + // connection is exempt from idle-timeout, load-balance and + // clear-generation eviction because closing it would abort its + // (possibly distributed) transaction. + break; + } + } + // Optimistically try to get an idle connection from the channel // Doesn't wait if the channel is empty, just returns null. connection ??= GetIdleConnection(); @@ -1062,9 +1295,8 @@ private async Task GetInternalConnection( connection = null; } } - while (connection is null); - PrepareConnection(owningConnection, connection); + PrepareConnection(owningConnection, connection, transaction); return connection; } @@ -1133,6 +1365,60 @@ private void PrepareConnection(DbConnection owningObject, DbConnectionInternal c } } + /// + /// Attempts to retrieve a connection that is already enlisted in the given transaction. + /// + /// The transaction the connection must already be enlisted in. + /// A live connection already enlisted in the transaction, or null. + private DbConnectionInternal? GetFromTransactedPool(Transaction transaction) + { + DbConnectionInternal? connection = TransactedConnectionPool.GetTransactedObject(transaction); + if (connection is null) + { + return null; + } + + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Connection {1}, Popped from transacted pool.", + Id, + connection.ObjectID); + + SqlClientDiagnostics.Metrics.ExitFreeConnection(); + + // Transacting connections are exempt from idle-timeout and clear-generation eviction + // (closing them would abort the transaction, which may be distributed), so only + // liveness is checked here rather than the full IsLiveConnection gate. + if (connection.IsTransactionRoot) + { + try + { + // A dead transaction root must surface the underlying failure to the caller: + // there is no way to recover the delegated transaction on another connection. + connection.IsConnectionAlive(throwOnException: true); + } + catch + { + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Connection {1}, found dead and removed.", + Id, + connection.ObjectID); + RemoveConnection(connection); + throw; + } + } + else if (!connection.IsConnectionAlive()) + { + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Connection {1}, found dead and removed.", + Id, + connection.ObjectID); + RemoveConnection(connection); + connection = null; + } + + return connection; + } + /// /// Validates that the connection is owned by the provided DbConnection and that it is in a valid state to be returned to the pool. /// diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/ChannelDbConnectionPoolTest.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/ChannelDbConnectionPoolTest.cs index 604964927e..90417c5e96 100644 --- a/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/ChannelDbConnectionPoolTest.cs +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/ChannelDbConnectionPoolTest.cs @@ -887,37 +887,6 @@ public void TestReplaceConnection() #endregion - #region Not Implemented Method Tests - - /// - /// Verifies that remains - /// unimplemented and throws . - /// - [Fact] - public void TestPutObjectFromTransactedPool() - { - // Arrange - var pool = ConstructPool(SuccessfulConnectionFactory); - - // Act & Assert - Assert.Throws(() => pool.PutObjectFromTransactedPool(null!)); - } - - /// - /// Verifies that - /// remains unimplemented and throws . - /// - [Fact] - public void TestTransactionEnded() - { - // Arrange - var pool = ConstructPool(SuccessfulConnectionFactory); - - // Act & Assert - Assert.Throws(() => pool.TransactionEnded(null!, null!)); - } - #endregion - #region Pool Clear Tests /// diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/ChannelDbConnectionPoolTransactionTest.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/ChannelDbConnectionPoolTransactionTest.cs new file mode 100644 index 0000000000..a69941bbbf --- /dev/null +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/ChannelDbConnectionPoolTransactionTest.cs @@ -0,0 +1,709 @@ +// 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.Common; +using System.Threading; +using System.Threading.Tasks; +using System.Transactions; +using Microsoft.Data.Common.ConnectionString; +using Microsoft.Data.ProviderBase; +using Microsoft.Data.SqlClient.ConnectionPool; +using Microsoft.Extensions.Time.Testing; +using Xunit; + +namespace Microsoft.Data.SqlClient.UnitTests.ConnectionPool; + +/// +/// Tests for transaction affinity: routing a returning +/// connection to the transacted store instead of the idle channel, vending an already-enlisted +/// connection back to the same transaction, releasing the connection when the transaction ends, +/// and flowing the ambient transaction on the asynchronous open path. +/// +public class ChannelDbConnectionPoolTransactionTest : IDisposable +{ + private const int DefaultMaxPoolSize = 50; + private const int DefaultMinPoolSize = 0; + private const int DefaultCreationTimeoutInMilliseconds = 15000; + + private IDbConnectionPool _pool; + + public ChannelDbConnectionPoolTransactionTest() + { + _pool = CreatePool(); + } + + public void Dispose() + { + // A transaction entry that outlives its transaction is a leak: the connections it holds + // are never returned to general circulation. + Assert.Empty(_pool.TransactedConnectionPool.TransactedConnections); + + _pool.Shutdown(); + _pool.Clear(); + } + + #region Helper Methods + + /// + /// Builds a pool for these tests. A frozen is injected so that + /// time-driven background maintenance (idle-timeout pruning, warmup and replenishment, + /// blocking-period expiry) cannot advance and race the assertions about pool contents. + /// + private ChannelDbConnectionPool CreatePool( + int maxPoolSize = DefaultMaxPoolSize, + int minPoolSize = DefaultMinPoolSize, + bool hasTransactionAffinity = true) + { + var poolGroupOptions = new DbConnectionPoolGroupOptions( + poolByIdentity: false, + minPoolSize: minPoolSize, + maxPoolSize: maxPoolSize, + creationTimeout: DefaultCreationTimeoutInMilliseconds, + loadBalanceTimeout: 0, + hasTransactionAffinity: hasTransactionAffinity, + idleTimeout: 0 + ); + + var dbConnectionPoolGroup = new DbConnectionPoolGroup( + new SqlConnectionOptions("Data Source=localhost;"), + new ConnectionPoolKey("TestDataSource", credential: null, accessToken: null, accessTokenCallback: null, sspiContextProvider: null), + poolGroupOptions + ); + + var pool = new ChannelDbConnectionPool( + new MockSqlConnectionFactory(), + dbConnectionPoolGroup, + DbConnectionPoolIdentity.NoIdentity, + new DbConnectionPoolProviderInfo(), + timeProvider: new FakeTimeProvider() + ); + + pool.Startup(); + return pool; + } + + /// + /// Tears down the pool built by the constructor and replaces it with one configured + /// differently, for the few tests that need non-default pool options. + /// + private void ReplaceFixturePool(bool hasTransactionAffinity) + { + _pool.Shutdown(); + _pool.Clear(); + _pool = CreatePool(hasTransactionAffinity: hasTransactionAffinity); + } + + /// + /// Opens a connection synchronously. The pool reads the ambient transaction off the calling + /// thread on this path. + /// + private DbConnectionInternal GetConnection(SqlConnection owner) + { + _pool.TryGetConnection( + owner, + taskCompletionSource: null, + TimeoutTimer.StartNew(TimeSpan.FromSeconds(15)), + out DbConnectionInternal? connection); + return connection!; + } + + /// + /// Opens a connection asynchronously, carrying in the + /// 's AsyncState exactly as + /// SqlConnection.InternalOpenAsync does. The pool must take the transaction from there, + /// because the open runs on a thread pool thread the ambient transaction may not flow to. + /// + private async Task GetConnectionAsync( + SqlConnection owner, + Transaction? transaction = null) + { + var tcs = new TaskCompletionSource(transaction); + _pool.TryGetConnection( + owner, + taskCompletionSource: tcs, + TimeoutTimer.StartNew(TimeSpan.FromSeconds(15)), + out DbConnectionInternal? connection); + return connection ?? await tcs.Task; + } + + private void ReturnConnection(DbConnectionInternal connection, SqlConnection owner) => + _pool.ReturnInternalConnection(connection, owner); + + /// + /// Asserts the pool's accounting after a step. + /// + /// Total connections owned by the pool, checked out or not. A connection + /// parked in the transacted store still holds its pool slot and so is counted here. + /// Connections sitting in the idle channel, available to any caller. + /// Connections parked in the transacted store across all + /// transactions. These hold a pool slot but are not available to other callers. + private void AssertPoolState(int count, int idleCount, int transactedCount) + { + Assert.Equal(count, _pool.Count); + Assert.Equal(idleCount, _pool.IdleCount); + Assert.Equal(transactedCount, TotalTransactedConnections()); + } + + private int TotalTransactedConnections() + { + int total = 0; + foreach (var entry in _pool.TransactedConnectionPool.TransactedConnections) + { + total += entry.Value.Count; + } + return total; + } + + private int TransactedConnectionsFor(Transaction transaction) => + _pool.TransactedConnectionPool.TransactedConnections.TryGetValue(transaction, out var connections) + ? connections.Count + : 0; + + #endregion + + #region Connection Return Routing + + /// + /// A connection returned while still enlisted must be parked in the transacted store rather + /// than the idle channel, so it cannot be vended to a caller in a different transaction. It + /// keeps its pool slot while parked. + /// + [Fact] + public void ReturnConnection_WhileEnlisted_ParksInTransactedStoreNotIdleChannel() + { + // Arrange + using var scope = new TransactionScope(); + Transaction? transaction = Transaction.Current; + Assert.NotNull(transaction); + + var owner = new SqlConnection(); + var connection = GetConnection(owner); + Assert.NotNull(connection); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + // Act + ReturnConnection(connection, owner); + + // Assert + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + Assert.Equal(1, TransactedConnectionsFor(transaction!)); + + scope.Complete(); + } + + /// + /// Without an ambient transaction there is nothing to enlist in, so a returning connection + /// goes straight back into the idle channel. + /// + [Fact] + public void ReturnConnection_WithoutTransaction_ReturnsToIdleChannel() + { + // Arrange + var owner = new SqlConnection(); + var connection = GetConnection(owner); + Assert.NotNull(connection); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + // Act + ReturnConnection(connection, owner); + + // Assert + AssertPoolState(count: 1, idleCount: 1, transactedCount: 0); + } + + /// + /// The pool deactivates a returning connection before reading its enlistment, and deactivation + /// is what detaches a completed transaction. A connection returned after its transaction has + /// already ended must therefore land in the idle channel, not be parked under a dead + /// transaction where nothing would ever release it. + /// + [Fact] + public void ReturnConnection_AfterTransactionCompleted_ReturnsToIdleChannel() + { + // Arrange + var owner = new SqlConnection(); + DbConnectionInternal connection; + using (var scope = new TransactionScope()) + { + connection = GetConnection(owner); + Assert.NotNull(connection); + scope.Complete(); + } + + // Act - the transaction is fully disposed by this point. + ReturnConnection(connection, owner); + + // Assert + AssertPoolState(count: 1, idleCount: 1, transactedCount: 0); + } + + /// + /// A pool with automatic enlistment disabled must never bind a connection to the ambient + /// transaction, so the transacted store stays out of the picture entirely. + /// + [Fact] + public void ReturnConnection_WithTransactionAffinityDisabled_ReturnsToIdleChannel() + { + // Arrange + ReplaceFixturePool(hasTransactionAffinity: false); + + using var scope = new TransactionScope(); + var owner = new SqlConnection(); + var connection = GetConnection(owner); + Assert.NotNull(connection); + + // Act + ReturnConnection(connection, owner); + + // Assert + AssertPoolState(count: 1, idleCount: 1, transactedCount: 0); + + scope.Complete(); + } + + /// + /// Returning to a pool that has already shut down destroys the connection instead of pooling + /// it, and must not throw. + /// + [Fact] + public void ReturnConnection_ToShutDownPool_DestroysConnection() + { + // Arrange + using var scope = new TransactionScope(); + var owner = new SqlConnection(); + var connection = GetConnection(owner); + Assert.NotNull(connection); + + _pool.Shutdown(); + + // Act + ReturnConnection(connection, owner); + + // Assert - Dispose() dooms the connection and clears its pool back-reference. + Assert.True(connection.IsConnectionDoomed, + "A connection returned to a shut-down pool should be destroyed, not pooled."); + Assert.Null(connection.Pool); + AssertPoolState(count: 0, idleCount: 0, transactedCount: 0); + } + + #endregion + + #region Vending From The Transacted Store + + /// + /// Round trip: a second request inside the same transaction must be served the connection that + /// is already enlisted in it, rather than a fresh connection. Reusing it is what keeps the + /// transaction from being promoted to a distributed one. + /// + [Fact] + public void GetConnection_UnderSameTransaction_VendsTheAlreadyEnlistedConnection() + { + // Arrange + using var scope = new TransactionScope(); + Transaction? transaction = Transaction.Current; + Assert.NotNull(transaction); + + var owner1 = new SqlConnection(); + var connection1 = GetConnection(owner1); + Assert.NotNull(connection1); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + ReturnConnection(connection1, owner1); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + // Act + var owner2 = new SqlConnection(); + var connection2 = GetConnection(owner2); + + // Assert + Assert.Same(connection1, connection2); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + ReturnConnection(connection2, owner2); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + Assert.Equal(1, TransactedConnectionsFor(transaction!)); + + scope.Complete(); + } + + /// + /// The asynchronous path must reuse the enlisted connection exactly as the synchronous path + /// does, even though it runs the open on a thread pool thread. + /// + [Fact] + public async Task GetConnectionAsync_UnderSameTransaction_VendsTheAlreadyEnlistedConnection() + { + // Arrange + using var scope = new TransactionScope(TransactionScopeAsyncFlowOption.Enabled); + Transaction? transaction = Transaction.Current; + Assert.NotNull(transaction); + + var owner1 = new SqlConnection(); + var connection1 = await GetConnectionAsync(owner1, transaction); + Assert.NotNull(connection1); + ReturnConnection(connection1, owner1); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + // Act + var owner2 = new SqlConnection(); + var connection2 = await GetConnectionAsync(owner2, transaction); + + // Assert + Assert.Same(connection1, connection2); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + ReturnConnection(connection2, owner2); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + scope.Complete(); + } + + /// + /// A connection enlisted in one transaction must never be handed to a caller in a different + /// transaction; each transaction gets its own entry in the transacted store. + /// + [Fact] + public void GetConnection_UnderDifferentTransaction_DoesNotVendTheEnlistedConnection() + { + // Arrange + using var outerScope = new TransactionScope(); + Transaction? outerTransaction = Transaction.Current; + Assert.NotNull(outerTransaction); + + var owner1 = new SqlConnection(); + var connection1 = GetConnection(owner1); + ReturnConnection(connection1, owner1); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + // Act - RequiresNew starts an unrelated transaction. + using (var innerScope = new TransactionScope(TransactionScopeOption.RequiresNew)) + { + Transaction? innerTransaction = Transaction.Current; + Assert.NotEqual(outerTransaction, innerTransaction); + + var owner2 = new SqlConnection(); + var connection2 = GetConnection(owner2); + + // Assert - a fresh connection, because connection1 belongs to the outer transaction. + Assert.NotSame(connection1, connection2); + AssertPoolState(count: 2, idleCount: 0, transactedCount: 1); + + ReturnConnection(connection2, owner2); + AssertPoolState(count: 2, idleCount: 0, transactedCount: 2); + Assert.Equal(1, TransactedConnectionsFor(outerTransaction!)); + Assert.Equal(1, TransactedConnectionsFor(innerTransaction!)); + + innerScope.Complete(); + } + + // Completing the inner transaction releases only its connection; the outer transaction's + // connection stays parked. + AssertPoolState(count: 2, idleCount: 1, transactedCount: 1); + Assert.Equal(1, TransactedConnectionsFor(outerTransaction!)); + + outerScope.Complete(); + } + + #endregion + + #region Transaction Completion + + /// + /// Committing releases the parked connection back into the idle channel, where it becomes + /// available to any caller. + /// + [Fact] + public void TransactionCommit_ReturnsParkedConnectionToIdleChannel() + { + // Arrange + DbConnectionInternal connection; + using (var scope = new TransactionScope()) + { + var owner = new SqlConnection(); + connection = GetConnection(owner); + ReturnConnection(connection, owner); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + // Act + scope.Complete(); + } + + // Assert + AssertPoolState(count: 1, idleCount: 1, transactedCount: 0); + + // The released connection is reusable by a caller with no ambient transaction. + var owner2 = new SqlConnection(); + var connection2 = GetConnection(owner2); + Assert.Same(connection, connection2); + ReturnConnection(connection2, owner2); + } + + /// + /// Rolling back must release the parked connection just as committing does; otherwise an + /// aborted transaction would strand its connection in the transacted store forever. + /// + [Fact] + public void TransactionRollback_ReturnsParkedConnectionToIdleChannel() + { + // Arrange + using (new TransactionScope()) + { + var owner = new SqlConnection(); + var connection = GetConnection(owner); + ReturnConnection(connection, owner); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + // Act - leaving the scope without calling Complete rolls the transaction back. + } + + // Assert + AssertPoolState(count: 1, idleCount: 1, transactedCount: 0); + } + + /// + /// A connection parked in the transacted store survives pool shutdown, because closing it + /// would abort a possibly distributed transaction. Once that transaction ends, the shut-down + /// pool must destroy the connection rather than return it to circulation. + /// + [Fact] + public void TransactionCompletion_AfterShutdown_DestroysConnection() + { + // Arrange + DbConnectionInternal connection; + using (var scope = new TransactionScope()) + { + var owner = new SqlConnection(); + connection = GetConnection(owner); + ReturnConnection(connection, owner); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + // Act - the shutdown drain must leave the transacted connection alone. + _pool.Shutdown(); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + scope.Complete(); + } + + // Assert + AssertPoolState(count: 0, idleCount: 0, transactedCount: 0); + Assert.True(connection.IsConnectionDoomed, + "A shut-down pool must destroy a connection released by its transaction."); + } + + /// + /// TransactionEnded for a connection that was never parked must be a no-op. It must not push a + /// connection that is still checked out into the idle channel, where a second caller could + /// pick it up while the first is still using it. + /// + [Fact] + public void TransactionEnded_ForConnectionThatWasNeverParked_LeavesItCheckedOut() + { + // Arrange + var owner = new SqlConnection(); + var connection = GetConnection(owner); + Assert.NotNull(connection); + + using var scope = new TransactionScope(); + Transaction? transaction = Transaction.Current; + Assert.NotNull(transaction); + + // Act + _pool.TransactionEnded(transaction!, connection); + + // Assert + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + ReturnConnection(connection, owner); + scope.Complete(); + } + + #endregion + + #region Connection Replacement + + /// + /// Replacing a connection mid-transaction must carry the enlistment across, so the replacement + /// is the one that parks in the transacted store when it is returned. + /// + [Fact] + public void ReplaceConnection_CarriesEnlistedTransactionToNewConnection() + { + // Arrange + using var scope = new TransactionScope(); + Transaction? transaction = Transaction.Current; + Assert.NotNull(transaction); + + var owner = new SqlConnection(); + var oldConnection = GetConnection(owner); + Assert.NotNull(oldConnection); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + // Act + var newConnection = _pool.ReplaceConnection( + owner, + oldConnection, + TimeoutTimer.StartNew(TimeSpan.FromSeconds(15))); + + // Assert - a distinct connection took over the old connection's slot. + Assert.NotNull(newConnection); + Assert.NotSame(oldConnection, newConnection); + Assert.True(oldConnection.IsConnectionDoomed); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + ReturnConnection(newConnection, owner); + + // The replacement inherited the transaction, so it parks in the transacted store. + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + Assert.Equal(1, TransactedConnectionsFor(transaction!)); + + scope.Complete(); + } + + #endregion + + #region Async Ambient Transaction Flow + + /// + /// A created without + /// keeps the ambient transaction in thread-static storage, so it is not observable from the + /// thread pool thread the pool opens on. The pool must therefore take the transaction from the + /// 's AsyncState. This test uses a transaction that + /// is never ambient anywhere, so AsyncState is the only way the pool can learn about it. + /// + [Fact] + public async Task GetConnectionAsync_AmbientTransactionNotFlowed_StillEnlistsFromAsyncState() + { + // Arrange + using var transaction = new CommittableTransaction(); + Assert.Null(Transaction.Current); + Assert.Null(await Task.Run(() => Transaction.Current)); + + // Act + var owner = new SqlConnection(); + var connection = await GetConnectionAsync(owner, transaction); + + // Assert + Assert.NotNull(connection); + Assert.Equal(transaction, connection.EnlistedTransaction); + + ReturnConnection(connection, owner); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + Assert.Equal(1, TransactedConnectionsFor(transaction)); + + transaction.Rollback(); + } + + /// + /// Assigning writes to thread-static storage that the + /// ExecutionContext does not unwind, so doing it on a thread pool thread would leave a stale + /// transaction behind for unrelated work later scheduled onto that same thread -- including + /// the login-time auto-enlistment that non-pooled connections perform against the ambient + /// transaction. The pool must pass the transaction explicitly instead of assigning it. + /// + [Fact] + public async Task GetConnectionAsync_DoesNotLeakAmbientTransactionOntoThreadPool() + { + // Arrange & Act - several async opens under a transaction, each on a thread pool thread. + for (int i = 0; i < 8; i++) + { + using var transaction = new CommittableTransaction(); + var owner = new SqlConnection(); + var connection = await GetConnectionAsync(owner, transaction); + ReturnConnection(connection, owner); + transaction.Rollback(); + } + + // Assert - no thread pool thread was left with an ambient transaction. + for (int i = 0; i < 16; i++) + { + Assert.Null(await Task.Run(() => Transaction.Current)); + } + } + + /// + /// The synchronous path runs on the caller's thread, where the ambient transaction set by a + /// TransactionScope is directly observable and must still be honored even though no + /// transaction is handed to the pool explicitly. + /// + [Fact] + public void GetConnection_Sync_UsesAmbientTransactionFromCallersThread() + { + // Arrange + using var scope = new TransactionScope(); + Transaction? transaction = Transaction.Current; + Assert.NotNull(transaction); + + // Act + var owner = new SqlConnection(); + var connection = GetConnection(owner); + + // Assert + Assert.NotNull(connection); + Assert.Equal(transaction, connection.EnlistedTransaction); + + ReturnConnection(connection, owner); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + scope.Complete(); + } + + #endregion + + #region Mock Classes + + internal class MockSqlConnectionFactory : SqlConnectionFactory + { + protected override DbConnectionInternal CreateConnection( + SqlConnectionOptions options, + ConnectionPoolKey poolKey, + DbConnectionPoolGroupProviderInfo poolGroupProviderInfo, + IDbConnectionPool pool, + DbConnection owningConnection, + TimeoutTimer timeout) + { + return new MockDbConnectionInternal(); + } + } + + internal class MockDbConnectionInternal : DbConnectionInternal + { + private static int s_nextId = 1; + public int MockId { get; } = Interlocked.Increment(ref s_nextId); + + public override string ServerVersion => "Mock"; + + public override ConnectionCapabilities Capabilities => new(); + + public override DbTransaction BeginTransaction(System.Data.IsolationLevel il) + { + throw new NotImplementedException(); + } + + public override void EnlistTransaction(Transaction? transaction) + { + if (transaction != null) + { + EnlistedTransaction = transaction; + } + } + + protected override void Activate(Transaction? transaction) + { + EnlistedTransaction = transaction; + } + + protected override void Deactivate() + { + } + + public override string ToString() => $"MockConnection_{MockId}"; + + internal override void ResetConnection() + { + } + } + + #endregion +}