Fix | Fixes concurrent connection speed issues due to 'Authentication Context' cache not maintained in NetCore

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

@@ -2380,6 +2380,19 @@ internal void OnFeatureExtAck(int featureId, byte[] data)
                                 throw SQL.ParsingErrorLibraryType(ParsingErrorState.FedAuthFeatureAckUnknownLibraryType, (int)_fedAuthFeatureExtensionData.Value.libraryType);
                         }
                         _federatedAuthenticationAcknowledged = true;
+
+                        // If a new authentication context was used as part of this login attempt, try to update the new context in the cache, i.e.dbConnectionPool.AuthenticationContexts.
+                        // ChooseAuthenticationContextToUpdate will take care that only the context which has more validity will remain in the cache, based on the Update logic.
+                        if (_newDbConnectionPoolAuthenticationContext != null)
+                        {
+                            Debug.Assert(_dbConnectionPool != null, "_dbConnectionPool should not be null when _newDbConnectionPoolAuthenticationContext != null.");
+
+                            DbConnectionPoolAuthenticationContext newAuthenticationContextInCacheAfterAddOrUpdate = _dbConnectionPool.AuthenticationContexts.AddOrUpdate(_dbConnectionPoolAuthenticationContextKey, _newDbConnectionPoolAuthenticationContext,
+                                                                                 (key, oldValue) => DbConnectionPoolAuthenticationContext.ChooseAuthenticationContextToUpdate(oldValue, _newDbConnectionPoolAuthenticationContext));
+
+                            Debug.Assert(newAuthenticationContextInCacheAfterAddOrUpdate != null, "newAuthenticationContextInCacheAfterAddOrUpdate should not be null.");
+                        }
+
                         break;
                     }
                 case TdsEnums.FEATUREEXT_TCE:

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/ConnectivityTests/AADConnectionTest.cs

@@ -2,8 +2,9 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using Xunit;
 using System;
+using System.Diagnostics;
+using Xunit;
 
 namespace Microsoft.Data.SqlClient.ManualTesting.Tests
 {
@@ -300,5 +301,34 @@ public static void NoCredentialsActiveDirectoryPassword()
             string expectedMessage = "Either Credential or both 'User ID' and 'Password' (or 'UID' and 'PWD') connection string keywords must be specified, if 'Authentication=Active Directory Password'.";
             Assert.Contains(expectedMessage, e.Message);
         }
+
+        [ConditionalFact(nameof(IsAADConnStringsSetup))]
+        public static void ConnectionSpeed()
+        {
+            //Ensure server endpoints are warm
+            using (var connectionDrill = new SqlConnection(DataTestUtility.AADPasswordConnectionString))
+            {
+                connectionDrill.Open();
+            }
+            SqlConnection.ClearAllPools();
+
+            using (var connectionDrill = new SqlConnection(DataTestUtility.AADPasswordConnectionString))
+            {
+                Stopwatch firstConnectionTime = new Stopwatch();
+                firstConnectionTime.Start();
+                connectionDrill.Open();
+                firstConnectionTime.Stop();
+                using (var connectionDrill2 = new SqlConnection(DataTestUtility.AADPasswordConnectionString))
+                {
+                    Stopwatch secondConnectionTime = new Stopwatch();
+                    secondConnectionTime.Start();
+                    connectionDrill2.Open();
+                    secondConnectionTime.Stop();
+                    // Subsequent AAD connections within a short timeframe should use an auth token cached from the connection pool
+                    // Second connection speed in tests was typically 10-15% of the first connection time. Using 30% since speeds may vary.
+                    Assert.True(secondConnectionTime.ElapsedMilliseconds / firstConnectionTime.ElapsedMilliseconds < .30, $"Second AAD connection too slow ({secondConnectionTime.ElapsedMilliseconds}ms)! (More than 30% of the first ({firstConnectionTime.ElapsedMilliseconds}ms).)");
+                }
+            }
+        }
     }
 }