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 3c6347da54..9f75c50c1a 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,17 @@ public ConcurrentDictionary< private int MinPoolSize => PoolGroupOptions.MinPoolSize; + /// + /// Indicates whether the pool automatically enlists connections in the ambient transaction. + /// When disabled, the ambient transaction is neither used to consult the + /// nor handed to activation. + /// + /// This governs the ambient transaction only. A caller may still enlist explicitly via + /// SqlConnection.EnlistTransaction, and such a connection is parked in the transacted store + /// on return regardless of this flag, since it is bound to a transaction either way. + /// + 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 +325,49 @@ public void Clear() /// public void PutObjectFromTransactedPool(DbConnectionInternal connection) { - throw new NotImplementedException(); + 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. + if (State is Running && connection.CanBePooled) + { + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Connection {1}, Transaction has ended; returning connection to pool.", + Id, + connection.ObjectID); + + connection.ResetConnection(); + + // probeLiveness: false because we are running on the System.Transactions + // transaction-completion callback thread (see DbConnectionInternal + // .DelegatedTransactionEnded, whose contract requires the caller to hold a lock on + // the connection). DbConnectionInternal.IsConnectionAlive polls the socket, and we + // do not want to do socket work while holding that lock on a thread we do not own. + // WaitHandleDbConnectionPool does not probe here either; it simply calls + // ResetConnection followed by PutNewObject. The connection is still validated on + // its way back out of the idle channel, so a connection that died during the + // transaction is detected before it is vended. The idle-expiry, load-balance and + // clear-generation checks all still run below; only the socket poll is skipped. + PutConnectionInIdleChannel(connection, probeLiveness: false); + } + else + { + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Connection {1}, Transaction has ended; destroying unpoolable connection.", + Id, + connection.ObjectID); + + // RemoveConnection triggers replenishment, which is the channel pool's equivalent + // of the wait handle pool's QueuePoolCreateRequest. + RemoveConnection(connection); + } } /// @@ -331,9 +385,30 @@ 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); + + // newConnection came from the idle channel, so it already holds a slot of its own. + // Releasing oldConnection's slot here keeps the pool's count accurate. This is + // deliberately different from the create-new branch below, which hands oldConnection's + // slot to the replacement via _connectionSlots.TryReplace and therefore only disposes + // it -- calling RemoveConnection there would signal a free slot that does not exist. oldConnection.DeactivateConnection(); + + // RemoveConnection early-returns without freeing the slot when the connection is a + // transaction root awaiting its transaction's end, which would leave the pool + // holding two reservations for one logical connection. That cannot happen here: + // IsTxRootWaitingForTxEnd is set only by SetInStasis, and the only path that puts a + // pooled connection in stasis is ReturnInternalConnection. oldConnection is checked + // out by owningObject at this point, so it has not been returned and cannot be in + // stasis. Being a transaction root is not sufficient; the connection must also have + // been returned. The slot is therefore always released here. + // + // If that ever changes, the outcome is a transient over-reservation rather than a + // leak: the connection comes back through PutObjectFromTransactedPool when its + // transaction ends, which calls RemoveConnection again once stasis has been + // terminated, and the slot is reclaimed then. RemoveConnection(oldConnection); } else @@ -373,7 +448,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 @@ -422,45 +498,159 @@ public void ReturnInternalConnection(DbConnectionInternal connection, DbConnecti { ValidateOwnershipAndSetPoolingState(connection, owningObject); + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Connection {1}, Deactivating.", + Id, + connection.ObjectID); + + // Deactivate before inspecting the connection, because DeactivateConnection mutates both + // of the gates we branch on below: + // - Deactivate() dooms the connection when async commands are still outstanding, which + // the IsConnectionDoomed check must observe. + // - DeactivateConnection dooms it via DoNotPoolThisConnection when the load-balance + // timeout has elapsed, which the CanBePooled check must observe. + // WaitHandleDbConnectionPool.DeactivateObject orders it the same way. + connection.DeactivateConnection(); + + if (connection.IsConnectionDoomed) + { + // The connection is not fit for reuse -- just dispose of it. + RemoveConnection(connection); + return; + } + + // Note: this logic mirrors WaitHandleDbConnectionPool.ReturnObject, minus one dead + // branch. Its Running path also checks IsTransactionRoot && Pool == null -> SetInStasis, + // under its own "how did we get here if the pool is null?" TODO. A connection cannot + // arrive here without a pool, because this method is called through the connection's own + // Pool reference. The branch is redundant in any case: putting a transaction root in + // stasis is exactly what the first case below does when the connection cannot be pooled. + ReturnDisposition disposition; + lock (connection) + { + if (State is not Running || !connection.CanBePooled) + { + // 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 + // call the appropriate callback when the transaction ends. + if (connection.IsTransactionRoot) + { + connection.SetInStasis(); + disposition = ReturnDisposition.HeldByTransaction; + } + else + { + disposition = ReturnDisposition.Destroy; + } + } + else if (connection.EnlistedTransaction is { } transaction) + { + // 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. + // + // 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. + TransactedConnectionPool.PutTransactedObject(transaction, connection); + disposition = ReturnDisposition.HeldByTransaction; + } + else + { + disposition = ReturnDisposition.Reuse; + } + } + + 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; + } + } + + /// + /// 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. + /// + /// Whether to poll the physical connection to confirm it is still alive. Pass false when + /// running on a thread that must not block, such as a System.Transactions completion + /// callback. The idle-expiry, load-balance timeout and clear-generation checks run + /// regardless; only the socket poll is suppressed. + /// + private void PutConnectionInIdleChannel(DbConnectionInternal connection, bool probeLiveness = true) + { // 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) { connection.SetReturnedTime(_timeProvider.GetUtcNow().UtcDateTime); } - if (!IsLiveConnection(connection)) + if (!IsLiveConnection(connection, probeLiveness)) { RemoveConnection(connection); 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); - } - } } /// @@ -614,7 +804,26 @@ public void Startup() /// public void TransactionEnded(Transaction transaction, DbConnectionInternal transactedObject) { - throw new NotImplementedException(); + // 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); + + // Removal from the transacted list happens synchronously inside this call, and + // TransactedConnectionPool.TransactionEnded calls back into PutObjectFromTransactedPool + // itself to return the connection to general circulation. + // + // We deliberately do not call PutObjectFromTransactedPool ourselves afterwards: that + // callback is conditional on the connection actually having been found in the list. + // A transaction can complete while the application still holds the connection, in which + // case the connection was never parked, and returning it here would hand a connection + // that is still in use to another caller. In that case the connection instead reaches + // the pool through the normal ReturnInternalConnection path when it is closed. + // This mirrors WaitHandleDbConnectionPool.TransactionEnded. + TransactedConnectionPool.TransactionEnded(transaction, transactedObject); } /// @@ -640,10 +849,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 @@ -672,6 +883,32 @@ 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 by the caller, on the caller's thread, and handed + // to us in the TaskCompletionSource's AsyncState (see SqlConnection.InternalOpenAsync). + // + // We must not read Transaction.Current inside the Task.Run below. A + // TransactionScope created with TransactionScopeAsyncFlowOption.Enabled stores the + // transaction in an AsyncLocal, which does flow onto the pool's worker thread, so that + // would appear to work. But Enabled is not the default: a plain TransactionScope keeps + // the transaction in thread-static storage, which does not flow, and reading + // Transaction.Current on the worker would silently fail to enlist. The WaitHandle pool + // enlists correctly in that case, so this is also a compatibility requirement. + // AsyncState is correct under both options. + // + // This does not make the suppressed-flow pattern work -- the caller's own scope is + // still broken past the first await -- but it keeps the connection in the transaction + // the caller intended rather than silently running outside it. + // + // Note that we deliberately do not assign Transaction.Current on the thread pool + // thread either. 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) @@ -679,10 +916,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 @@ -690,7 +923,8 @@ public bool TryGetConnection( connection = await GetInternalConnection( owningObject, async: true, - timeout + timeout, + ambientTransaction ).ConfigureAwait(false); if (!taskCompletionSource.TrySetResult(connection)) @@ -901,8 +1135,12 @@ _connectionCreationRateLimiter is not null && /// Checks that the provided connection is live and unexpired and closes it if needed. /// /// + /// + /// Whether to poll the physical connection to confirm it is still alive. Pass false when + /// running on a thread that must not block; the remaining checks are all cheap and local. + /// /// Returns true if the connection is live and unexpired, otherwise returns false. - private bool IsLiveConnection(DbConnectionInternal connection) + private bool IsLiveConnection(DbConnectionInternal connection, bool probeLiveness = true) { // Connection has been sitting idle longer than the configured idle timeout. // Checked before the (potentially expensive) liveness probe so an idle-expired @@ -922,8 +1160,9 @@ private bool IsLiveConnection(DbConnectionInternal connection) return false; } - // Broken physical connection - if (!connection.IsConnectionAlive()) + // Broken physical connection. Skipped when the caller cannot afford to block: this + // polls the socket, so it must not run on a thread we do not own. + if (probeLiveness && !connection.IsConnectionAlive()) { return false; } @@ -944,11 +1183,40 @@ private bool IsLiveConnection(DbConnectionInternal connection) } /// - /// Closes the provided connection and removes it from the pool. + /// Closes the provided connection and removes it from the pool, freeing its slot. /// + /// + /// + /// Not guaranteed: a connection that is a transaction root awaiting its transaction's end + /// is left alone and its slot stays reserved, because disposing it would abort a possibly + /// distributed transaction. Callers that treat this method as "the slot is now free" must + /// tolerate that, including and the drain, which + /// can therefore complete while the pool still owns a connection. + /// + /// + /// In practice no current caller reaches that early return, because a connection in stasis + /// is filed in neither the idle channel nor the transacted store, and every caller sources + /// its connection from one of those or from a connection that is still checked out. The + /// guard is retained to match WaitHandleDbConnectionPool.DestroyObject and to keep the + /// invariant enforced if a future caller does reach it. + /// + /// /// 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. @@ -1003,6 +1271,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 @@ -1015,20 +1287,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(); @@ -1070,9 +1368,8 @@ private async Task GetInternalConnection( connection = null; } } - while (connection is null); - PrepareConnection(owningConnection, connection); + PrepareConnection(owningConnection, connection, transaction); return connection; } @@ -1141,6 +1438,55 @@ 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}, Transaction {1}, Connection {2}, Popped from transacted pool.", + Id, + transaction.GetHashCode(), + 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. + bool isAlive = false; + try + { + // A dead transaction root must surface the underlying failure to the caller, since + // there is no way to recover the delegated transaction on another connection. Any + // other dead connection is simply reported so the caller can pick up or open a + // different one. Either way the connection is dropped in the finally below. + isAlive = connection.IsConnectionAlive(throwOnException: connection.IsTransactionRoot); + } + finally + { + if (!isAlive) + { + 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 b3c9815eb8..228e949cd0 100644 --- a/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/ChannelDbConnectionPoolTest.cs +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/ChannelDbConnectionPoolTest.cs @@ -888,37 +888,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..84f581ea2a --- /dev/null +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/ConnectionPool/ChannelDbConnectionPoolTransactionTest.cs @@ -0,0 +1,1269 @@ +// 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; + private MockSqlConnectionFactory _connectionFactory = null!; + + 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 + ); + + _connectionFactory = new MockSqlConnectionFactory(); + + var pool = new ChannelDbConnectionPool( + _connectionFactory, + 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); + } + + /// + /// Replaces the fixture's pool with one capped at , so a test can + /// exhaust it and observe what happens to a caller that has to wait. + /// + private void ReplaceFixturePool(int maxPoolSize) + { + _pool.Shutdown(); + _pool.Clear(); + _pool = CreatePool(maxPoolSize: maxPoolSize); + } + + /// + /// 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 the way SqlConnection.InternalOpenAsync does: the ambient + /// transaction is captured here, on the caller's thread, and handed to the pool in the + /// 's AsyncState. The pool must take it from there, + /// because the open itself runs on a thread pool thread the ambient transaction may not flow + /// to, and on a retry the pool is re-entered from a continuation on an arbitrary thread. + /// + /// The owning connection. + /// The transaction to hand to the pool. Defaults to the caller's + /// ambient transaction, matching what InternalOpenAsync captures. + private async Task GetConnectionAsync( + SqlConnection owner, + Transaction? transaction = null) + { + var tcs = new TaskCompletionSource(transaction ?? Transaction.Current); + _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); + + /// + /// Awaits a task with a timeout so a test that expects a waiter to be served fails with a clear + /// message instead of hanging. Written against because + /// Task.WaitAsync does not exist on .NET Framework. + /// + private static async Task WithTimeout(Task task, TimeSpan timeout) + { + Task completed = await Task.WhenAny(task, Task.Delay(timeout)).ConfigureAwait(false); + Assert.True( + ReferenceEquals(completed, task), + $"Timed out after {timeout.TotalSeconds:0.#}s waiting for the task to complete."); + return await task.ConfigureAwait(false); + } + + /// + /// 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; + + /// + /// Asserts which transaction a connection is bound to, complementing the pool-state assertions: + /// AssertPoolState says where the connection was filed, this says what it is actually enlisted + /// in. Pass null to assert the connection is not enlisted at all. + /// + /// + /// Compared by equality rather than reference because the EnlistedTransaction setter stores a + /// clone, so that the connection does not hold the caller's transaction past the end of its + /// using block. + /// + private static void AssertEnlistedIn(Transaction? expected, DbConnectionInternal connection) => + Assert.Equal(expected, connection.EnlistedTransaction); + + #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); + AssertEnlistedIn(transaction, 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!)); + + // The connection stays bound to the transaction while parked, which is what makes it + // ineligible for a caller in any other transaction. + AssertEnlistedIn(transaction, connection); + + 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); + AssertEnlistedIn(null, connection); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + // Act + ReturnConnection(connection, owner); + + // Assert + AssertEnlistedIn(null, connection); + AssertPoolState(count: 1, idleCount: 1, transactedCount: 0); + } + + /// + /// A connection whose transaction has already ended must land in the idle channel, not be + /// parked under a dead transaction where nothing would ever release it. The enlistment is + /// dropped when the transaction completes, so by the time the connection is returned there is + /// nothing left to park it under. + /// + [Fact] + public void ReturnConnection_AfterTransactionCompleted_ReturnsToIdleChannel() + { + // Arrange + var owner = new SqlConnection(); + DbConnectionInternal connection; + using (var scope = new TransactionScope()) + { + connection = GetConnection(owner); + Assert.NotNull(connection); + AssertEnlistedIn(Transaction.Current, connection); + scope.Complete(); + } + + // The completed transaction detached itself from the connection. + AssertEnlistedIn(null, connection); + + // 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); + + // The ambient transaction exists but must not reach the connection. + Assert.NotNull(Transaction.Current); + AssertEnlistedIn(null, connection); + + // Act + ReturnConnection(connection, owner); + + // Assert + AssertEnlistedIn(null, connection); + 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); + } + + /// + /// A doomed connection is destroyed even when it is still enlisted. The doomed check runs + /// before the transacted decision precisely so that a connection unfit for reuse is never + /// parked in the transacted store, where it would be handed to the next caller in that same + /// transaction. + /// + [Fact] + public void ReturnConnection_WhenDoomedWhileEnlisted_DestroysWithoutParking() + { + // Arrange + using var scope = new TransactionScope(); + Transaction? transaction = Transaction.Current; + Assert.NotNull(transaction); + + var owner = new SqlConnection(); + var connection = GetConnection(owner); + Assert.NotNull(connection); + AssertEnlistedIn(transaction, connection); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + // The connection breaks while it is enlisted and checked out. + ((MockDbConnectionInternal)connection).Doom(); + Assert.True(connection.IsConnectionDoomed); + + // Act + ReturnConnection(connection, owner); + + // Assert - destroyed rather than parked under its transaction, and its slot is released. + Assert.Equal(0, TransactedConnectionsFor(transaction!)); + AssertPoolState(count: 0, idleCount: 0, transactedCount: 0); + Assert.Null(connection.Pool); + + scope.Complete(); + } + + #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); + AssertEnlistedIn(transaction, 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); + AssertEnlistedIn(transaction, 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); + AssertEnlistedIn(transaction, 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); + AssertEnlistedIn(outerTransaction, connection1); + ReturnConnection(connection1, owner1); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + // Act - RequiresNew starts an unrelated transaction. + Transaction? innerTransaction; + using (var innerScope = new TransactionScope(TransactionScopeOption.RequiresNew)) + { + 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); + AssertEnlistedIn(innerTransaction, connection2); + AssertEnlistedIn(outerTransaction, connection1); + AssertPoolState(count: 2, idleCount: 0, transactedCount: 1); + + // The store holds only the outer transaction's connection, so connection2 was newly + // created rather than taken from the store. + Assert.Equal(1, TransactedConnectionsFor(outerTransaction!)); + Assert.Equal(0, TransactedConnectionsFor(innerTransaction!)); + + 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!)); + Assert.Equal(0, TransactedConnectionsFor(innerTransaction!)); + + 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 transaction released it, so it carries no enlistment into general circulation. + AssertEnlistedIn(null, connection); + + // The released connection is reusable by a caller with no ambient transaction. + var owner2 = new SqlConnection(); + var connection2 = GetConnection(owner2); + Assert.Same(connection, connection2); + AssertEnlistedIn(null, 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 + DbConnectionInternal connection; + using (new TransactionScope()) + { + var owner = new SqlConnection(); + 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); + AssertEnlistedIn(null, connection); + } + + /// + /// 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 on its own, with none of the surrounding lifecycle events, must be a no-op + /// for a connection that was never parked. Only parking hands the connection to the pool, so + /// even an enlisted connection is still checked out here, and pushing it into the idle channel + /// would let a second caller pick it up while the first is still using it. + /// + [Fact] + public void TransactionEnded_ForEnlistedConnectionThatWasNeverParked_LeavesItCheckedOut() + { + // Arrange + using var scope = new TransactionScope(); + Transaction? transaction = Transaction.Current; + Assert.NotNull(transaction); + + var owner = new SqlConnection(); + var connection = GetConnection(owner); + Assert.NotNull(connection); + + // The connection is enlisted, but checked out rather than parked. + AssertEnlistedIn(transaction, connection); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + // Act - the completion notification arrives on its own, with no detach and no return. + _pool.TransactionEnded(transaction!, connection); + + // Assert - nothing to release, and the enlistment is untouched. + AssertEnlistedIn(transaction, connection); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + // It still reaches the pool through the normal return path. + ReturnConnection(connection, owner); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + scope.Complete(); + } + + /// + /// A transaction can complete while the application still holds its connection. The connection + /// was never parked, so there is nothing for the completion to release, and pushing it into the + /// idle channel would hand a connection that is still in use to a second caller. + /// + [Fact] + public void TransactionCompletes_WhileConnectionStillCheckedOut_LeavesItCheckedOut() + { + // Arrange + var owner = new SqlConnection(); + DbConnectionInternal connection; + using (var scope = new TransactionScope()) + { + connection = GetConnection(owner); + Assert.NotNull(connection); + AssertEnlistedIn(Transaction.Current, connection); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + // Act - the transaction completes with the connection still checked out. + scope.Complete(); + } + + // Assert + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + // It reaches the pool only when its owner returns it. + ReturnConnection(connection, owner); + AssertPoolState(count: 1, idleCount: 1, transactedCount: 0); + } + + /// + /// Releasing a parked connection runs on the System.Transactions transaction-completion + /// callback thread, under DbConnectionInternal.DelegatedTransactionEnded's requirement that the + /// caller holds a lock on the connection. The pool must not poll the socket there, so the + /// liveness probe is suppressed on that path and the connection is pooled without it. + /// WaitHandleDbConnectionPool likewise does not probe when returning a connection from the + /// transacted store. A connection that died during the transaction is still caught, just later: + /// it is validated on its way back out of the idle channel, before it can be vended. + /// + [Fact] + public void TransactionCommit_ReturningParkedConnection_DoesNotProbeLiveness() + { + // Arrange + MockDbConnectionInternal connection; + using (var scope = new TransactionScope()) + { + var owner = new SqlConnection(); + connection = Assert.IsType(GetConnection(owner)); + ReturnConnection(connection, owner); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + // Report the connection as dead and ignore any probes made while arranging, so the + // assertions below describe only what the completion path itself did. + connection.IsDead = true; + connection.LivenessProbeCount = 0; + + // Act + scope.Complete(); + } + + // Assert - the completion path pooled it without asking whether it was alive. + Assert.Equal(0, connection.LivenessProbeCount); + AssertPoolState(count: 1, idleCount: 1, transactedCount: 0); + + // The deferred check still catches it: vending probes the connection, finds it dead, and + // discards it rather than handing a broken connection to the caller. + var owner2 = new SqlConnection(); + var connection2 = GetConnection(owner2); + Assert.NotSame(connection, connection2); + Assert.True(connection.LivenessProbeCount > 0); + ReturnConnection(connection2, owner2); + } + + /// + /// The counterpart to : + /// suppressing the probe is specific to the transaction-completion path. An ordinary return + /// runs on the caller's own thread, where blocking is acceptable, so it still probes and drops + /// a dead connection immediately instead of leaving it to be discovered later. + /// + [Fact] + public void ReturnConnection_WhenDeadAndNotEnlisted_ProbesLivenessAndDestroysIt() + { + // Arrange + var owner = new SqlConnection(); + var connection = Assert.IsType(GetConnection(owner)); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + connection.IsDead = true; + connection.LivenessProbeCount = 0; + + // Act + ReturnConnection(connection, owner); + + // Assert + Assert.True(connection.LivenessProbeCount > 0); + AssertPoolState(count: 0, idleCount: 0, transactedCount: 0); + } + + /// + /// A connection parked in the transacted store still holds its pool slot, so with MaxPoolSize + /// reached and everything parked, a caller outside the transaction cannot be served and must + /// wait. Ending the transaction has to wake it: the release path writes the connection into the + /// idle channel, which completes the waiter. If it did not, the caller would block until its + /// own timeout even though a connection had become available. + /// + [Fact] + public async Task TransactionEnd_WithPoolExhaustedByParkedConnection_WakesWaitingCaller() + { + // Arrange - a pool of exactly one connection, parked under a transaction. + ReplaceFixturePool(maxPoolSize: 1); + + DbConnectionInternal parkedConnection; + var waiter = new TaskCompletionSource(state: null); + + using (var scope = new TransactionScope()) + { + var owner = new SqlConnection(); + parkedConnection = GetConnection(owner); + ReturnConnection(parkedConnection, owner); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + // A second caller with no ambient transaction. The pool is at MaxPoolSize and the idle + // channel is empty, and the parked connection belongs to a transaction this caller is + // not in, so it cannot be served now. Suppress the ambient transaction so the request + // is genuinely unrelated rather than being handed the parked connection by affinity. + using (new TransactionScope(TransactionScopeOption.Suppress)) + { + bool completedSynchronously = _pool.TryGetConnection( + new SqlConnection(), + waiter, + TimeoutTimer.StartNew(TimeSpan.FromSeconds(15)), + out DbConnectionInternal? immediateConnection); + + Assert.False(completedSynchronously); + Assert.Null(immediateConnection); + } + + Assert.False(waiter.Task.IsCompleted, "The caller must wait while the only connection is parked."); + + // Act + scope.Complete(); + } + + // Assert - the waiter was handed the released connection rather than timing out. + DbConnectionInternal servedConnection = await WithTimeout(waiter.Task, TimeSpan.FromSeconds(15)); + Assert.Same(parkedConnection, servedConnection); + AssertEnlistedIn(null, servedConnection); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + } + + /// + /// Clear cannot destroy a connection parked in the transacted store, because doing so would + /// abort a possibly distributed transaction, and the parked connection is not in the idle + /// channel that Clear drains. It is instead retired lazily: Clear bumps the generation counter, + /// and when the transaction ends the release path sees the stale generation and destroys the + /// connection rather than returning it to circulation. This is the channel pool's substitute + /// for WaitHandleDbConnectionPool's eager DoNotPoolThisConnection sweep, so the lag between + /// Clear returning and the connection actually going away is expected. + /// + [Fact] + public void Clear_WhileConnectionIsParked_RetiresItOnTransactionEndInsteadOfRepooling() + { + // 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 - Clear runs while the connection is parked and its transaction is live. + _pool.Clear(); + + // Assert - it survives Clear, still parked and still holding its slot. + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + Assert.False(connection.IsConnectionDoomed); + + scope.Complete(); + } + + // Assert - on release the stale generation is detected, so the connection is destroyed + // instead of being handed to the next caller. + Assert.True(connection.IsConnectionDoomed); + Assert.Null(connection.Pool); + AssertPoolState(count: 0, idleCount: 0, transactedCount: 0); + + // The next caller therefore gets a freshly created connection, not the cleared one. + var owner2 = new SqlConnection(); + var connection2 = GetConnection(owner2); + Assert.NotSame(connection, connection2); + ReturnConnection(connection2, owner2); + } + + #endregion + + #region Stasis + /// + /// A transaction root that cannot be pooled is put in stasis rather than destroyed: closing it + /// would orphan the root transaction with no way to commit or roll back. Stasis is not the + /// transacted store. The connection is filed nowhere, so it is invisible to the idle channel, + /// to and to pruning, but it still holds its pool + /// slot. The slot comes back when the transaction ends and the connection unwinds through + /// PutObjectFromTransactedPool. + /// + [Fact] + public void ReturnConnection_NonPoolableTransactionRoot_EntersStasisThenReleasesSlotOnTransactionEnd() + { + // Arrange + MockDbConnectionInternal connection; + using (var scope = new TransactionScope()) + { + var owner = new SqlConnection(); + connection = Assert.IsType(GetConnection(owner)); + + // A root that cannot be pooled. Deliberately not doomed: IsConnectionDoomed would + // short-circuit the return path before the stasis decision is ever reached. + connection.IsRoot = true; + connection.MarkDoNotPool(); + Assert.False(connection.CanBePooled); + Assert.False(connection.IsConnectionDoomed); + + // Act + ReturnConnection(connection, owner); + + // Assert - in stasis: holding its slot, but filed in neither the idle channel nor the + // transacted store, so no other caller can reach it. + Assert.True(connection.IsTxRootWaitingForTxEnd); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + // It survives a Clear that would have destroyed any pooled connection, because Clear + // drains the idle channel and a connection in stasis is not in it. + _pool.Clear(); + Assert.True(connection.IsTxRootWaitingForTxEnd); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + scope.Complete(); + } + + // Assert - the transaction ended, stasis terminated, and because the connection still + // cannot be pooled it was destroyed rather than returned to circulation. Its slot is free. + Assert.False(connection.IsTxRootWaitingForTxEnd); + Assert.Null(connection.Pool); + AssertPoolState(count: 0, idleCount: 0, transactedCount: 0); + } + + /// + /// The transaction-root counterpart to + /// . A shut-down pool destroys + /// an ordinary connection on return, but a transaction root must still go to stasis instead, or + /// shutting down a pool would abort a live transaction. The slot stays outstanding until the + /// transaction ends, which is why a completed Shutdown does not imply the pool owns nothing. + /// + [Fact] + public void ReturnConnection_TransactionRootToShutDownPool_EntersStasisRatherThanBeingDestroyed() + { + // Arrange + MockDbConnectionInternal connection; + using (var scope = new TransactionScope()) + { + var owner = new SqlConnection(); + connection = Assert.IsType(GetConnection(owner)); + connection.IsRoot = true; + + _pool.Shutdown(); + + // Act + ReturnConnection(connection, owner); + + // Assert - not destroyed, unlike the non-root case. + Assert.True(connection.IsTxRootWaitingForTxEnd); + Assert.False(connection.IsConnectionDoomed); + Assert.NotNull(connection.Pool); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 0); + + scope.Complete(); + } + + // Assert - the pool is not Running, so the connection is destroyed on the way out of + // stasis rather than pooled, and only now does the pool's count reach zero. + Assert.False(connection.IsTxRootWaitingForTxEnd); + Assert.True(connection.IsConnectionDoomed); + Assert.Null(connection.Pool); + AssertPoolState(count: 0, idleCount: 0, transactedCount: 0); + } + + #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); + AssertEnlistedIn(transaction, 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, carrying the + // enlistment with it. + Assert.NotNull(newConnection); + Assert.NotSame(oldConnection, newConnection); + AssertEnlistedIn(transaction, 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 + /// -- the default -- keeps its ambient + /// transaction in thread-static storage, so it does not flow onto the thread pool thread the + /// pool opens on. The connection must still enlist, because the transaction is captured on the + /// caller's thread before the open is scheduled (see SqlConnection.InternalOpenAsync). + /// + /// This is one of three guards on that capture. With Enabled the transaction rides an + /// AsyncLocal and is live on the pool's worker thread, so reading Transaction.Current there + /// would pass every other test in this class -- and silently stop enlisting for the default + /// option, which the WaitHandle pool handles correctly. This case is the realistic shape of + /// that regression; see also + /// + /// and . + /// + /// The open is started inside the scope but awaited outside it. That is not incidental: + /// awaiting inside resumes the continuation on a thread pool thread, and disposing the scope + /// there throws "A TransactionScope must be disposed on the same thread that it was created." + /// That is the very limitation TransactionScopeAsyncFlowOption.Enabled exists to remove, so a + /// test of the suppressed case cannot avoid it. The scope is given an explicit + /// so that leaving the scope ends the scope without + /// ending the transaction the pool is still enlisting in. + /// + [Fact] + public async Task GetConnectionAsync_WithAsyncFlowDisabled_StillEnlistsInAmbientTransaction() + { + // Arrange + using var transaction = new CommittableTransaction(); + var owner = new SqlConnection(); + Task openTask; + + using (var scope = new TransactionScope(transaction)) + { + Assert.Equal(transaction, Transaction.Current); + + // Act - starting the open captures the ambient transaction synchronously, here, and + // hands the rest of the work to a thread pool thread. + openTask = GetConnectionAsync(owner); + scope.Complete(); + } + + var connection = await openTask; + + // 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(); + } + + /// + /// The pool must take the transaction from AsyncState rather than from whatever is ambient on + /// the thread that happens to enter it. TryGetConnection is normally called synchronously from + /// the caller's thread, where the two agree -- but on the retry path + /// (SqlConnection.OpenAsyncRetry.Retry) the pool is re-entered from a continuation running on + /// an arbitrary thread, which has no ambient transaction. Reading Transaction.Current there + /// would silently drop the enlistment, or worse, pick up an unrelated transaction. + /// + /// No TransactionScope is opened here, so this thread already stands in for that continuation: + /// the transaction exists but is not ambient anywhere, and AsyncState is the only channel + /// carrying it. + /// + [Fact] + public async Task GetConnectionAsync_EnteredFromThreadWithoutAmbientTransaction_EnlistsFromAsyncState() + { + // Arrange + using var transaction = new CommittableTransaction(); + var owner = new SqlConnection(); + Assert.Null(Transaction.Current); + + // Act - carry the transaction only in AsyncState, the way a retry continuation does. + 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(); + } + + /// + /// The pool must enlist the connection in the caller's transaction without making that + /// transaction ambient on the thread it opens on. Assigning Transaction.Current writes to + /// thread-static storage that 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. + /// + /// The connection factory runs inside the pool's own open work, so it observes what the pool + /// made ambient there. That is an observation of the code under test rather than of whichever + /// thread happened to run it, so it does not depend on thread identity or scheduling. + /// + [Fact] + public async Task GetConnectionAsync_DoesNotSetAmbientTransactionOnPoolWorkerThread() + { + // Arrange - a transaction the pool must enlist in but must not make ambient. + using var transaction = new CommittableTransaction(); + Assert.Null(Transaction.Current); + + // Act + var owner = new SqlConnection(); + var connection = await GetConnectionAsync(owner, transaction); + + // Assert - the pool left the worker's ambient transaction alone while still enlisting the + // connection. + Assert.Equal(1, _connectionFactory.CreateCount); + Assert.Null(_connectionFactory.AmbientTransactionAtCreate); + Assert.Equal(transaction, connection.EnlistedTransaction); + + ReturnConnection(connection, owner); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + + transaction.Rollback(); + } + + /// + /// 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(); + } + + /// + /// The asynchronous equivalent of the case above. The pool cannot read the caller's ambient + /// transaction itself, because the open runs on a thread pool thread it does not flow to, so + /// the caller captures it and passes it in AsyncState. Even with a scope created with + /// , where the transaction is genuinely + /// ambient on the calling thread, that capture is what has to carry it through. + /// + [Fact] + public async Task GetConnectionAsync_UsesAmbientTransactionCapturedOnCallersThread() + { + // Arrange + using var scope = new TransactionScope(TransactionScopeAsyncFlowOption.Enabled); + Transaction? transaction = Transaction.Current; + Assert.NotNull(transaction); + + // Act - no transaction is passed explicitly, so the helper captures the ambient one on + // this thread exactly as SqlConnection.InternalOpenAsync does. + var owner = new SqlConnection(); + var connection = await GetConnectionAsync(owner); + + // Assert + Assert.NotNull(connection); + Assert.Equal(transaction, connection.EnlistedTransaction); + + ReturnConnection(connection, owner); + AssertPoolState(count: 1, idleCount: 0, transactedCount: 1); + Assert.Equal(1, TransactedConnectionsFor(transaction!)); + + scope.Complete(); + } + + #endregion + + #region Mock Classes + + internal class MockSqlConnectionFactory : SqlConnectionFactory + { + /// + /// The value of observed on the thread the pool used to + /// create the connection. On the asynchronous path that is the thread pool thread the pool + /// runs its open work on, so this is a direct observation of whether the pool assigned the + /// ambient transaction there. + /// + public Transaction? AmbientTransactionAtCreate { get; private set; } + + public int CreateCount { get; private set; } + + protected override DbConnectionInternal CreateConnection( + SqlConnectionOptions options, + ConnectionPoolKey poolKey, + DbConnectionPoolGroupProviderInfo poolGroupProviderInfo, + IDbConnectionPool pool, + DbConnection owningConnection, + TimeoutTimer timeout) + { + AmbientTransactionAtCreate = Transaction.Current; + CreateCount++; + 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() + { + } + + internal void Doom() => DoomThisConnection(); + + /// + /// Marks the connection unfit for pooling without dooming it. This is the distinction the + /// pool cares about: short-circuits + /// the return path before the transacted decision, whereas a merely non-poolable connection + /// still reaches it and can be put in stasis. + /// + internal void MarkDoNotPool() => DoNotPoolThisConnection(); + + /// + /// When true, the connection reports itself as the root of a delegated transaction, which + /// is what makes the pool put it in stasis rather than destroy it. + /// + internal bool IsRoot { get; set; } + + internal override bool IsTransactionRoot => IsRoot; + + /// + /// When true, the connection reports itself as dead from + /// , simulating a physical connection that died while the + /// pool was not looking. + /// + internal bool IsDead { get; set; } + + /// + /// Counts calls to . Settable so a test can zero it after + /// arranging and observe only the probes made by the operation under test. + /// + internal int LivenessProbeCount { get; set; } + + internal override bool IsConnectionAlive(bool throwOnException = false) + { + LivenessProbeCount++; + return !IsDead; + } + + public override string ToString() => $"MockConnection_{MockId}"; + + internal override void ResetConnection() + { + } + } + + #endregion +}