Merge branch 'dev/mdaigle/disposable-reservations' into dev/mdaigle/get-return-pooled-connections

Author: mdaigle

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

@@ -133,7 +133,20 @@ internal DbConnectionInternal CreatePooledConnection(IDbConnectionPool pool, DbC
 
                 newConnection.MakePooledConnection(pool);
             }
+
+            if (newConnection == null)
+            {
+                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.Log.TryTraceEvent("<prov.DbConnectionFactory.CreatePooledConnection|RES|CPOOL> {0}, Pooled database connection created.", ObjectID);
+
+            newConnection.PrePush(null);
+
             return newConnection;
         }
 

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPool.cs

@@ -300,22 +300,21 @@ public bool TryGetConnection(
         /// <param name="cancellationToken">The cancellation token to cancel the operation.</param>
         /// <returns>A task representing the asynchronous operation, with a result of the new internal connection.</returns>
         /// <throws>InvalidOperationException - when the newly created connection is invalid or already in the pool.</throws>
-        private Task<DbConnectionInternal?> OpenNewInternalConnection(
+        private DbConnectionInternal? OpenNewInternalConnection(
             DbConnection? owningConnection, 
             DbConnectionOptions userOptions, 
-            bool async, 
+            bool async,
             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. 
-            if (_connectionSlots.TryReserve())
-            {
-                DbConnectionInternal? newConnection = null;
-                try
-                {
-                    var startTime = Stopwatch.GetTimestamp();
+            // in case of an exception.
 
+            return _connectionSlots.Add(
+                createCallback: () =>
+                {
                     // 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 
@@ -324,44 +323,20 @@ public bool TryGetConnection(
                     // 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.
-                    newConnection = ConnectionFactory.CreatePooledConnection(
+                    return ConnectionFactory.CreatePooledConnection(
                         this,
                         owningConnection,
                         PoolGroup.ConnectionOptions,
                         PoolGroup.PoolKey,
                         userOptions);
-
-                    // We don't expect these conditions to happen, but we need to check them to be thorough.
-                    if (newConnection == null)
-                    {
-                        throw ADP.InternalError(ADP.InternalErrorCode.CreateObjectReturnedNull);
-                    }
-                    if (!newConnection.CanBePooled)
-                    {
-                        throw ADP.InternalError(ADP.InternalErrorCode.NewObjectCannotBePooled);
-                    }
-
-                    ValidateOwnershipAndSetPoolingState(newConnection, null);
-
-                    _connectionSlots.Add(newConnection);
-
-                    return Task.FromResult<DbConnectionInternal?>(newConnection);
-                }
-                catch
+                },
+                cleanupCallback: (newConnection) =>
                 {
-                    // If we failed to open a new connection, we need to reset any state we modified.
-                    // Clear the reservation we made, dispose of the connection if it was created, and wake up any waiting
-                    // attempts on the idle connection channel.
                     if (newConnection is not null)
                     {
                         RemoveConnection(newConnection);
                     }
-
-                    throw;
-                }
-            }
-
-            return Task.FromResult<DbConnectionInternal?>(null);
+                });
         }
 
         /// <summary>
@@ -392,7 +367,6 @@ private void RemoveConnection(DbConnectionInternal connection)
         {
             _connectionSlots.TryRemove(connection);
 
-            _connectionSlots.ReleaseReservation();
             // Closing a connection opens a free spot in the pool.
             // 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.
@@ -456,11 +430,11 @@ private async Task<DbConnectionInternal> GetInternalConnection(
 
 
                     // If we didn't find an idle connection, try to open a new one.  
-                    connection ??= await OpenNewInternalConnection(
+                    connection ??= OpenNewInternalConnection(
                         owningConnection,
                         userOptions,
                         async,
-                        cancellationToken).ConfigureAwait(false);
+                        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

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/ConnectionPoolSlots.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Diagnostics;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Data.ProviderBase;
 
 #nullable enable
@@ -43,6 +44,41 @@ namespace Microsoft.Data.SqlClient.ConnectionPool
     /// </example>
     internal sealed class ConnectionPoolSlots
     {
+
+        private sealed class Reservation : IDisposable
+        {
+            private readonly ConnectionPoolSlots _slots;
+            private bool _retain = false;
+            private bool _disposed = false;
+
+            internal Reservation(ConnectionPoolSlots slots)
+            {
+                _slots = slots;
+            }
+
+            public void Dispose()
+            {
+                if (_disposed)
+                {
+                    return;
+                }
+
+                if (!_retain)
+                {
+                    _slots.ReleaseReservation();
+                    _disposed = true;
+                }
+            }
+
+            internal void Keep()
+            {
+                _retain = true;
+            }
+
+        }
+
+        internal delegate void CleanupCallback(DbConnectionInternal? connection);
+
         private readonly DbConnectionInternal?[] _connections;
         private readonly uint _capacity;
         private volatile int _reservations;
@@ -66,30 +102,53 @@ internal ConnectionPoolSlots(uint fixedCapacity)
         /// <summary>
         /// Adds a connection to the collection. Can only be called after a reservation has been made.
         /// </summary>
-        /// <param name="connection">The connection to add to the collection.</param>
+        /// <param name="createCallback">The connection to add to the collection.</param>
+        /// <param name="cleanupCallback">Callback to clean up resources if an exception occurs.</param>
         /// <exception cref="InvalidOperationException">
         /// Thrown when unable to find an empty slot. 
         /// This can occur if a reservation is not taken before adding a connection.
         /// </exception>
-        internal void Add(DbConnectionInternal connection)
+        internal DbConnectionInternal? Add(Func<DbConnectionInternal?> createCallback, CleanupCallback cleanupCallback)
         {
-            int i;
-            for (i = 0; i < _capacity; i++)
+            DbConnectionInternal? connection = null;
+            try
             {
-                if (Interlocked.CompareExchange(ref _connections[i], connection, null) == null)
+                using var reservation = TryReserve();
+                if (reservation is null)
                 {
-                    return;
+                    return null;
                 }
-            }
 
-            throw new InvalidOperationException("Couldn't find an empty slot.");
+                connection = createCallback();
+
+                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);
+                throw new Exception("Failed to create or add connection", e);
+            }
         }
 
         /// <summary>
         /// Releases a reservation that was previously obtained.
         /// Must be called after removing a connection from the collection or if an exception occurs.
         /// </summary>
-        internal void ReleaseReservation()
+        private void ReleaseReservation()
         {
             Interlocked.Decrement(ref _reservations);
             Debug.Assert(_reservations >= 0, "Released a reservation that wasn't held");
@@ -106,6 +165,7 @@ internal bool TryRemove(DbConnectionInternal connection)
             {
                 if (Interlocked.CompareExchange(ref _connections[i], null, connection) == connection)
                 {
+                    ReleaseReservation();
                     return true;
                 }
             }
@@ -117,7 +177,7 @@ internal bool TryRemove(DbConnectionInternal connection)
         /// Attempts to reserve a spot in the collection.
         /// </summary>
         /// <returns>True if a reservation was successfully obtained.</returns>
-        internal bool TryReserve()
+        private Reservation? TryReserve()
         {
             for (var expected = _reservations; expected < _capacity; expected = _reservations)
             {
@@ -131,9 +191,9 @@ internal bool TryReserve()
                     continue;
                 }
 
-                return true;
+                return new Reservation(this);
             }
-            return false;
+            return null;
         }
     }
 }

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/ConnectionPool/WaitHandleDbConnectionPool.cs

@@ -745,15 +745,6 @@ private DbConnectionInternal CreateObject(DbConnection owningObject, DbConnectio
             try
             {
                 newObj = _connectionFactory.CreatePooledConnection(this, owningObject, _connectionPoolGroup.ConnectionOptions, _connectionPoolGroup.PoolKey, 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)
                 {

src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/ConnectionPool/ChannelDbConnectionPoolTest.cs

@@ -477,7 +477,7 @@ public void GetConnectionTimeout_ShouldThrowTimeoutException()
             DbConnectionInternal internalConnection = null;
 
             // Act & Assert
-            var ex = Assert.Throws<InvalidOperationException>(() =>
+            var ex = Assert.Throws<Exception>(() =>
             {
                 var completed = pool.TryGetConnection(
                     new SqlConnection(),
@@ -487,7 +487,9 @@ out internalConnection
                 );
             });
 
-            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.IsType<InvalidOperationException>(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]
@@ -499,7 +501,7 @@ public async Task GetConnectionAsyncTimeout_ShouldThrowTimeoutException()
             TaskCompletionSource<DbConnectionInternal> taskCompletionSource = new TaskCompletionSource<DbConnectionInternal>();
 
             // Act & Assert
-            var ex = await Assert.ThrowsAsync<InvalidOperationException>(async () =>
+            var ex = await Assert.ThrowsAsync<Exception>(async () =>
             {
                 var completed = pool.TryGetConnection(
                     new SqlConnection(),
@@ -511,7 +513,9 @@ out internalConnection
                 await taskCompletionSource.Task;
             });
 
-            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.IsType<InvalidOperationException>(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]