diff --git a/src/Directory.Packages.props b/src/Directory.Packages.props index f897e55f77..63cb90936a 100644 --- a/src/Directory.Packages.props +++ b/src/Directory.Packages.props @@ -11,6 +11,7 @@ + diff --git a/src/Microsoft.Data.SqlClient/netcore/ref/Microsoft.Data.SqlClient.csproj b/src/Microsoft.Data.SqlClient/netcore/ref/Microsoft.Data.SqlClient.csproj index 60eff24c1c..c02c44829c 100644 --- a/src/Microsoft.Data.SqlClient/netcore/ref/Microsoft.Data.SqlClient.csproj +++ b/src/Microsoft.Data.SqlClient/netcore/ref/Microsoft.Data.SqlClient.csproj @@ -45,6 +45,11 @@ + + + + + diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj index 36de19a385..d1d6da1f4e 100644 --- a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj +++ b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj @@ -102,6 +102,9 @@ Microsoft\Data\SqlClient\ConnectionPool\ChannelDbConnectionPool.cs + + Microsoft\Data\SqlClient\ConnectionPool\ConnectionPoolSlots.cs + Microsoft\Data\SqlClient\ConnectionPool\DbConnectionPoolAuthenticationContext.cs diff --git a/src/Microsoft.Data.SqlClient/netfx/ref/Microsoft.Data.SqlClient.csproj b/src/Microsoft.Data.SqlClient/netfx/ref/Microsoft.Data.SqlClient.csproj index 6b507b5a0a..bb2fe8ee4d 100644 --- a/src/Microsoft.Data.SqlClient/netfx/ref/Microsoft.Data.SqlClient.csproj +++ b/src/Microsoft.Data.SqlClient/netfx/ref/Microsoft.Data.SqlClient.csproj @@ -46,6 +46,7 @@ + diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj index 0c2fe65f93..b578e2f78f 100644 --- a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj +++ b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj @@ -291,6 +291,9 @@ Microsoft\Data\SqlClient\ConnectionPool\ChannelDbConnectionPool.cs + + Microsoft\Data\SqlClient\ConnectionPool\ConnectionPoolSlots.cs + Microsoft\Data\SqlClient\ConnectionPool\DbConnectionPoolAuthenticationContext.cs @@ -966,6 +969,7 @@ + diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft.Data.SqlClient.csproj b/src/Microsoft.Data.SqlClient/src/Microsoft.Data.SqlClient.csproj index aef9ceeb7e..15e3947b7f 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft.Data.SqlClient.csproj +++ b/src/Microsoft.Data.SqlClient/src/Microsoft.Data.SqlClient.csproj @@ -21,7 +21,8 @@ - + + diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/ProviderBase/DbConnectionInternal.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/ProviderBase/DbConnectionInternal.cs index 31f207e5ce..a01e0c16ca 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/ProviderBase/DbConnectionInternal.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/ProviderBase/DbConnectionInternal.cs @@ -46,11 +46,6 @@ internal abstract class DbConnectionInternal /// private bool _cannotBePooled; - /// - /// When the connection was created. - /// - private DateTime _createTime; - /// /// [usage must be thread-safe] the transaction that we're enlisted in, either manually or automatically. /// @@ -93,10 +88,16 @@ internal DbConnectionInternal(ConnectionState state, bool hidePassword, bool all AllowSetConnectionString = allowSetConnectionString; ShouldHidePassword = hidePassword; State = state; + CreateTime = DateTime.UtcNow; } #region Properties + /// + /// When the connection was created. + /// + internal DateTime CreateTime { get; } + internal bool AllowSetConnectionString { get; } internal bool CanBePooled => !IsConnectionDoomed && !_cannotBePooled && !_owningObject.TryGetTarget(out _); @@ -531,7 +532,7 @@ internal void DeactivateConnection() // If we're not already doomed, check the connection's lifetime and // doom it if it's lifetime has elapsed. DateTime now = DateTime.UtcNow; - if (now.Ticks - _createTime.Ticks > Pool.LoadBalanceTimeout.Ticks) + if (now.Ticks - CreateTime.Ticks > Pool.LoadBalanceTimeout.Ticks) { DoNotPoolThisConnection(); } @@ -701,7 +702,6 @@ internal void MakeNonPooledObject(DbConnection owningObject) /// internal void MakePooledConnection(IDbConnectionPool connectionPool) { - _createTime = DateTime.UtcNow; Pool = connectionPool; } @@ -756,7 +756,7 @@ internal virtual void PrepareForReplaceConnection() // By default, there is no preparation required } - internal void PrePush(object expectedOwner) + internal void PrePush(DbConnection expectedOwner) { // Called by IDbConnectionPool when we're about to be put into it's pool, we take this // opportunity to ensure ownership and pool counts are legit. 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 8b0efe8763..4270a2deb1 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 @@ -3,12 +3,18 @@ // See the LICENSE file in the project root for more information. using System; 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; using System.Threading.Tasks; using System.Transactions; using Microsoft.Data.Common; using Microsoft.Data.Common.ConnectionString; using Microsoft.Data.ProviderBase; +using static Microsoft.Data.SqlClient.ConnectionPool.DbConnectionPoolState; #nullable enable @@ -20,52 +26,107 @@ namespace Microsoft.Data.SqlClient.ConnectionPool /// internal sealed class ChannelDbConnectionPool : IDbConnectionPool { + #region Fields + // Limits synchronous operations which depend on async operations on managed + // threads from blocking on all available threads, which would stop async tasks + // from being scheduled and cause deadlocks. Use ProcessorCount/2 as a balance + // between sync and async tasks. + private static SemaphoreSlim _syncOverAsyncSemaphore = new(Math.Max(1, Environment.ProcessorCount / 2)); + + /// + /// Tracks the number of instances of this class. Used to generate unique IDs for each instance. + /// + private static int _instanceCount; + + private readonly int _instanceId = Interlocked.Increment(ref _instanceCount); + + /// + /// Tracks all connections currently managed by this pool, whether idle or busy. + /// Only updated rarely - when physical connections are opened/closed - but is read in perf-sensitive contexts. + /// + private readonly ConnectionPoolSlots _connectionSlots; + + /// + /// Reader side for the idle connection channel. Contains nulls in order to release waiting attempts after + /// a connection has been physically closed/broken. + /// + private readonly ChannelReader _idleConnectionReader; + private readonly ChannelWriter _idleConnectionWriter; + #endregion + + /// + /// Initializes a new PoolingDataSource. + /// + internal ChannelDbConnectionPool( + SqlConnectionFactory connectionFactory, + DbConnectionPoolGroup connectionPoolGroup, + DbConnectionPoolIdentity identity, + DbConnectionPoolProviderInfo connectionPoolProviderInfo) + { + ConnectionFactory = connectionFactory; + PoolGroup = connectionPoolGroup; + PoolGroupOptions = connectionPoolGroup.PoolGroupOptions; + ProviderInfo = connectionPoolProviderInfo; + Identity = identity; + AuthenticationContexts = new(); + MaxPoolSize = Convert.ToUInt32(PoolGroupOptions.MaxPoolSize); + + _connectionSlots = new(MaxPoolSize); + + // We enforce Max Pool Size, so no need to create a bounded channel (which is less efficient) + // On the consuming side, we have the multiplexing write loop but also non-multiplexing Rents + // On the producing side, we have connections being released back into the pool (both multiplexing and not) + var idleChannel = Channel.CreateUnbounded(); + _idleConnectionReader = idleChannel.Reader; + _idleConnectionWriter = idleChannel.Writer; + + State = Running; + } + #region Properties /// - public ConcurrentDictionary AuthenticationContexts => throw new NotImplementedException(); + public ConcurrentDictionary< + DbConnectionPoolAuthenticationContextKey, + DbConnectionPoolAuthenticationContext> AuthenticationContexts { get; } /// - public SqlConnectionFactory ConnectionFactory => throw new NotImplementedException(); + public SqlConnectionFactory ConnectionFactory { get; } /// - public int Count => throw new NotImplementedException(); + public int Count => _connectionSlots.ReservationCount; /// public bool ErrorOccurred => throw new NotImplementedException(); /// - public int Id => throw new NotImplementedException(); + public int Id => _instanceId; /// - public DbConnectionPoolIdentity Identity => throw new NotImplementedException(); + public DbConnectionPoolIdentity Identity { get; } /// - public bool IsRunning => throw new NotImplementedException(); + public bool IsRunning => State == Running; /// - public TimeSpan LoadBalanceTimeout => throw new NotImplementedException(); + public TimeSpan LoadBalanceTimeout => PoolGroupOptions.LoadBalanceTimeout; /// - public DbConnectionPoolGroup PoolGroup => throw new NotImplementedException(); + public DbConnectionPoolGroup PoolGroup { get; } /// - public DbConnectionPoolGroupOptions PoolGroupOptions => throw new NotImplementedException(); + public DbConnectionPoolGroupOptions PoolGroupOptions { get; } /// - public DbConnectionPoolProviderInfo ProviderInfo => throw new NotImplementedException(); + public DbConnectionPoolProviderInfo ProviderInfo { get; } /// - public DbConnectionPoolState State - { - get => throw new NotImplementedException(); - set => throw new NotImplementedException(); - } + public DbConnectionPoolState State { get; private set; } /// - public bool UseLoadBalancing => throw new NotImplementedException(); - #endregion - + public bool UseLoadBalancing => PoolGroupOptions.UseLoadBalancing; + private uint MaxPoolSize { get; } + #endregion #region Methods /// @@ -75,21 +136,48 @@ public void Clear() } /// - public void PutObjectFromTransactedPool(DbConnectionInternal obj) + public void PutObjectFromTransactedPool(DbConnectionInternal connection) { throw new NotImplementedException(); } /// - public DbConnectionInternal ReplaceConnection(DbConnection owningObject, DbConnectionOptions userOptions, DbConnectionInternal oldConnection) + public DbConnectionInternal ReplaceConnection( + DbConnection owningObject, + DbConnectionOptions userOptions, + DbConnectionInternal oldConnection) { throw new NotImplementedException(); } /// - public void ReturnInternalConnection(DbConnectionInternal obj, object owningObject) + public void ReturnInternalConnection(DbConnectionInternal connection, DbConnection owningObject) { - throw new NotImplementedException(); + ValidateOwnershipAndSetPoolingState(connection, owningObject); + + if (!IsLiveConnection(connection)) + { + RemoveConnection(connection); + return; + } + + SqlClientEventSource.Log.TryPoolerTraceEvent( + " {0}, Connection {1}, Deactivating.", + Id, + connection.ObjectID); + connection.DeactivateConnection(); + + if (connection.IsConnectionDoomed || + !connection.CanBePooled || + State == ShuttingDown) + { + RemoveConnection(connection); + } + else + { + var written = _idleConnectionWriter.TryWrite(connection); + Debug.Assert(written, "Failed to write returning connection to the idle channel."); + } } /// @@ -111,9 +199,378 @@ public void TransactionEnded(Transaction transaction, DbConnectionInternal trans } /// - public bool TryGetConnection(DbConnection owningObject, TaskCompletionSource taskCompletionSource, DbConnectionOptions userOptions, out DbConnectionInternal connection) + public bool TryGetConnection( + DbConnection owningObject, + TaskCompletionSource taskCompletionSource, + DbConnectionOptions userOptions, + out DbConnectionInternal? connection) { - throw new NotImplementedException(); + var timeout = TimeSpan.FromSeconds(owningObject.ConnectionTimeout); + + // If taskCompletionSource is null, we are in a sync context. + if (taskCompletionSource is null) + { + var task = GetInternalConnection( + owningObject, + userOptions, + async: false, + timeout); + + // 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 + // to GetInternalConnection, which means it will not use Task.Run or any async-await logic that would + // schedule tasks on the managed threadpool. + connection = task.ConfigureAwait(false).GetAwaiter().GetResult(); + return connection is not null; + } + + // Early exit if the task is already completed. + if (taskCompletionSource.Task.IsCompleted) + { + connection = null; + return false; + } + + // This is ugly, but async anti-patterns above and below us in the stack necessitate a fresh task to be + // created. Ideally we would just return the Task from GetInternalConnection and let the caller await + // it as needed, but instead we need to signal to the provided TaskCompletionSource when the connection + // is established. This pattern has implications for connection open retry logic that are intricate + // enough to merit dedicated work. For now, callers that need to open many connections asynchronously + // and in parallel *must* pre-prevision threads in the managed thread pool to avoid exhaustion and + // timeouts. + // + // Also note that we don't have access to the cancellation token passed by the caller to the original + // 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. + Task.Run(async () => + { + if (taskCompletionSource.Task.IsCompleted) + { + 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 + { + connection = await GetInternalConnection( + owningObject, + userOptions, + async: true, + timeout + ).ConfigureAwait(false); + + if (!taskCompletionSource.TrySetResult(connection)) + { + // We were able to get a connection, but the task was cancelled out from under us. + // This can happen if the caller's CancellationToken is cancelled while we're waiting for a connection. + // Check the success to avoid an unnecessary exception. + ReturnInternalConnection(connection, owningObject); + } + } + catch (Exception e) + { + if (connection != null) + { + ReturnInternalConnection(connection, owningObject); + } + + // It's possible to fail to set an exception on the TaskCompletionSource if the task is already + // completed. In that case, this exception will be swallowed because nobody directly awaits this + // task. + taskCompletionSource.TrySetException(e); + } + }); + + connection = null; + return false; + } + + private struct CreateState + { + internal ChannelDbConnectionPool pool; + internal DbConnection? owningConnection; + internal DbConnectionOptions userOptions; + } + + /// + /// Opens a new internal connection to the database. + /// + /// The owning connection. + /// The options for the connection. + /// The cancellation token to cancel the operation. + /// A task representing the asynchronous operation, with a result of the new internal connection. + /// + /// Thrown when the cancellation token is cancelled before the connection operation completes. + /// + private DbConnectionInternal? OpenNewInternalConnection( + DbConnection? owningConnection, + DbConnectionOptions userOptions, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + + // Opening a connection can be a slow operation and we don't want to hold a lock for the duration. + // Instead, we reserve a connection slot prior to attempting to open a new connection and release the slot + // in case of an exception. + + return _connectionSlots.Add( + createCallback: static (state) => + { + // https://github.com/dotnet/SqlClient/issues/3459 + // TODO: This blocks the thread for several network calls! + // When running async, the blocked thread is one allocated from the managed thread pool (due to + // use of Task.Run in TryGetConnection). This is why it's critical for async callers to + // pre-provision threads in the managed thread pool. Our options are limited because + // DbConnectionInternal doesn't support an async open. It's better to block this thread and keep + // throughput high than to queue all of our opens onto a single worker thread. Add an async path + // when this support is added to DbConnectionInternal. + return state.pool.ConnectionFactory.CreatePooledConnection( + state.owningConnection, + state.pool, + state.pool.PoolGroup.PoolKey, + state.pool.PoolGroup.ConnectionOptions, + state.userOptions); + }, + cleanupCallback: static (newConnection, idleConnectionWriter) => + { + idleConnectionWriter.TryWrite(null); + newConnection?.Dispose(); + }, + new CreateState + { + pool = this, + owningConnection = owningConnection, + userOptions = userOptions + }, + _idleConnectionWriter); + } + + /// + /// Checks that the provided connection is live and unexpired and closes it if needed. + /// + /// + /// Returns true if the connection is live and unexpired, otherwise returns false. + private bool IsLiveConnection(DbConnectionInternal connection) + { + if (!connection.IsConnectionAlive()) + { + return false; + } + + if (LoadBalanceTimeout != TimeSpan.Zero && DateTime.UtcNow > connection.CreateTime + LoadBalanceTimeout) + { + return false; + } + + return true; + } + + /// + /// Closes the provided connection and removes it from the pool. + /// + /// The connection to be closed. + private void RemoveConnection(DbConnectionInternal connection) + { + _connectionSlots.TryRemove(connection); + + // Removing a connection from the pool opens a free slot. + // Write a null to the idle connection channel to wake up a waiter, who can now open a new + // connection. Statement order is important since we have synchronous completions on the channel. + _idleConnectionWriter.TryWrite(null); + + connection.Dispose(); + } + + /// + /// Tries to read a connection from the idle connection channel. + /// + /// A connection from the idle channel, or null if the channel is empty. + private DbConnectionInternal? GetIdleConnection() + { + // The channel may contain nulls. Read until we find a non-null connection or exhaust the channel. + while (_idleConnectionReader.TryRead(out DbConnectionInternal? connection)) + { + if (connection is null) + { + continue; + } + + if (!IsLiveConnection(connection)) + { + RemoveConnection(connection); + continue; + } + + return connection; + } + + return null; + } + + /// + /// Gets an internal connection from the pool, either by retrieving an idle connection or opening a new one. + /// + /// The DbConnection that will own this internal connection + /// The user options to set on the internal connection + /// A boolean indicating whether the operation should be asynchronous. + /// The timeout for the operation. + /// Returns a DbConnectionInternal that is retrieved from the pool. + /// + /// Thrown when an OperationCanceledException is caught, indicating that the timeout period + /// elapsed prior to obtaining a connection from the pool. + /// + /// + /// Thrown when a ChannelClosedException is caught, indicating that the connection pool + /// has been shut down. + /// + private async Task GetInternalConnection( + DbConnection owningConnection, + DbConnectionOptions userOptions, + bool async, + TimeSpan timeout) + { + DbConnectionInternal? connection = null; + using CancellationTokenSource cancellationTokenSource = new(timeout); + CancellationToken cancellationToken = cancellationTokenSource.Token; + + // Continue looping until we create or retrieve a connection + do + { + try + { + // Optimistically try to get an idle connection from the channel + // Doesn't wait if the channel is empty, just returns null. + connection ??= GetIdleConnection(); + + + // If we didn't find an idle connection, try to open a new one. + connection ??= OpenNewInternalConnection( + owningConnection, + userOptions, + cancellationToken); + + // If we're at max capacity and couldn't open a connection. Block on the idle channel with a + // timeout. Note that Channels guarantee fair FIFO behavior to callers of ReadAsync + // (first-come, first-served), which is crucial to us. + if (async) + { + connection ??= await _idleConnectionReader.ReadAsync(cancellationToken).ConfigureAwait(false); + } + else + { + connection ??= ReadChannelSyncOverAsync(cancellationToken); + } + } + catch (OperationCanceledException) + { + throw ADP.PooledOpenTimeout(); + } + catch (ChannelClosedException) + { + //TODO: exceptions from resource file + throw new Exception("The connection pool has been shut down."); + } + + if (connection is not null && !IsLiveConnection(connection)) + { + // If the connection is not live, we need to remove it from the pool and try again. + RemoveConnection(connection); + connection = null; + } + } + while (connection is null); + + PrepareConnection(owningConnection, connection); + return connection; + } + + /// + /// Performs a blocking synchronous read from the idle connection channel. + /// + /// Cancels the read operation. + /// The connection read from the channel. + private DbConnectionInternal? ReadChannelSyncOverAsync(CancellationToken cancellationToken) + { + // If there are no connections in the channel, then ReadAsync will block until one is available. + // Channels doesn't offer a sync API, so running ReadAsync synchronously on this thread may spawn + // additional new async work items in the managed thread pool if there are no items available in the + // channel. We need to ensure that we don't block all available managed threads with these child + // tasks or we could deadlock. Prefer to block the current user-owned thread, and limit throughput + // to the managed threadpool. + + _syncOverAsyncSemaphore.Wait(cancellationToken); + try + { + ConfiguredValueTaskAwaitable.ConfiguredValueTaskAwaiter awaiter = + _idleConnectionReader.ReadAsync(cancellationToken).ConfigureAwait(false).GetAwaiter(); + using ManualResetEventSlim mres = new ManualResetEventSlim(false, 0); + + // Cancellation happens through the ReadAsync call, which will complete the task. + // Even a failed task will complete and set the ManualResetEventSlim. + awaiter.UnsafeOnCompleted(() => mres.Set()); + mres.Wait(CancellationToken.None); + return awaiter.GetResult(); + } + finally + { + _syncOverAsyncSemaphore.Release(); + } + } + + /// + /// Sets connection state and activates the connection for use. Should always be called after a connection is + /// created or retrieved from the pool. + /// + /// The owning DbConnection instance. + /// The DbConnectionInternal to be activated. + /// + /// Thrown when any exception occurs during connection activation. + /// + private void PrepareConnection(DbConnection owningObject, DbConnectionInternal connection) + { + lock (connection) + { + // Protect against Clear which calls IsEmancipated, which is affected by PrePush and PostPop + connection.PostPop(owningObject); + } + + try + { + //TODO: pass through transaction + connection.ActivateConnection(null); + } + catch + { + // At this point, the connection is "out of the pool" (the call to postpop). If we hit a transient + // error anywhere along the way when enlisting the connection in the transaction, we need to get + // the connection back into the pool so that it isn't leaked. + ReturnInternalConnection(connection, owningObject); + throw; + } + } + + /// + /// Validates that the connection is owned by the provided DbConnection and that it is in a valid state to be returned to the pool. + /// + /// The owning DbConnection instance. + /// The DbConnectionInternal to be validated. + private void ValidateOwnershipAndSetPoolingState(DbConnectionInternal connection, DbConnection? owningObject) + { + lock (connection) + { + // Calling PrePush prevents the object from being reclaimed + // once we leave the lock, because it sets _pooledCount such + // that it won't appear to be out of the pool. What that + // means, is that we're now responsible for this connection: + // it won't get reclaimed if it gets lost. + connection.PrePush(owningObject); + } } #endregion } diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ConnectionPoolSlots.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ConnectionPoolSlots.cs new file mode 100644 index 0000000000..e5935ddee5 --- /dev/null +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ConnectionPoolSlots.cs @@ -0,0 +1,200 @@ +// 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.Diagnostics; +using System.Threading; +using Microsoft.Data.ProviderBase; + +#nullable enable + +namespace Microsoft.Data.SqlClient.ConnectionPool +{ + /// + /// A thread-safe collection with a fixed capacity. Avoids wasted work by reserving a slot before adding an item. + /// + internal sealed class ConnectionPoolSlots + { + /// + /// Represents a reservation that manages a resource and ensures cleanup when no longer needed. + /// + /// The type of the resource being managed by the reservation. + private sealed class Reservation : IDisposable + { + private Action cleanupCallback; + private T state; + private bool _retain = false; + private bool _disposed = false; + + internal Reservation(T state, Action cleanupCallback) + { + this.state = state; + this.cleanupCallback = cleanupCallback; + } + + public void Dispose() + { + if (_disposed) + { + return; + } + + _disposed = true; + + if (!_retain) + { + cleanupCallback(state); + } + } + + internal void Keep() + { + _retain = true; + } + } + + internal delegate T CreateCallback(S state); + internal delegate void CleanupCallback(DbConnectionInternal? connection, T state); + + private readonly DbConnectionInternal?[] _connections; + private readonly uint _capacity; + private volatile int _reservations; + + /// + /// Constructs a ConnectionPoolSlots instance with the given fixed capacity. + /// + /// The fixed capacity of the collection. + /// + /// Thrown when fixedCapacity is greater than Int32.MaxValue or equal to zero. + /// + internal ConnectionPoolSlots(uint fixedCapacity) + { + if (fixedCapacity > int.MaxValue) + { + throw new ArgumentOutOfRangeException(nameof(fixedCapacity), "Capacity must be less than or equal to Int32.MaxValue."); + } + + if (fixedCapacity == 0) + { + throw new ArgumentOutOfRangeException(nameof(fixedCapacity), "Capacity must be greater than zero."); + } + + _capacity = fixedCapacity; + _reservations = 0; + _connections = new DbConnectionInternal?[fixedCapacity]; + } + + /// + /// Gets the total number of reservations currently held. + /// + internal int ReservationCount => _reservations; + + /// + /// Adds a connection to the collection. + /// + /// Callback that provides the connection to add to the collection. This callback + /// *must not* call any other ConnectionPoolSlots methods. + /// Callback to clean up resources if an exception occurs. This callback *must + /// not* call any other ConnectionPoolSlots methods. This callback *must not* throw exceptions. + /// State made available to the create callback. + /// State made available to the cleanup callback. + /// + /// Throws when createCallback throws an exception. + /// Throws when a reservation is successfully made, but an empty slot cannot be found. This condition is + /// unexpected and indicates a bug. + /// + /// Returns the new connection, or null if there was not available space. + internal DbConnectionInternal? Add( + CreateCallback createCallback, + CleanupCallback cleanupCallback, + T createState, + S cleanupState) + { + DbConnectionInternal? connection = null; + try + { + using var reservation = TryReserve(); + if (reservation is null) + { + return null; + } + + connection = createCallback(createState); + + if (connection is null) + { + return null; + } + + for (int i = 0; i < _capacity; i++) + { + if (Interlocked.CompareExchange(ref _connections[i], connection, null) == null) + { + reservation.Keep(); + return connection; + } + } + + throw new InvalidOperationException("Couldn't find an empty slot."); + } + catch (Exception e) + { + cleanupCallback(connection, cleanupState); + throw; + } + } + + /// + /// Releases a reservation that was previously obtained. + /// Must be called after removing a connection from the collection or if an exception occurs. + /// + private void ReleaseReservation() + { + Interlocked.Decrement(ref _reservations); + Debug.Assert(_reservations >= 0, "Released a reservation that wasn't held"); + } + + /// + /// Removes a connection from the collection. + /// + /// The connection to remove from the collection. + /// True if the connection was found and removed; otherwise, false. + internal bool TryRemove(DbConnectionInternal connection) + { + for (int i = 0; i < _connections.Length; i++) + { + if (Interlocked.CompareExchange(ref _connections[i], null, connection) == connection) + { + ReleaseReservation(); + return true; + } + } + + return false; + } + + /// + /// Attempts to reserve a spot in the collection. + /// + /// A Reservation if successful, otherwise returns null. + private Reservation? TryReserve() + { + for (var expected = _reservations; expected < _capacity; expected = _reservations) + { + // Try to reserve a spot in the collection by incrementing _reservations. + // If _reservations changed underneath us, then another thread already reserved the spot we were trying to take. + // Cycle back through the check above to reset expected and to make sure we don't go + // over capacity. + // Note that we purposefully don't use SpinWait for this: https://github.com/dotnet/coreclr/pull/21437 + if (Interlocked.CompareExchange(ref _reservations, expected + 1, expected) != expected) + { + continue; + } + + return new Reservation(this, (slots) => slots.ReleaseReservation()); + } + return null; + } + } +} diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/DbConnectionPoolState.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/DbConnectionPoolState.cs index 1790e38a57..ca4353f2c6 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/DbConnectionPoolState.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/DbConnectionPoolState.cs @@ -6,7 +6,6 @@ namespace Microsoft.Data.SqlClient.ConnectionPool { internal enum DbConnectionPoolState { - Initializing, Running, ShuttingDown, } diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/IDbConnectionPool.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/IDbConnectionPool.cs index 066318fce2..b684bb24bb 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/IDbConnectionPool.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/IDbConnectionPool.cs @@ -10,6 +10,8 @@ using Microsoft.Data.Common.ConnectionString; using Microsoft.Data.ProviderBase; +#nullable enable + namespace Microsoft.Data.SqlClient.ConnectionPool { /// @@ -81,7 +83,7 @@ internal interface IDbConnectionPool /// /// The current state of the connection pool. /// - DbConnectionPoolState State { get; set; } + DbConnectionPoolState State { get; } /// /// Indicates whether the connection pool is using load balancing. @@ -104,7 +106,7 @@ internal interface IDbConnectionPool /// The user options to use if a new connection must be opened. /// The retrieved connection will be passed out via this parameter. /// True if a connection was set in the out parameter, otherwise returns false. - bool TryGetConnection(DbConnection owningObject, TaskCompletionSource taskCompletionSource, DbConnectionOptions userOptions, out DbConnectionInternal connection); + bool TryGetConnection(DbConnection owningObject, TaskCompletionSource taskCompletionSource, DbConnectionOptions userOptions, out DbConnectionInternal? connection); /// /// Replaces the internal connection currently associated with owningObject with a new internal connection from the pool. @@ -120,7 +122,7 @@ internal interface IDbConnectionPool /// /// The internal connection to return to the pool. /// The connection that currently owns this internal connection. Used to verify ownership. - void ReturnInternalConnection(DbConnectionInternal obj, object owningObject); + void ReturnInternalConnection(DbConnectionInternal obj, DbConnection owningObject); /// /// Puts an internal connection from a transacted pool back into the general pool. diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/WaitHandleDbConnectionPool.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/WaitHandleDbConnectionPool.cs index bd6ebf708e..e622bc1bff 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/WaitHandleDbConnectionPool.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/WaitHandleDbConnectionPool.cs @@ -451,8 +451,6 @@ internal WaitHandleDbConnectionPool( throw ADP.InternalError(ADP.InternalErrorCode.AttemptingToPoolOnRestrictedToken); } - State = Initializing; - lock (s_random) { // Random.Next is not thread-safe @@ -750,17 +748,8 @@ private DbConnectionInternal CreateObject(DbConnection owningObject, DbConnectio owningObject, this, _connectionPoolGroup.PoolKey, - _connectionPoolGroup.ConnectionOptions, + _connectionPoolGroup.ConnectionOptions, userOptions); - if (newObj == null) - { - throw ADP.InternalError(ADP.InternalErrorCode.CreateObjectReturnedNull); // CreateObject succeeded, but null object - } - if (!newObj.CanBePooled) - { - throw ADP.InternalError(ADP.InternalErrorCode.NewObjectCannotBePooled); // CreateObject succeeded, but non-poolable object - } - newObj.PrePush(null); lock (_objectList) { @@ -859,10 +848,6 @@ private void DeactivateObject(DbConnectionInternal obj) } else { - // NOTE: constructor should ensure that current state cannot be State.Initializing, so it can only - // be State.Running or State.ShuttingDown - Debug.Assert(State is Running or ShuttingDown); - lock (obj) { // A connection with a delegated transaction cannot currently @@ -1666,7 +1651,7 @@ private void PutNewObject(DbConnectionInternal obj) } - public void ReturnInternalConnection(DbConnectionInternal obj, object owningObject) + public void ReturnInternalConnection(DbConnectionInternal obj, DbConnection owningObject) { Debug.Assert(obj != null, "null obj?"); diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlConnectionFactory.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlConnectionFactory.cs index 6f75f8c75f..4b59dec944 100644 --- a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlConnectionFactory.cs +++ b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlConnectionFactory.cs @@ -49,7 +49,7 @@ internal class SqlConnectionFactory #region Constructors - private SqlConnectionFactory() + protected SqlConnectionFactory() { _connectionPoolGroups = new Dictionary(); _poolsToRelease = new List(); @@ -111,7 +111,7 @@ internal DbConnectionPoolProviderInfo CreateConnectionPoolProviderInfo(DbConnect ? new SqlConnectionPoolProviderInfo() : null; - internal SqlInternalConnectionTds CreateNonPooledConnection( + internal DbConnectionInternal CreateNonPooledConnection( DbConnection owningConnection, DbConnectionPoolGroup poolGroup, DbConnectionOptions userOptions) @@ -119,7 +119,7 @@ internal SqlInternalConnectionTds CreateNonPooledConnection( Debug.Assert(owningConnection is not null, "null owningConnection?"); Debug.Assert(poolGroup is not null, "null poolGroup?"); - SqlInternalConnectionTds newConnection = CreateConnection( + DbConnectionInternal newConnection = CreateConnection( poolGroup.ConnectionOptions, poolGroup.PoolKey, poolGroup.ProviderInfo, @@ -136,7 +136,7 @@ internal SqlInternalConnectionTds CreateNonPooledConnection( return newConnection; } - internal SqlInternalConnectionTds CreatePooledConnection( + internal DbConnectionInternal CreatePooledConnection( DbConnection owningConnection, IDbConnectionPool pool, DbConnectionPoolKey poolKey, @@ -145,20 +145,31 @@ internal SqlInternalConnectionTds CreatePooledConnection( { Debug.Assert(pool != null, "null pool?"); - SqlInternalConnectionTds newConnection = CreateConnection( + DbConnectionInternal newConnection = CreateConnection( options, poolKey, // @TODO: is pool.PoolGroup.Key the same thing? pool.PoolGroup.ProviderInfo, pool, owningConnection, userOptions); - if (newConnection is not null) + + if (newConnection is null) { - SqlClientEventSource.Metrics.HardConnectRequest(); - newConnection.MakePooledConnection(pool); + throw ADP.InternalError(ADP.InternalErrorCode.CreateObjectReturnedNull); // CreateObject succeeded, but null object } - + + if (!newConnection.CanBePooled) + { + throw ADP.InternalError(ADP.InternalErrorCode.NewObjectCannotBePooled); // CreateObject succeeded, but non-poolable object + } + + SqlClientEventSource.Metrics.HardConnectRequest(); + newConnection.MakePooledConnection(pool); + SqlClientEventSource.Log.TryTraceEvent(" {0}, Pooled database connection created.", ObjectId); + + newConnection.PrePush(null); + return newConnection; } @@ -576,7 +587,7 @@ internal void SetInnerConnectionTo(DbConnection owningObject, DbConnectionIntern #region Private Methods // @TODO: I think this could be broken down into methods more specific to use cases above - private static SqlInternalConnectionTds CreateConnection( + protected virtual DbConnectionInternal CreateConnection( DbConnectionOptions options, DbConnectionPoolKey poolKey, DbConnectionPoolGroupProviderInfo poolGroupProviderInfo, diff --git a/src/Microsoft.Data.SqlClient/src/System/Runtime/CompilerServices/IsExternalInit.netfx.cs b/src/Microsoft.Data.SqlClient/src/System/Runtime/CompilerServices/IsExternalInit.netfx.cs index 0d0181ba6d..530d1b4096 100644 --- a/src/Microsoft.Data.SqlClient/src/System/Runtime/CompilerServices/IsExternalInit.netfx.cs +++ b/src/Microsoft.Data.SqlClient/src/System/Runtime/CompilerServices/IsExternalInit.netfx.cs @@ -6,7 +6,7 @@ using System.ComponentModel; - +// The `init` accessor was introduced in C# 9.0 and is not natively supported in .NET Framework. // This class enables the use of the `init` property accessor in .NET framework. namespace System.Runtime.CompilerServices { diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPoolTest.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPoolTest.cs index 2dcfe476fe..d258250327 100644 --- a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPoolTest.cs +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPoolTest.cs @@ -3,153 +3,911 @@ // See the LICENSE file in the project root for more information. using System; +using System.Collections.Concurrent; using System.Data.Common; +using System.Threading; using System.Threading.Tasks; using System.Transactions; +using Microsoft.Data.Common; +using Microsoft.Data.Common.ConnectionString; +using Microsoft.Data.ProviderBase; using Microsoft.Data.SqlClient.ConnectionPool; using Xunit; -namespace Microsoft.Data.SqlClient.UnitTests +namespace Microsoft.Data.SqlClient.UnitTests.ConnectionPool { public class ChannelDbConnectionPoolTest { - private readonly ChannelDbConnectionPool _pool; + private ChannelDbConnectionPool pool; + private SqlConnectionFactory connectionFactory; + private DbConnectionPoolGroup dbConnectionPoolGroup; + private DbConnectionPoolGroupOptions poolGroupOptions; + private DbConnectionPoolIdentity identity; + private DbConnectionPoolProviderInfo connectionPoolProviderInfo; - public ChannelDbConnectionPoolTest() + private static readonly SqlConnectionFactory SuccessfulConnectionFactory = new SuccessfulSqlConnectionFactory(); + private static readonly SqlConnectionFactory TimeoutConnectionFactory = new TimeoutSqlConnectionFactory(); + + private void Setup(SqlConnectionFactory connectionFactory) + { + this.connectionFactory = connectionFactory; + identity = DbConnectionPoolIdentity.NoIdentity; + connectionPoolProviderInfo = new DbConnectionPoolProviderInfo(); + poolGroupOptions = new DbConnectionPoolGroupOptions( + poolByIdentity: false, + minPoolSize: 0, + maxPoolSize: 50, + creationTimeout: 15, + loadBalanceTimeout: 0, + hasTransactionAffinity: true + ); + dbConnectionPoolGroup = new DbConnectionPoolGroup( + new DbConnectionOptions("DataSource=localhost;", null), + new DbConnectionPoolKey("TestDataSource"), + poolGroupOptions + ); + pool = new ChannelDbConnectionPool( + connectionFactory, + dbConnectionPoolGroup, + identity, + connectionPoolProviderInfo + ); + } + + [Theory] + [InlineData(1)] + [InlineData(5)] + [InlineData(10)] + public void GetConnectionEmptyPool_ShouldCreateNewConnection(int numConnections) + { + // Arrange + Setup(SuccessfulConnectionFactory); + + // Act + for (int i = 0; i < numConnections; i++) + { + DbConnectionInternal internalConnection = null; + var completed = pool.TryGetConnection( + new SqlConnection(), + taskCompletionSource: null, + new DbConnectionOptions("", null), + out internalConnection + ); + + // Assert + Assert.True(completed); + Assert.NotNull(internalConnection); + } + + + // Assert + Assert.Equal(numConnections, pool.Count); + } + + [Theory] + [InlineData(1)] + [InlineData(5)] + [InlineData(10)] + public async Task GetConnectionAsyncEmptyPool_ShouldCreateNewConnection(int numConnections) + { + // Arrange + Setup(SuccessfulConnectionFactory); + + // Act + for (int i = 0; i < numConnections; i++) + { + var tcs = new TaskCompletionSource(); + DbConnectionInternal internalConnection = null; + var completed = pool.TryGetConnection( + new SqlConnection(), + tcs, + new DbConnectionOptions("", null), + out internalConnection + ); + + // Assert + Assert.False(completed); + Assert.Null(internalConnection); + Assert.NotNull(await tcs.Task); + } + + + // Assert + Assert.Equal(numConnections, pool.Count); + } + + [Fact] + public void GetConnectionMaxPoolSize_ShouldTimeoutAfterPeriod() + { + // Arrange + Setup(SuccessfulConnectionFactory); + + for (int i = 0; i < poolGroupOptions.MaxPoolSize; i++) + { + DbConnectionInternal internalConnection = null; + var completed = pool.TryGetConnection( + new SqlConnection(), + taskCompletionSource: null, + new DbConnectionOptions("", null), + out internalConnection + ); + + Assert.True(completed); + Assert.NotNull(internalConnection); + } + + try + { + // Act + DbConnectionInternal extraConnection = null; + var exceeded = pool.TryGetConnection( + new SqlConnection("Timeout=1"), + taskCompletionSource: null, + new DbConnectionOptions("", null), + out extraConnection + ); + } + catch (Exception ex) + { + // Assert + Assert.IsType(ex); + Assert.Equal("Timeout expired. The timeout period elapsed prior to obtaining a connection from the pool. This may have occurred because all pooled connections were in use and max pool size was reached.", ex.Message); + } + + // Assert + Assert.Equal(poolGroupOptions.MaxPoolSize, pool.Count); + } + + [Fact] + public async Task GetConnectionAsyncMaxPoolSize_ShouldTimeoutAfterPeriod() + { + // Arrange + Setup(SuccessfulConnectionFactory); + + for (int i = 0; i < poolGroupOptions.MaxPoolSize; i++) + { + DbConnectionInternal internalConnection = null; + var completed = pool.TryGetConnection( + new SqlConnection(), + taskCompletionSource: null, + new DbConnectionOptions("", null), + out internalConnection + ); + + Assert.True(completed); + Assert.NotNull(internalConnection); + } + + try + { + // Act + TaskCompletionSource taskCompletionSource = new TaskCompletionSource(); + DbConnectionInternal extraConnection = null; + var exceeded = pool.TryGetConnection( + new SqlConnection("Timeout=1"), + taskCompletionSource, + new DbConnectionOptions("", null), + out extraConnection + ); + await taskCompletionSource.Task; + } + catch (Exception ex) + { + // Assert + Assert.IsType(ex); + Assert.Equal("Timeout expired. The timeout period elapsed prior to obtaining a connection from the pool. This may have occurred because all pooled connections were in use and max pool size was reached.", ex.Message); + } + + // Assert + Assert.Equal(poolGroupOptions.MaxPoolSize, pool.Count); + } + + [Fact] + public async Task GetConnectionMaxPoolSize_ShouldReuseAfterConnectionReleased() + { + // Arrange + Setup(SuccessfulConnectionFactory); + DbConnectionInternal firstConnection = null; + SqlConnection firstOwningConnection = new SqlConnection(); + + pool.TryGetConnection( + firstOwningConnection, + taskCompletionSource: null, + new DbConnectionOptions("", null), + out firstConnection + ); + + for (int i = 1; i < poolGroupOptions.MaxPoolSize; i++) + { + DbConnectionInternal internalConnection = null; + var completed = pool.TryGetConnection( + new SqlConnection(), + taskCompletionSource: null, + new DbConnectionOptions("", null), + out internalConnection + ); + + Assert.True(completed); + Assert.NotNull(internalConnection); + } + + TaskCompletionSource tcs = new TaskCompletionSource(); + + // Act + var task = Task.Run(() => + { + DbConnectionInternal extraConnection = null; + var exceeded = pool.TryGetConnection( + new SqlConnection(""), + taskCompletionSource: null, + new DbConnectionOptions("", null), + out extraConnection + ); + return extraConnection; + }); + pool.ReturnInternalConnection(firstConnection, firstOwningConnection); + var extraConnection = await task; + + // Assert + Assert.Equal(firstConnection, extraConnection); + } + + [Fact] + public async Task GetConnectionAsyncMaxPoolSize_ShouldReuseAfterConnectionReleased() + { + // Arrange + Setup(SuccessfulConnectionFactory); + DbConnectionInternal firstConnection = null; + SqlConnection firstOwningConnection = new SqlConnection(); + + pool.TryGetConnection( + firstOwningConnection, + taskCompletionSource: null, + new DbConnectionOptions("", null), + out firstConnection + ); + + for (int i = 1; i < poolGroupOptions.MaxPoolSize; i++) + { + DbConnectionInternal internalConnection = null; + var completed = pool.TryGetConnection( + new SqlConnection(), + taskCompletionSource: null, + new DbConnectionOptions("", null), + out internalConnection + ); + + Assert.True(completed); + Assert.NotNull(internalConnection); + } + + TaskCompletionSource taskCompletionSource = new TaskCompletionSource(); + + // Act + DbConnectionInternal recycledConnection = null; + var exceeded = pool.TryGetConnection( + new SqlConnection(""), + taskCompletionSource, + new DbConnectionOptions("", null), + out recycledConnection + ); + pool.ReturnInternalConnection(firstConnection, firstOwningConnection); + recycledConnection = await taskCompletionSource.Task; + + // Assert + Assert.Equal(firstConnection, recycledConnection); + } + + [Fact] + public async Task GetConnectionMaxPoolSize_ShouldRespectOrderOfRequest() + { + // Arrange + Setup(SuccessfulConnectionFactory); + DbConnectionInternal firstConnection = null; + SqlConnection firstOwningConnection = new SqlConnection(); + + pool.TryGetConnection( + firstOwningConnection, + taskCompletionSource: null, + new DbConnectionOptions("", null), + out firstConnection + ); + + for (int i = 1; i < poolGroupOptions.MaxPoolSize; i++) + { + DbConnectionInternal internalConnection = null; + var completed = pool.TryGetConnection( + new SqlConnection(), + taskCompletionSource: null, + new DbConnectionOptions("", null), + out internalConnection + ); + + Assert.True(completed); + Assert.NotNull(internalConnection); + } + + // Use ManualResetEventSlim to synchronize the tasks + // and force the request queueing order. + using ManualResetEventSlim mresQueueOrder = new ManualResetEventSlim(); + using CountdownEvent allRequestsQueued = new CountdownEvent(2); + + // Act + var recycledTask = Task.Run(() => + { + DbConnectionInternal recycledConnection = null; + mresQueueOrder.Set(); + allRequestsQueued.Signal(); + pool.TryGetConnection( + new SqlConnection(""), + null, + new DbConnectionOptions("", null), + out recycledConnection + ); + return recycledConnection; + }); + var failedTask = Task.Run(() => + { + DbConnectionInternal failedConnection = null; + // Force this request to be second in the queue. + mresQueueOrder.Wait(); + allRequestsQueued.Signal(); + pool.TryGetConnection( + new SqlConnection("Timeout=1"), + null, + new DbConnectionOptions("", null), + out failedConnection + ); + return failedConnection; + }); + + allRequestsQueued.Wait(); + pool.ReturnInternalConnection(firstConnection, firstOwningConnection); + var recycledConnection = await recycledTask; + + // Assert + Assert.Equal(firstConnection, recycledConnection); + await Assert.ThrowsAsync(async () => await failedTask); + } + + [Fact] + public async Task GetConnectionAsyncMaxPoolSize_ShouldRespectOrderOfRequest() + { + // Arrange + Setup(SuccessfulConnectionFactory); + DbConnectionInternal firstConnection = null; + SqlConnection firstOwningConnection = new SqlConnection(); + + pool.TryGetConnection( + firstOwningConnection, + taskCompletionSource: null, + new DbConnectionOptions("", null), + out firstConnection + ); + + for (int i = 1; i < poolGroupOptions.MaxPoolSize; i++) + { + DbConnectionInternal internalConnection = null; + var completed = pool.TryGetConnection( + new SqlConnection(), + taskCompletionSource: null, + new DbConnectionOptions("", null), + out internalConnection + ); + + Assert.True(completed); + Assert.NotNull(internalConnection); + } + + TaskCompletionSource recycledTaskCompletionSource = new TaskCompletionSource(); + TaskCompletionSource failedCompletionSource = new TaskCompletionSource(); + + // Act + DbConnectionInternal recycledConnection = null; + var exceeded = pool.TryGetConnection( + new SqlConnection(""), + recycledTaskCompletionSource, + new DbConnectionOptions("", null), + out recycledConnection + ); + + // Gives time for the recycled connection to be queued before the failed request is initiated. + await Task.Delay(1000); + + DbConnectionInternal failedConnection = null; + var exceeded2 = pool.TryGetConnection( + new SqlConnection("Timeout=1"), + failedCompletionSource, + new DbConnectionOptions("", null), + out failedConnection + ); + + pool.ReturnInternalConnection(firstConnection, firstOwningConnection); + recycledConnection = await recycledTaskCompletionSource.Task; + + // Assert + Assert.Equal(firstConnection, recycledConnection); + await Assert.ThrowsAsync(async () => failedConnection = await failedCompletionSource.Task); + } + + [Fact] + public void ConnectionsAreReused() { - _pool = new ChannelDbConnectionPool(); + // Arrange + Setup(SuccessfulConnectionFactory); + SqlConnection owningConnection = new SqlConnection(); + DbConnectionInternal internalConnection1 = null; + DbConnectionInternal internalConnection2 = null; + + // Act: Get the first connection + var completed1 = pool.TryGetConnection( + owningConnection, + null, + new DbConnectionOptions("", null), + out internalConnection1 + ); + + // Assert: First connection should succeed + Assert.True(completed1); + Assert.NotNull(internalConnection1); + + // Act: Return the first connection to the pool + pool.ReturnInternalConnection(internalConnection1, owningConnection); + + // Act: Get the second connection (should reuse the first one) + var completed2 = pool.TryGetConnection( + owningConnection, + null, + new DbConnectionOptions("", null), + out internalConnection2 + ); + + // Assert: Second connection should succeed and reuse the first connection + Assert.True(completed2); + Assert.NotNull(internalConnection2); + Assert.Same(internalConnection1, internalConnection2); } [Fact] - public void TestAuthenticationContexts() + public void GetConnectionTimeout_ShouldThrowTimeoutException() { - Assert.Throws(() => _ = _pool.AuthenticationContexts); + // Arrange + Setup(TimeoutConnectionFactory); + DbConnectionInternal internalConnection = null; + + // Act & Assert + var ex = Assert.Throws(() => + { + var completed = pool.TryGetConnection( + new SqlConnection(), + taskCompletionSource: null, + new DbConnectionOptions("", null), + out internalConnection + ); + }); + + Assert.IsType(ex.InnerException); + + Assert.Equal("Timeout expired. The timeout period elapsed prior to obtaining a connection from the pool. This may have occurred because all pooled connections were in use and max pool size was reached.", ex.InnerException.Message); } + [Fact] + public async Task GetConnectionAsyncTimeout_ShouldThrowTimeoutException() + { + // Arrange + Setup(TimeoutConnectionFactory); + DbConnectionInternal internalConnection = null; + TaskCompletionSource taskCompletionSource = new TaskCompletionSource(); + + // Act & Assert + var ex = await Assert.ThrowsAsync(async () => + { + var completed = pool.TryGetConnection( + new SqlConnection(), + taskCompletionSource, + new DbConnectionOptions("", null), + out internalConnection + ); + + await taskCompletionSource.Task; + }); + + Assert.IsType(ex.InnerException); + + Assert.Equal("Timeout expired. The timeout period elapsed prior to obtaining a connection from the pool. This may have occurred because all pooled connections were in use and max pool size was reached.", ex.InnerException.Message); + } + + [Fact] + public void StressTest() + { + //Arrange + Setup(SuccessfulConnectionFactory); + ConcurrentBag tasks = new ConcurrentBag(); + + + for (int i = 1; i < poolGroupOptions.MaxPoolSize * 3; i++) + { + var t = Task.Run(() => + { + DbConnectionInternal internalConnection = null; + SqlConnection owningObject = new SqlConnection(); + var completed = pool.TryGetConnection( + owningObject, + taskCompletionSource: null, + new DbConnectionOptions("", null), + out internalConnection + ); + if (completed) + { + pool.ReturnInternalConnection(internalConnection, owningObject); + } + + Assert.True(completed); + Assert.NotNull(internalConnection); + }); + tasks.Add(t); + } + + Task.WaitAll(tasks.ToArray()); + Assert.True(pool.Count <= poolGroupOptions.MaxPoolSize, "Pool size exceeded max pool size after stress test."); + } + + [Fact] + public void StressTestAsync() + { + //Arrange + Setup(SuccessfulConnectionFactory); + ConcurrentBag tasks = new ConcurrentBag(); + + + for (int i = 1; i < poolGroupOptions.MaxPoolSize * 3; i++) + { + var t = Task.Run(async () => + { + DbConnectionInternal internalConnection = null; + SqlConnection owningObject = new SqlConnection(); + TaskCompletionSource taskCompletionSource = new TaskCompletionSource(); + var completed = pool.TryGetConnection( + owningObject, + taskCompletionSource, + new DbConnectionOptions("", null), + out internalConnection + ); + internalConnection = await taskCompletionSource.Task; + pool.ReturnInternalConnection(internalConnection, owningObject); + + Assert.NotNull(internalConnection); + }); + tasks.Add(t); + } + + Task.WaitAll(tasks.ToArray()); + Assert.True(pool.Count <= poolGroupOptions.MaxPoolSize, "Pool size exceeded max pool size after stress test."); + } + + + #region Property Tests + [Fact] public void TestConnectionFactory() { - Assert.Throws(() => _ = _pool.ConnectionFactory); + Setup(SuccessfulConnectionFactory); + Assert.Equal(connectionFactory, pool.ConnectionFactory); } [Fact] public void TestCount() { - Assert.Throws(() => _ = _pool.Count); + Setup(SuccessfulConnectionFactory); + Assert.Equal(0, pool.Count); } [Fact] public void TestErrorOccurred() { - Assert.Throws(() => _ = _pool.ErrorOccurred); + Setup(SuccessfulConnectionFactory); + Assert.Throws(() => _ = pool.ErrorOccurred); } [Fact] public void TestId() { - Assert.Throws(() => _ = _pool.Id); + Setup(SuccessfulConnectionFactory); + Assert.True(pool.Id >= 1); } [Fact] public void TestIdentity() { - Assert.Throws(() => _ = _pool.Identity); + Setup(SuccessfulConnectionFactory); + Assert.Equal(identity, pool.Identity); } [Fact] public void TestIsRunning() { - Assert.Throws(() => _ = _pool.IsRunning); + Setup(SuccessfulConnectionFactory); + Assert.True(pool.IsRunning); } [Fact] public void TestLoadBalanceTimeout() { - Assert.Throws(() => _ = _pool.LoadBalanceTimeout); + Setup(SuccessfulConnectionFactory); + Assert.Equal(poolGroupOptions.LoadBalanceTimeout, pool.LoadBalanceTimeout); } [Fact] public void TestPoolGroup() { - Assert.Throws(() => _ = _pool.PoolGroup); + Setup(SuccessfulConnectionFactory); + Assert.Equal(dbConnectionPoolGroup, pool.PoolGroup); } [Fact] public void TestPoolGroupOptions() { - Assert.Throws(() => _ = _pool.PoolGroupOptions); + Setup(SuccessfulConnectionFactory); + Assert.Equal(poolGroupOptions, pool.PoolGroupOptions); } [Fact] public void TestProviderInfo() { - Assert.Throws(() => _ = _pool.ProviderInfo); + Setup(SuccessfulConnectionFactory); + Assert.Equal(connectionPoolProviderInfo, pool.ProviderInfo); } [Fact] public void TestStateGetter() { - Assert.Throws(() => _ = _pool.State); + Setup(SuccessfulConnectionFactory); + Assert.Equal(DbConnectionPoolState.Running, pool.State); } [Fact] public void TestStateSetter() { - Assert.Throws(() => _pool.State = DbConnectionPoolState.Running); + Setup(SuccessfulConnectionFactory); + Assert.Equal(DbConnectionPoolState.Running, pool.State); } [Fact] public void TestUseLoadBalancing() { - Assert.Throws(() => _ = _pool.UseLoadBalancing); + Setup(SuccessfulConnectionFactory); + Assert.Equal(poolGroupOptions.UseLoadBalancing, pool.UseLoadBalancing); } + #endregion + + #region Not Implemented Method Tests + [Fact] public void TestClear() { - Assert.Throws(() => _pool.Clear()); + Setup(SuccessfulConnectionFactory); + Assert.Throws(() => pool.Clear()); } [Fact] public void TestPutObjectFromTransactedPool() { - Assert.Throws(() => _pool.PutObjectFromTransactedPool(null!)); + Setup(SuccessfulConnectionFactory); + Assert.Throws(() => pool.PutObjectFromTransactedPool(null!)); } [Fact] public void TestReplaceConnection() { - Assert.Throws(() => _pool.ReplaceConnection(null!, null!, null!)); + Setup(SuccessfulConnectionFactory); + Assert.Throws(() => pool.ReplaceConnection(null!, null!, null!)); } [Fact] - public void TestReturnInternalConnection() + public void TestShutdown() { - Assert.Throws(() => _pool.ReturnInternalConnection(null!, null!)); + Setup(SuccessfulConnectionFactory); + Assert.Throws(() => pool.Shutdown()); } [Fact] - public void TestShutdown() + public void TestStartup() { - Assert.Throws(() => _pool.Shutdown()); + Setup(SuccessfulConnectionFactory); + Assert.Throws(() => pool.Startup()); } [Fact] - public void TestStartup() + public void TestTransactionEnded() + { + Setup(SuccessfulConnectionFactory); + Assert.Throws(() => pool.TransactionEnded(null!, null!)); + } + #endregion + + #region Test classes + internal class SuccessfulSqlConnectionFactory : SqlConnectionFactory + { + protected override DbConnectionInternal CreateConnection( + DbConnectionOptions options, + DbConnectionPoolKey poolKey, + DbConnectionPoolGroupProviderInfo poolGroupProviderInfo, + IDbConnectionPool pool, + DbConnection owningConnection, + DbConnectionOptions userOptions) + { + return new StubDbConnectionInternal(); + } + } + + internal class TimeoutSqlConnectionFactory : SqlConnectionFactory + { + protected override DbConnectionInternal CreateConnection( + DbConnectionOptions options, + DbConnectionPoolKey poolKey, + DbConnectionPoolGroupProviderInfo poolGroupProviderInfo, + IDbConnectionPool pool, + DbConnection owningConnection, + DbConnectionOptions userOptions) + { + throw ADP.PooledOpenTimeout(); + } + } + + internal class StubDbConnectionInternal : DbConnectionInternal { - Assert.Throws(() => _pool.Startup()); + #region Not Implemented Members + public override string ServerVersion => throw new NotImplementedException(); + + public override DbTransaction BeginTransaction(System.Data.IsolationLevel il) + { + throw new NotImplementedException(); + } + + public override void EnlistTransaction(Transaction transaction) + { + return; + } + + protected override void Activate(Transaction transaction) + { + return; + } + + protected override void Deactivate() + { + return; + } + #endregion } + #endregion [Fact] - public void TestTransactionEnded() + public void Constructor_WithZeroMaxPoolSize_ThrowsArgumentOutOfRangeException() + { + // Arrange + var poolGroupOptions = new DbConnectionPoolGroupOptions( + poolByIdentity: false, + minPoolSize: 0, + maxPoolSize: 0, // This should cause an exception + creationTimeout: 15, + loadBalanceTimeout: 0, + hasTransactionAffinity: true + ); + var dbConnectionPoolGroup = new DbConnectionPoolGroup( + new DbConnectionOptions("DataSource=localhost;", null), + new DbConnectionPoolKey("TestDataSource"), + poolGroupOptions + ); + + // Act & Assert + var exception = Assert.Throws(() => + new ChannelDbConnectionPool( + SuccessfulConnectionFactory, + dbConnectionPoolGroup, + DbConnectionPoolIdentity.NoIdentity, + new DbConnectionPoolProviderInfo() + )); + + Assert.Equal("fixedCapacity", exception.ParamName); + Assert.Contains("Capacity must be greater than zero", exception.Message); + } + + [Fact] + public void Constructor_WithLargeMaxPoolSize() { - Assert.Throws(() => _pool.TransactionEnded(null!, null!)); + // Arrange - Test that Int32.MaxValue is accepted as a valid pool size + var poolGroupOptions = new DbConnectionPoolGroupOptions( + poolByIdentity: false, + minPoolSize: 0, + maxPoolSize: 10000, + creationTimeout: 15, + loadBalanceTimeout: 0, + hasTransactionAffinity: true + ); + var dbConnectionPoolGroup = new DbConnectionPoolGroup( + new DbConnectionOptions("DataSource=localhost;", null), + new DbConnectionPoolKey("TestDataSource"), + poolGroupOptions + ); + + try + { + // Act & Assert - This should not throw ArgumentOutOfRangeException, but may throw OutOfMemoryException + var pool = new ChannelDbConnectionPool( + SuccessfulConnectionFactory, + dbConnectionPoolGroup, + DbConnectionPoolIdentity.NoIdentity, + new DbConnectionPoolProviderInfo() + ); + + Assert.NotNull(pool); + Assert.Equal(0, pool.Count); + } + catch (OutOfMemoryException) + { + // OutOfMemoryException is acceptable when trying to allocate an array of int.MaxValue size + // This test is primarily checking that ArgumentOutOfRangeException is not thrown for the capacity validation + // The fact that we reach the OutOfMemoryException means the capacity validation passed + } } [Fact] - public void TestTryGetConnection() + public void Constructor_WithValidSmallPoolSizes_WorksCorrectly() { - Assert.Throws(() => _pool.TryGetConnection(null!, null!, null!, out _)); + // Arrange - Test various small pool sizes that should work correctly + + // Test with pool size of 1 + var poolGroupOptions1 = new DbConnectionPoolGroupOptions( + poolByIdentity: false, + minPoolSize: 0, + maxPoolSize: 1, + creationTimeout: 15, + loadBalanceTimeout: 0, + hasTransactionAffinity: true + ); + var dbConnectionPoolGroup1 = new DbConnectionPoolGroup( + new DbConnectionOptions("DataSource=localhost;", null), + new DbConnectionPoolKey("TestDataSource"), + poolGroupOptions1 + ); + + // Act & Assert - Pool size of 1 should work + var pool1 = new ChannelDbConnectionPool( + SuccessfulConnectionFactory, + dbConnectionPoolGroup1, + DbConnectionPoolIdentity.NoIdentity, + new DbConnectionPoolProviderInfo() + ); + + Assert.NotNull(pool1); + Assert.Equal(0, pool1.Count); + + // Test with pool size of 2 + var poolGroupOptions2 = new DbConnectionPoolGroupOptions( + poolByIdentity: false, + minPoolSize: 0, + maxPoolSize: 2, + creationTimeout: 15, + loadBalanceTimeout: 0, + hasTransactionAffinity: true + ); + var dbConnectionPoolGroup2 = new DbConnectionPoolGroup( + new DbConnectionOptions("DataSource=localhost;", null), + new DbConnectionPoolKey("TestDataSource"), + poolGroupOptions2 + ); + + var pool2 = new ChannelDbConnectionPool( + SuccessfulConnectionFactory, + dbConnectionPoolGroup2, + DbConnectionPoolIdentity.NoIdentity, + new DbConnectionPoolProviderInfo() + ); + + Assert.NotNull(pool2); + Assert.Equal(0, pool2.Count); } } } diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ConnectionPool/ConnectionPoolSlotsTest.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ConnectionPool/ConnectionPoolSlotsTest.cs new file mode 100644 index 0000000000..698391ef8f --- /dev/null +++ b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ConnectionPool/ConnectionPoolSlotsTest.cs @@ -0,0 +1,550 @@ +// 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.Threading.Tasks; +using Microsoft.Data.SqlClient.ConnectionPool; +using Microsoft.Data.ProviderBase; +using Xunit; +using System.Data; +using System.Data.Common; +using System.Transactions; + +#nullable enable + +namespace Microsoft.Data.SqlClient.UnitTests.ConnectionPool +{ + public class ConnectionPoolSlotsTest + { + // Mock implementation of DbConnectionInternal for testing + private class MockDbConnectionInternal : DbConnectionInternal + { + public MockDbConnectionInternal() : base(ConnectionState.Open, true, false) { } + + public override string ServerVersion => "Mock Server 1.0"; + + public override DbTransaction BeginTransaction(System.Data.IsolationLevel il) + { + throw new NotImplementedException(); + } + + public override void EnlistTransaction(Transaction transaction) + { + // Mock implementation - do nothing + } + + protected override void Activate(Transaction transaction) + { + // Mock implementation - do nothing + } + + protected override void Deactivate() + { + // Mock implementation - do nothing + } + } + + [Fact] + public void Constructor_ValidCapacity_SetsReservationCountToZero() + { + // Arrange & Act + var poolSlots = new ConnectionPoolSlots(5); + + // Assert + Assert.Equal(0, poolSlots.ReservationCount); + } + + [Fact] + public void Constructor_ZeroCapacity_ThrowsArgumentOutOfRangeException() + { + // Act & Assert + var exception = Assert.Throws(() => new ConnectionPoolSlots(0)); + Assert.Equal("fixedCapacity", exception.ParamName); + Assert.Contains("Capacity must be greater than zero", exception.Message); + } + + [Fact] + public void Constructor_CapacityGreaterThanIntMaxValue_ThrowsArgumentOutOfRangeException() + { + // Arrange + uint invalidCapacity = (uint)int.MaxValue + 1; + + // Act & Assert + var exception = Assert.Throws(() => new ConnectionPoolSlots(invalidCapacity)); + Assert.Equal("fixedCapacity", exception.ParamName); + Assert.Contains("Capacity must be less than or equal to Int32.MaxValue", exception.Message); + } + + [Theory] + [InlineData(10000u)] + public void Constructor_LargeCapacityValues_SetsReservationCountToZero(uint capacity) + { + // Arrange & Act + var poolSlots = new ConnectionPoolSlots(capacity); + + // Assert + Assert.Equal(0, poolSlots.ReservationCount); + } + + [Fact] + public void Add_ValidConnection_ReturnsConnectionAndIncrementsReservationCount() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(5); + var createCallbackCount = 0; + + // Act + var connection = poolSlots.Add( + createCallback: (state) => { + createCallbackCount++; + return new MockDbConnectionInternal(); + }, + cleanupCallback: (conn, state) => Assert.Fail(), + createState: "test", + cleanupState: "cleanup"); + + // Assert + Assert.NotNull(connection); + Assert.Equal(1, poolSlots.ReservationCount); + Assert.Equal(1, createCallbackCount); + } + + [Fact] + public void Add_NullFromCreateCallback_ReturnsNullAndDoesNotIncrementReservationCount() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(5); + var createCallbackCount = 0; + + // Act + var connection = poolSlots.Add( + createCallback: state => { + createCallbackCount++; + return null; + }, + cleanupCallback: (conn, state) => Assert.Fail(), + createState: "test", + cleanupState: "cleanup"); + + // Assert + Assert.Null(connection); + Assert.Equal(0, poolSlots.ReservationCount); + Assert.Equal(1, createCallbackCount); + } + + [Fact] + public void Add_AtCapacity_ReturnsNullForAdditionalConnections() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(1); + var createCallbackCount = 0; + + // Act - Add first connection + var connection1 = poolSlots.Add( + createCallback: state => { + createCallbackCount++; + return new MockDbConnectionInternal(); + }, + cleanupCallback: (conn, state) => Assert.Fail(), + createState: "test", + cleanupState: "cleanup"); + + // Act - Try to add second connection beyond capacity + var connection2 = poolSlots.Add( + createCallback: state => + { + Assert.Fail(); + return null; + }, + cleanupCallback: (conn, state) => { + Assert.Fail(); + }, + createState: "test", + cleanupState: "cleanup"); + + // Assert + Assert.NotNull(connection1); + Assert.Null(connection2); + Assert.Equal(1, poolSlots.ReservationCount); + Assert.Equal(1, createCallbackCount); + } + + [Fact] + public void Add_CreateCallbackThrowsException_CallsCleanupCallbackAndRethrowsException() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(5); + var createCallbackCount = 0; + var cleanupCallbackCount = 0; + + // Act & Assert + var exception = Assert.Throws(() => + poolSlots.Add( + createCallback: state => { + createCallbackCount++; + throw new InvalidOperationException("Test exception"); + }, + cleanupCallback: (conn, state) => cleanupCallbackCount++, + createState: "test", + cleanupState: "cleanup")); + + Assert.Contains("Failed to create or add connection", exception.Message); + Assert.Equal(1, cleanupCallbackCount); + Assert.Equal(0, poolSlots.ReservationCount); + Assert.Equal(1, createCallbackCount); + Assert.Equal(1, cleanupCallbackCount); + } + + [Fact] + public void Add_MultipleConnections_IncrementsReservationCountCorrectly() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(5); + var createCallbackCount = 0; + var createCallbackCount2 = 0; + + // Act + var connection1 = poolSlots.Add( + createCallback: state => + { + createCallbackCount++; + return new MockDbConnectionInternal(); + }, + cleanupCallback: (conn, state) => Assert.Fail(), + createState: "test", + cleanupState: "cleanup"); + + var connection2 = poolSlots.Add( + createCallback: state => + { + createCallbackCount2++; + return new MockDbConnectionInternal(); + }, + cleanupCallback: (conn, state) => Assert.Fail(), + createState: "test", + cleanupState: "cleanup"); + + // Assert + Assert.NotNull(connection1); + Assert.NotNull(connection2); + Assert.Equal(2, poolSlots.ReservationCount); + Assert.Equal(1, createCallbackCount); + Assert.Equal(1, createCallbackCount2); + } + + [Fact] + public void TryRemove_ExistingConnection_ReturnsTrueAndDecrementsReservationCount() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(5); + var connection = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "test", + cleanupState: "cleanup"); + + var reservationCountBeforeRemove = poolSlots.ReservationCount; + + // Act + var removed = poolSlots.TryRemove(connection!); + + // Assert + Assert.Equal(1, reservationCountBeforeRemove); + Assert.True(removed); + Assert.Equal(0, poolSlots.ReservationCount); + } + + [Fact] + public void TryRemove_NonExistentConnection_ReturnsFalseAndDoesNotChangeReservationCount() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(5); + var connection = new MockDbConnectionInternal(); + var connection2 = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "test", + cleanupState: "cleanup"); + var reservationCountBeforeRemove = poolSlots.ReservationCount; + + // Act + var removed = poolSlots.TryRemove(connection); + + // Assert + Assert.Equal(1, reservationCountBeforeRemove); + Assert.False(removed); + Assert.Equal(1, poolSlots.ReservationCount); + } + + [Fact] + public void TryRemove_SameConnectionTwice_ReturnsFalseOnSecondAttempt() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(5); + var connection = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "test", + cleanupState: "cleanup"); + var reservationCountBeforeRemove = poolSlots.ReservationCount; + + // Act + var firstRemove = poolSlots.TryRemove(connection!); + var secondRemove = poolSlots.TryRemove(connection!); + + // Assert + Assert.Equal(1, reservationCountBeforeRemove); + Assert.True(firstRemove); + Assert.False(secondRemove); + Assert.Equal(0, poolSlots.ReservationCount); + } + + [Fact] + public void TryRemove_SameConnectionTwice_ReturnsTrueWhenAddedTwice() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(5); + var commonConnection = new MockDbConnectionInternal(); + var connection = poolSlots.Add( + createCallback: state => commonConnection, + cleanupCallback: (conn, state) => { }, + createState: "test", + cleanupState: "cleanup"); + var connection2 = poolSlots.Add( + createCallback: state => commonConnection, + cleanupCallback: (conn, state) => { }, + createState: "test", + cleanupState: "cleanup"); + var reservationCountBeforeRemove = poolSlots.ReservationCount; + + // Act + var firstRemove = poolSlots.TryRemove(connection!); + var secondRemove = poolSlots.TryRemove(connection2!); + + // Assert + Assert.Equal(2, reservationCountBeforeRemove); + Assert.True(firstRemove); + Assert.True(secondRemove); + Assert.Equal(0, poolSlots.ReservationCount); + } + + [Fact] + public void TryRemove_MultipleConnections_RemovesCorrectConnection() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(5); + var connection1 = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "test1", + cleanupState: "cleanup1"); + + var connection2 = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "test2", + cleanupState: "cleanup2"); + + // Act + var removed = poolSlots.TryRemove(connection1!); + + // Assert + Assert.True(removed); + Assert.Equal(1, poolSlots.ReservationCount); + + // Act + var removed2 = poolSlots.TryRemove(connection1!); + + // Assert + Assert.False(removed2); // Should return false since connection1 was already removed + Assert.Equal(1, poolSlots.ReservationCount); + + // Act + var removed3 = poolSlots.TryRemove(connection2!); + + // Assert + Assert.True(removed3); + Assert.Equal(0, poolSlots.ReservationCount); + } + + [Fact] + public void ConcurrentAddAndRemove_MaintainsCorrectReservationCount() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(100); + const int operationCount = 50; + var connections = new DbConnectionInternal[operationCount]; + var addTasks = new Task[operationCount]; + + // Act - Add connections concurrently + for (int i = 0; i < operationCount; i++) + { + int index = i; + addTasks[i] = Task.Run(() => + { + connections[index] = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: $"test{index}", + cleanupState: $"cleanup{index}")!; + }); + } + + // Wait for all add operations to complete + Task.WaitAll(addTasks); + + // Verify all connections were added + Assert.Equal(operationCount, poolSlots.ReservationCount); + + var removeTasks = new Task[operationCount]; + + // Act - Remove connections concurrently + for (int i = 0; i < operationCount; i++) + { + int index = i; + removeTasks[i] = Task.Run(() => + { + poolSlots.TryRemove(connections[index]); + }); + } + + // Wait for all remove operations to complete + Task.WaitAll(removeTasks); + + // Assert + Assert.Equal(0, poolSlots.ReservationCount); + } + + [Theory] + [InlineData(1u)] + [InlineData(5u)] + [InlineData(10u)] + [InlineData(100u)] + public void Add_FillToCapacity_RespectsCapacityLimits(uint capacity) + { + // Arrange + var poolSlots = new ConnectionPoolSlots(capacity); + var connections = new DbConnectionInternal[capacity]; + + // Act - Fill to capacity + for (int i = 0; i < capacity; i++) + { + connections[i] = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: $"test{i}", + cleanupState: $"cleanup{i}")!; + Assert.NotNull(connections[i]); + } + + // Try to add one more beyond capacity + var extraConnection = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "overflow", + cleanupState: "overflow"); + + // Assert + Assert.Equal((int)capacity, poolSlots.ReservationCount); + Assert.Null(extraConnection); // The overflow connection should be null + } + + [Fact] + public void ReservationCount_AfterAddAndRemoveOperations_ReflectsCurrentState() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(10); + + // Act & Assert - Start with 0 + Assert.Equal(0, poolSlots.ReservationCount); + + // Add 3 connections + var conn1 = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "test1", + cleanupState: "cleanup1"); + Assert.Equal(1, poolSlots.ReservationCount); + + var conn2 = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "test2", + cleanupState: "cleanup2"); + Assert.Equal(2, poolSlots.ReservationCount); + + var conn3 = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "test3", + cleanupState: "cleanup3"); + Assert.Equal(3, poolSlots.ReservationCount); + + // Remove 1 connection + poolSlots.TryRemove(conn2!); + Assert.Equal(2, poolSlots.ReservationCount); + + // Remove remaining connections + poolSlots.TryRemove(conn1!); + Assert.Equal(1, poolSlots.ReservationCount); + + poolSlots.TryRemove(conn3!); + Assert.Equal(0, poolSlots.ReservationCount); + } + + [Fact] + public void Add_StateParametersPassedCorrectly_UsesProvidedState() + { + // Arrange + var poolSlots = new ConnectionPoolSlots(5); + string? receivedCreateState = null; + string? receivedCleanupState = null; + + // Act + try + { + poolSlots.Add( + createCallback: state => { receivedCreateState = state; throw new InvalidOperationException("Test"); }, + cleanupCallback: (conn, state) => { receivedCleanupState = state; }, + createState: "createState", + cleanupState: "cleanupState"); + } + catch + { + // Expected due to exception in create callback + } + + // Assert + Assert.Equal("createState", receivedCreateState); + Assert.Equal("cleanupState", receivedCleanupState); + } + + [Fact] + public void Constructor_EdgeCase_CapacityOfOne_WorksCorrectly() + { + // Arrange & Act + var poolSlots = new ConnectionPoolSlots(1); + + // Assert - Should be able to add one connection + var connection = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "test", + cleanupState: "cleanup"); + + Assert.NotNull(connection); + Assert.Equal(1, poolSlots.ReservationCount); + + // Should not be able to add a second connection + var connection2 = poolSlots.Add( + createCallback: state => new MockDbConnectionInternal(), + cleanupCallback: (conn, state) => { }, + createState: "test2", + cleanupState: "cleanup2"); + + Assert.Null(connection2); + Assert.Equal(1, poolSlots.ReservationCount); + } + } +} diff --git a/tools/specs/Microsoft.Data.SqlClient.nuspec b/tools/specs/Microsoft.Data.SqlClient.nuspec index 0f029b4c19..ad04728342 100644 --- a/tools/specs/Microsoft.Data.SqlClient.nuspec +++ b/tools/specs/Microsoft.Data.SqlClient.nuspec @@ -40,6 +40,7 @@ + @@ -76,6 +77,7 @@ +