[7.1.1 Cherry-pick] Fix access token expiry eviction in connection pool V2

doc/snippets/Microsoft.Data.SqlClient/SqlConnection.xml

@@ -361,9 +361,14 @@ The following example creates a <xref:Microsoft.Data.SqlClient.SqlCommand> and a
           are the same, they will be grouped into the same connection pool.
         </para>
         <para>
-          When using a token callback function, the connection manages
-          refreshing the tokens returned by the callback.  The application is
-          not responsible for knowing when tokens expire.
+          The driver manages token refresh and discards pooled connections with
+          expired or nearly expired tokens before reuse. Connections in use or
+          reused within the same active transaction are unaffected.
+        </para>
+        <para>
+          New physical connections invoke the callback only when a cached token
+          is missing or needs refreshing. Reusing a pooled connection does not
+          invoke the callback.
         </para>
         <para>
           This property is mutually exclusive with the

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

@@ -704,7 +704,9 @@ private void PutConnectionInIdleChannel(DbConnectionInternal connection, bool pr
                 connection.SetReturnedTime(_timeProvider.GetUtcNow().UtcDateTime);
             }
 
-            if (!IsLiveConnection(connection, probeLiveness))
+            // Match V1: check token expiry on general checkout, not return. An expired connection
+            // may remain idle, but will be discarded before it can be reused outside its transaction.
+            if (!IsLiveConnection(connection, probeLiveness, checkAccessTokenExpiry: false))
             {
                 RemoveConnection(connection);
                 return;
@@ -1264,9 +1266,21 @@ _connectionCreationRateLimiter is not null &&
         /// Whether to poll the physical connection to confirm it is still alive. Pass false when
         /// running on a thread that must not block; the remaining checks are all cheap and local.
         /// </param>
+        /// <param name="checkAccessTokenExpiry">
+        /// Validate the token before general checkout, but not when returning a connection to the pool.
+        /// </param>
         /// <returns>Returns true if the connection is live and unexpired, otherwise returns false.</returns>
-        private bool IsLiveConnection(DbConnectionInternal connection, bool probeLiveness = true)
+        private bool IsLiveConnection(DbConnectionInternal connection, bool probeLiveness = true, bool checkAccessTokenExpiry = true)
         {
+            if (checkAccessTokenExpiry && connection.IsAccessTokenExpired)
+            {
+                SqlClientEventSource.Log.TryPoolerTraceEvent(
+                    "ChannelDbConnectionPool.IsLiveConnection | INFO | {0}, Connection {1}, will not be reused because its access token has expired or is about to expire.",
+                    Id,
+                    connection.ObjectID);
+                return false;
+            }
+
             // Connection has been sitting idle longer than the configured idle timeout.
             // Checked before the (potentially expensive) liveness probe so an idle-expired
             // connection is discarded without an SNI round-trip.
@@ -1548,7 +1562,7 @@ private async Task<DbConnectionInternal> GetInternalConnection(
                         {
                             // Skip the liveness/idle/generation gate at the bottom of the loop:
                             // GetFromTransactedPool has already probed liveness, and a transacted
-                            // connection is exempt from idle-timeout, load-balance and
+                            // connection is exempt from token-expiry, idle-timeout, load-balance and
                             // clear-generation eviction because closing it would abort its
                             // (possibly distributed) transaction.
                             break;

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/AADFedAuthTokenRefreshTest/AADFedAuthTokenRefreshTest.cs

@@ -4,7 +4,13 @@
 
 using System;
 using System.Diagnostics;
+using System.Reflection;
+using System.Threading;
+using System.Threading.Tasks;
+using Azure.Core;
 using Microsoft.Data.SqlClient.ManualTesting.Tests.SQL.Common.SystemDataInternals;
+using Microsoft.Data.SqlClient.ManualTesting.Tests.SystemDataInternals;
+using Microsoft.Data.SqlClient.Tests.Common;
 using Xunit;
 using Xunit.Abstractions;
 
@@ -87,6 +93,107 @@ public void FedAuthTokenRefreshTest()
             }
         }
 
+        /// <summary>
+        /// Verifies both pools replace connections with expired or nearly expired tokens and
+        /// invoke the callback when the cached token also needs refreshing, for Open and OpenAsync.
+        /// </summary>
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.IsAADPasswordConnStrSetup))]
+        [InlineData(false, false, -1)]
+        [InlineData(false, true, -1)]
+        [InlineData(true, false, -1)]
+        [InlineData(true, true, -1)]
+        [InlineData(false, false, 5)]
+        [InlineData(false, true, 5)]
+        [InlineData(true, false, 5)]
+        [InlineData(true, true, 5)]
+        public async Task AccessTokenCallback_PooledConnectionIsReplacedOnExpiry(bool usePoolV2, bool async, int expiresInSeconds)
+        {
+            using var poolVersion = new ConnectionPoolVersionScope(usePoolV2);
+            string[] credentialKeys = { "Authentication", "User ID", "Password", "UID", "PWD" };
+            var builder = new SqlConnectionStringBuilder(
+                DataTestUtility.RemoveKeysInConnStr(DataTestUtility.AADPasswordConnectionString, credentialKeys))
+            {
+                Pooling = true,
+                MinPoolSize = 0,
+                MaxPoolSize = 1,
+                ConnectTimeout = 30,
+                Enlist = false
+            };
+            var credential = DataTestUtility.GetTokenCredential();
+            SqlAuthenticationToken callbackToken = null;
+            int callbackInvocations = 0;
+            using var connection = new SqlConnection(builder.ConnectionString)
+            {
+                AccessTokenCallback = async (parameters, cancellationToken) =>
+                {
+                    Interlocked.Increment(ref callbackInvocations);
+                    const string suffix = "/.default";
+                    string scope = parameters.Resource.EndsWith(suffix) ? parameters.Resource : parameters.Resource + suffix;
+                    AccessToken token = await credential.GetTokenAsync(new TokenRequestContext(new[] { scope }), cancellationToken);
+                    callbackToken = new SqlAuthenticationToken(token.Token, token.ExpiresOn);
+                    return callbackToken;
+                }
+            };
+
+            Task OpenConnection()
+            {
+                if (async)
+                {
+                    return connection.OpenAsync();
+                }
+                connection.Open();
+                return Task.CompletedTask;
+            }
+
+            // The empty pool requires a physical login, which invokes the callback and caches its token.
+            await OpenConnection();
+            object original = connection.GetInternalConnection();
+            Assert.NotNull(callbackToken);
+            int callbackCountAfterLogin = callbackInvocations;
+            Assert.True(callbackCountAfterLogin > 0);
+            // Close returns the physical connection to the pool; reopening reuses it without authentication.
+            connection.Close();
+            await OpenConnection();
+            Assert.Same(original, connection.GetInternalConnection());
+            Assert.Equal(callbackCountAfterLogin, callbackInvocations);
+
+            // The connection's expiry controls eviction; the cached token's expiry controls callback refresh.
+            // Age both without changing the real token or waiting. Five seconds is within the 30-second
+            // checkout buffer; minus one second covers an already expired token.
+            DateTimeOffset expiry = DateTimeOffset.UtcNow.AddSeconds(expiresInSeconds);
+            object cachedContext = FedAuthTokenHelper.GetAuthenticationContextValue(connection);
+            FieldInfo cacheExpiryField = cachedContext.GetType().GetField("_expirationTime", BindingFlags.Instance | BindingFlags.NonPublic);
+            Assert.NotNull(cacheExpiryField);
+            cacheExpiryField.SetValue(cachedContext, expiry.UtcDateTime);
+            FieldInfo tokenField = original.GetType().GetField("_fedAuthToken", BindingFlags.Instance | BindingFlags.NonPublic);
+            Assert.NotNull(tokenField);
+            object expiringToken = Activator.CreateInstance(tokenField.FieldType,
+                BindingFlags.Instance | BindingFlags.NonPublic, binder: null,
+                args: new object[] { new SqlAuthenticationToken(callbackToken.AccessToken, expiry) },
+                culture: null);
+            tokenField.SetValue(original, expiringToken);
+            // Expiry is checked on checkout, not return, so Close does not invoke the callback.
+            connection.Close();
+            Assert.Equal(callbackCountAfterLogin, callbackInvocations);
+
+            // Checkout discards the expired physical connection rather than reauthenticating it.
+            // Its replacement needs a login, and the expiring cache entry forces another callback.
+            // The credential may still return the same valid token; callback invocation is what matters.
+            await OpenConnection();
+            Assert.NotSame(original, connection.GetInternalConnection());
+            Assert.True(callbackInvocations > callbackCountAfterLogin);
+            object replacement = connection.GetInternalConnection();
+            int callbackCountAfterRefresh = callbackInvocations;
+            // The replacement now has a valid token, so another reopen reuses it without another callback.
+            connection.Close();
+            await OpenConnection();
+            Assert.Same(replacement, connection.GetInternalConnection());
+            Assert.Equal(callbackCountAfterRefresh, callbackInvocations);
+            using SqlCommand command = connection.CreateCommand();
+            command.CommandText = "SELECT 1";
+            Assert.Equal(1, async ? await command.ExecuteScalarAsync() : command.ExecuteScalar());
+        }
+
         [Conditional("DEBUG")]
         private void LogInfo(string message)
         {