Address review follow-ups on the async enclave attestation gate

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/AzureAttestationBasedEnclaveProvider.cs

@@ -11,7 +11,6 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.Data.Common;
 using Microsoft.Extensions.Caching.Memory;
 using Microsoft.IdentityModel.JsonWebTokens;
 using Microsoft.IdentityModel.Logging;
@@ -399,7 +398,7 @@ await GetOpenIdConfigForSigningKeysAsync(attestationInstanceUrl, shouldForceUpda
                 {
                     // Wait for SigningKeyRetryInSec sec before retrying to download the signing keys again.
                     // The synchronous path blocks the calling thread here; the async path must not.
-                    await Task.Delay(SigningKeyRetryInSec * 1000, cancellationToken).ConfigureAwait(false);
+                    await Task.Delay(TimeSpan.FromSeconds(SigningKeyRetryInSec), cancellationToken).ConfigureAwait(false);
                 }
 
                 // In cases if we fail to validate the token, since we are using the old signing keys

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/EnclaveProviderBase.cs

@@ -109,9 +109,19 @@ internal abstract class EnclaveProviderBase : SqlColumnEncryptionEnclaveProvider
         // attestation info from the server; only the attestation service call itself is shared. This
         // is a deliberate trade: it buys an ownership model that is sound under ConfigureAwait(false)
         // at the cost of some redundant per-caller work during a cold start.
+        //
+        // This gate uses LockTimeoutMaxInMilliseconds directly and does not participate in the
+        // synchronous path's adaptive 'lockTimeoutInMilliseconds', which sync callers drive to zero on
+        // a failed acquisition and restore on a successful one. The decoupling is deliberate and
+        // symmetric: async callers neither degrade that value for sync callers nor are degraded by it,
+        // so the async timeout cannot be collapsed to zero by unrelated synchronous contention.
         private static readonly SemaphoreSlim s_asyncAttestationGate = new SemaphoreSlim(1, 1);
 
-        // It is used to save the attestation url and nonce value across API calls
+        // Records the managed thread IDs that are part way through an attestation sequence, so that
+        // GetEnclaveSessionHelper can tell a re-entrant call on such a thread from a fresh caller and
+        // avoid making it wait on a gate it effectively already holds. The entry is keyed by, and
+        // holds, the thread ID; entries expire after s_threadRetryCacheTimeout in case the sequence is
+        // abandoned. Used by the synchronous path only.
         protected static readonly MemoryCache ThreadRetryCache = new MemoryCache(new MemoryCacheOptions());
         private static readonly TimeSpan s_threadRetryCacheTimeout = TimeSpan.FromMinutes(10);
         #endregion
@@ -250,6 +260,11 @@ protected void GetEnclaveSessionHelper(EnclaveSessionParameters enclaveSessionPa
             return Task.FromResult((sqlEnclaveSession, counter, customData, customDataLength));
         }
 
+        // How long a caller waits for the async attestation gate before giving up and attesting on its
+        // own. Overridable so that tests can exercise the timeout fallthrough without waiting out the
+        // full production timeout; production providers use the default.
+        protected virtual int AsyncAttestationGateTimeoutInMilliseconds => LockTimeoutMaxInMilliseconds;
+
         // Indicates whether this provider's attestation protocol uses a client-generated nonce.
         // Mirrors the value each provider passes to GetEnclaveSessionHelper on the synchronous path.
         protected abstract bool GeneratesNonceForAttestation { get; }
@@ -293,7 +308,7 @@ protected void GetEnclaveSessionHelper(EnclaveSessionParameters enclaveSessionPa
             // rather than failing. This mirrors the synchronous design's deliberate choice to favour
             // progress over strict collapsing when the gate holder is unusually slow.
             bool gateAcquired = await s_asyncAttestationGate
-                .WaitAsync(LockTimeoutMaxInMilliseconds, cancellationToken)
+                .WaitAsync(AsyncAttestationGateTimeoutInMilliseconds, cancellationToken)
                 .ConfigureAwait(false);
 
             try

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/VirtualSecureModeEnclaveProvider.cs

@@ -106,7 +106,7 @@ protected override async Task<byte[]> MakeRequestAsync(string url, CancellationT
                 {
                     if (n != 0)
                     {
-                        await Task.Delay(EnclaveRetrySleepInSeconds * 1000, cancellationToken).ConfigureAwait(false);
+                        await Task.Delay(TimeSpan.FromSeconds(EnclaveRetrySleepInSeconds), cancellationToken).ConfigureAwait(false);
                     }
 
 #if NET

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/VirtualSecureModeEnclaveProviderBase.cs

@@ -9,7 +9,6 @@
 using System.Security.Cryptography.Pkcs;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.Data.Common;
 using Microsoft.Extensions.Caching.Memory;
 
 namespace Microsoft.Data.SqlClient

src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/AlwaysEncrypted/SqlColumnEncryptionEnclaveProviderAsyncShould.cs

@@ -8,6 +8,7 @@
 #nullable disable
 
 using System;
+using System.Diagnostics;
 using System.Collections.Generic;
 using System.Linq;
 using System.Security.Cryptography;
@@ -328,11 +329,99 @@ public async Task CreateEnclaveSessionAsync_ConcurrentColdStart_CompletesWithout
             // parameters. Only the attestation service call itself is shared.
             Assert.Equal(sessions.Length, provider.AttestationParametersCount);
 
+            // The gate serialized the round trips, so no two ever overlapped.
+            Assert.Equal(1, provider.MaxConcurrentAttestations);
+
             // The gate must still be usable for a subsequent, unrelated attestation.
             SqlEnclaveSession next = await AttestAsync(provider, NewSessionParameters());
             Assert.NotNull(next);
         }
 
+        /// <summary>
+        /// Verifies that a caller which cannot take the async gate within the timeout falls through and
+        /// attests on its own rather than failing or deadlocking, and that it does not release a gate
+        /// it never took.
+        /// </summary>
+        /// <remarks>
+        /// Over-releasing would throw <see cref="SemaphoreFullException"/> and, worse, would admit an
+        /// extra caller into the gate. The follow-up attestation asserts the gate is still balanced.
+        /// </remarks>
+        [Fact]
+        public async Task CreateEnclaveSessionAsync_WhenGateWaitTimesOut_AttestsAnyway()
+        {
+            FakeAttestationEnclaveProvider provider = new FakeAttestationEnclaveProvider(TimeSpan.Zero);
+            TaskCompletionSource<bool> hold =
+                new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            provider.HoldAttestation = hold;
+
+            // Distinct parameters, so the blocked caller cannot be satisfied by the post-gate cache
+            // re-check and is forced down the timeout fallthrough path.
+            EnclaveSessionParameters holderParameters = NewSessionParameters();
+            EnclaveSessionParameters blockedParameters = NewSessionParameters();
+
+            Task<SqlEnclaveSession> holder = Task.Run(() => AttestAsync(provider, holderParameters));
+
+            // The holder now owns the gate and is parked inside its attestation until we release it.
+            Assert.True(provider.AttestationStarted.Wait(TimeSpan.FromSeconds(30)));
+            provider.GateTimeoutInMilliseconds = 1;
+
+            Task<SqlEnclaveSession> blocked = Task.Run(() => AttestAsync(provider, blockedParameters));
+
+            // Reaching two attestations while the first is still parked is only possible if the second
+            // caller gave up on the gate. Without the fallthrough this wait times out.
+            await provider.WaitForAttestationCountAsync(2);
+            Assert.Equal(2, provider.MaxConcurrentAttestations);
+
+            hold.SetResult(true);
+            provider.HoldAttestation = null;
+
+            SqlEnclaveSession heldSession = await holder;
+            SqlEnclaveSession blockedSession = await blocked;
+
+            Assert.NotNull(heldSession);
+            Assert.NotNull(blockedSession);
+            Assert.NotEqual(heldSession.SessionId, blockedSession.SessionId);
+
+            // The gate must still be balanced and usable. An over-release would have thrown
+            // SemaphoreFullException; a lost release would hang this call.
+            provider.GateTimeoutInMilliseconds = 15 * 1000;
+            Assert.NotNull(await AttestAsync(provider, NewSessionParameters()));
+            Assert.Equal(3, provider.AttestationCount);
+        }
+
+        /// <summary>
+        /// Verifies that after a cached session is invalidated, the next async caller re-attests and
+        /// obtains a new session rather than returning the evicted one.
+        /// </summary>
+        [Fact]
+        public async Task GetEnclaveSessionAsync_AfterInvalidation_ReattestsAndReturnsNewSession()
+        {
+            FakeAttestationEnclaveProvider provider = new FakeAttestationEnclaveProvider(TimeSpan.Zero);
+            EnclaveSessionParameters parameters = NewSessionParameters();
+
+            SqlEnclaveSession original = await AttestAsync(provider, parameters);
+            Assert.NotNull(original);
+            Assert.Equal(1, provider.AttestationCount);
+
+            await provider.InvalidateEnclaveSessionAsync(parameters, original);
+
+            // The evicted session must no longer be visible to a plain lookup.
+            (SqlEnclaveSession afterInvalidation, _, _, _) =
+                await provider.GetEnclaveSessionAsync(parameters, generateCustomData: false, isRetry: false);
+            Assert.Null(afterInvalidation);
+
+            SqlEnclaveSession replacement = await AttestAsync(provider, parameters);
+
+            Assert.NotNull(replacement);
+            Assert.NotEqual(original.SessionId, replacement.SessionId);
+            Assert.Equal(2, provider.AttestationCount);
+
+            // The replacement is cached, so a subsequent caller reuses it.
+            SqlEnclaveSession cached = await AttestAsync(provider, parameters);
+            Assert.Equal(replacement.SessionId, cached.SessionId);
+            Assert.Equal(2, provider.AttestationCount);
+        }
+
         /// <summary>
         /// Verifies that sync and async callers can attest concurrently and that every caller ends up
         /// with a usable session.
@@ -631,6 +720,8 @@ private sealed class FakeAttestationEnclaveProvider : EnclaveProviderBase
             private static long s_nextSessionId;
 
             private readonly TimeSpan _attestationDelay;
+            private int _concurrentAttestations;
+            private int _maxConcurrentAttestations;
 
             private int _attestationCount;
 
@@ -649,8 +740,74 @@ internal FakeAttestationEnclaveProvider(TimeSpan attestationDelay)
 
             protected override bool GeneratesNonceForAttestation => true;
 
+            /// <summary>
+            /// How long this provider waits for the async attestation gate. Tests that exercise the
+            /// timeout fallthrough set a small value so they do not wait out the production timeout.
+            /// </summary>
+            internal int GateTimeoutInMilliseconds { get; set; } = 15 * 1000;
+
+            protected override int AsyncAttestationGateTimeoutInMilliseconds => GateTimeoutInMilliseconds;
+
             internal int AttestationCount => Volatile.Read(ref _attestationCount);
 
+            /// <summary>
+            /// The high-water mark of attestations running at the same time. A value of one proves the
+            /// gate serialized every caller; a value above one proves at least one caller entered
+            /// without holding the gate.
+            /// </summary>
+            internal int MaxConcurrentAttestations => Volatile.Read(ref _maxConcurrentAttestations);
+
+            /// <summary>
+            /// Signalled once an attestation round trip is actually under way. Tests use this instead of
+            /// a sleep so they do not race against first-call JIT and key generation costs.
+            /// </summary>
+            internal ManualResetEventSlim AttestationStarted { get; } = new ManualResetEventSlim(false);
+
+            /// <summary>
+            /// When set, every asynchronous attestation parks on this source until the test completes it.
+            /// This lets a test pin callers inside the attestation step and observe overlap without
+            /// depending on sleeps or scheduling luck.
+            /// </summary>
+            internal TaskCompletionSource<bool> HoldAttestation { get; set; }
+
+            /// <summary>
+            /// Spins until <see cref="AttestationCount"/> reaches <paramref name="count"/>.
+            /// </summary>
+            internal async Task WaitForAttestationCountAsync(int count)
+            {
+                Stopwatch stopwatch = Stopwatch.StartNew();
+                while (AttestationCount < count)
+                {
+                    Assert.True(
+                        stopwatch.Elapsed < TimeSpan.FromSeconds(30),
+                        $"Timed out waiting for {count} attestations; saw {AttestationCount}.");
+                    await Task.Delay(10).ConfigureAwait(false);
+                }
+            }
+
+            private void EnterAttestation()
+            {
+                int concurrent = Interlocked.Increment(ref _concurrentAttestations);
+
+                int observed = Volatile.Read(ref _maxConcurrentAttestations);
+                while (concurrent > observed)
+                {
+                    int previous = Interlocked.CompareExchange(
+                        ref _maxConcurrentAttestations, concurrent, observed);
+                    if (previous == observed)
+                    {
+                        break;
+                    }
+
+                    observed = previous;
+                }
+
+                Interlocked.Increment(ref _attestationCount);
+                AttestationStarted.Set();
+            }
+
+            private void ExitAttestation() => Interlocked.Decrement(ref _concurrentAttestations);
+
             /// <summary>
             /// How many times attestation parameters (including the client Diffie-Hellman key) were
             /// generated. Unlike <see cref="AttestationCount"/>, this step sits outside the async gate,
@@ -706,8 +863,15 @@ internal override void CreateEnclaveSession(
                     sqlEnclaveSession = GetEnclaveSessionFromCache(enclaveSessionParameters, out counter);
                     if (sqlEnclaveSession == null)
                     {
-                        Interlocked.Increment(ref _attestationCount);
-                        Thread.Sleep(_attestationDelay);
+                        EnterAttestation();
+                        try
+                        {
+                            Thread.Sleep(_attestationDelay);
+                        }
+                        finally
+                        {
+                            ExitAttestation();
+                        }
                         sqlEnclaveSession = AddEnclaveSessionToCache(
                             enclaveSessionParameters,
                             SharedSecret,
@@ -736,8 +900,21 @@ internal override void InvalidateEnclaveSession(
                 int customDataLength,
                 CancellationToken cancellationToken)
             {
-                Interlocked.Increment(ref _attestationCount);
-                await Task.Delay(_attestationDelay, cancellationToken).ConfigureAwait(false);
+                EnterAttestation();
+                try
+                {
+                    await Task.Delay(_attestationDelay, cancellationToken).ConfigureAwait(false);
+
+                    TaskCompletionSource<bool> hold = HoldAttestation;
+                    if (hold != null)
+                    {
+                        await hold.Task.ConfigureAwait(false);
+                    }
+                }
+                finally
+                {
+                    ExitAttestation();
+                }
 
                 if (FailNextAttestation)
                 {