Switch for metadata refresh to be blocking (#3193)

* Adds the ability for the metadata refresh to be done as a blocking call, as per 8.0.1 behavior.

This is done through the Switch.Microsoft.IdentityModel.UpdateConfigAsBlocking switch.

If set, configuration calls will be blocking on update, and exceptions when requesting new metadata will be returned to the caller
---------

Co-authored-by: Keegan Caruso <[email protected]>

Author: keegan-caruso

src/Microsoft.IdentityModel.Protocols/Configuration/ConfigurationManager.cs

@@ -17,7 +17,7 @@ namespace Microsoft.IdentityModel.Protocols
     /// </summary>
     /// <typeparam name="T">The type of <see cref="IDocumentRetriever"/>.</typeparam>
     [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Design", "CA1001:TypesThatOwnDisposableFieldsShouldBeDisposable")]
-    public class ConfigurationManager<T> : BaseConfigurationManager, IConfigurationManager<T> where T : class
+    public partial class ConfigurationManager<T> : BaseConfigurationManager, IConfigurationManager<T> where T : class
     {
         internal Action _onBackgroundTaskFinish;
 
@@ -138,7 +138,7 @@ public ConfigurationManager(string metadataAddress, IConfigurationRetriever<T> c
         /// <summary>
         /// Obtains an updated version of Configuration.
         /// </summary>
-        /// <returns>Configuration of type T.</returns>
+        /// <returns>Configuration of type <typeparamref name="T"/>.</returns>
         /// <remarks>If the time since the last call is less than <see cref="BaseConfigurationManager.AutomaticRefreshInterval"/> then <see cref="IConfigurationRetriever{T}.GetConfigurationAsync"/> is not called and the current Configuration is returned.</remarks>
         public async Task<T> GetConfigurationAsync()
         {
@@ -149,13 +149,40 @@ public async Task<T> GetConfigurationAsync()
         /// Obtains an updated version of Configuration.
         /// </summary>
         /// <param name="cancel">CancellationToken</param>
-        /// <returns>Configuration of type T.</returns>
-        /// <remarks>If the time since the last call is less than <see cref="BaseConfigurationManager.AutomaticRefreshInterval"/> then <see cref="IConfigurationRetriever{T}.GetConfigurationAsync"/> is not called and the current Configuration is returned.</remarks>
+        /// <returns>Configuration of type <typeparamref name="T"/>.</returns>
+        /// <remarks>
+        /// <para>
+        /// If the time since the last call is less than <see cref="BaseConfigurationManager.AutomaticRefreshInterval"/>
+        /// then <see cref="IConfigurationRetriever{T}.GetConfigurationAsync"/> is not called and the current Configuration is returned.
+        /// By default, this method blocks until the configuration is retrieved the first time. After the configuration was retrieved once,
+        /// updates will happen in the background. Failures to retrieve the configuration on the background thread will be logged.
+        /// </para>
+        /// <para>
+        /// If this operation is configured to be blocking through the switch 'Switch.Microsoft.IdentityModel.UpdateConfigAsBlocking'
+        /// then this method will block each time the configuration needs to be updated or hasn't been retrieved. If the configuration
+        /// cannot be initially retrieved an exception will be thrown. If the configuration has been retrieved, but cannot be updated,
+        /// then the exception will be logged and the current configuration will be returned.
+        /// </para>
+        /// <para>
+        /// By using the app context switch you choose what works best for you when there is a signing key update:
+        /// either block requests from being validated until the new key is retrieved, or allow requests to be validated
+        /// with the current key until the new key is retrieved. If blocking, a service receiving high concurrent request
+        /// may experience thread starvation.
+        /// </para>
+        /// </remarks>
         public virtual async Task<T> GetConfigurationAsync(CancellationToken cancel)
         {
             if (_currentConfiguration != null && _syncAfter > TimeProvider.GetUtcNow())
                 return _currentConfiguration;
 
+            if (AppContextSwitches.UpdateConfigAsBlocking)
+                return await GetConfigurationWithBlockingAsync(cancel).ConfigureAwait(false);
+            else
+                return await GetConfigurationNonBlockingAsync(cancel).ConfigureAwait(false);
+        }
+
+        private async Task<T> GetConfigurationNonBlockingAsync(CancellationToken cancel)
+        {
             Exception fetchMetadataFailure = null;
 
             // LOGIC
@@ -174,7 +201,6 @@ public virtual async Task<T> GetConfigurationAsync(CancellationToken cancel)
                     return _currentConfiguration;
                 }
 
-#pragma warning disable CA1031 // Do not catch general exception types
                 try
                 {
                     // Don't use the individual CT here, this is a shared operation that shouldn't be affected by an individual's cancellation.
@@ -205,6 +231,7 @@ public virtual async Task<T> GetConfigurationAsync(CancellationToken cancel)
 
                     UpdateConfiguration(configuration);
                 }
+#pragma warning disable CA1031 // Do not catch general exception types
                 catch (Exception ex)
                 {
                     fetchMetadataFailure = ex;
@@ -221,11 +248,11 @@ public virtual async Task<T> GetConfigurationAsync(CancellationToken cancel)
                                 LogHelper.MarkAsNonPII(ex)),
                             ex));
                 }
+#pragma warning restore CA1031 // Do not catch general exception types
                 finally
                 {
                     _configurationNullLock.Release();
                 }
-#pragma warning restore CA1031 // Do not catch general exception types
             }
             else
             {
@@ -260,7 +287,6 @@ public virtual async Task<T> GetConfigurationAsync(CancellationToken cancel)
         /// </summary>
         private void UpdateCurrentConfiguration()
         {
-#pragma warning disable CA1031 // Do not catch general exception types
             long startTimestamp = TimeProvider.GetTimestamp();
 
             try
@@ -293,6 +319,7 @@ private void UpdateCurrentConfiguration()
                         UpdateConfiguration(configuration);
                 }
             }
+#pragma warning disable CA1031 // Do not catch general exception types
             catch (Exception ex)
             {
                 var elapsedTime = TimeProvider.GetElapsedTime(startTimestamp);
@@ -309,13 +336,13 @@ private void UpdateCurrentConfiguration()
                             ex),
                         ex));
             }
+#pragma warning restore CA1031 // Do not catch general exception types
             finally
             {
                 Interlocked.Exchange(ref _configurationRetrieverState, ConfigurationRetrieverIdle);
             }
 
             _onBackgroundTaskFinish?.Invoke();
-#pragma warning restore CA1031 // Do not catch general exception types
         }
 
         private void UpdateConfiguration(T configuration)
@@ -343,7 +370,20 @@ public override async Task<BaseConfiguration> GetBaseConfigurationAsync(Cancella
         /// <para>2. The time between when this method was called and DateTimeOffset.Now is greater than <see cref="BaseConfigurationManager.RefreshInterval"/>.</para>
         /// <para>If <see cref="BaseConfigurationManager.RefreshInterval"/> == <see cref="TimeSpan.MaxValue"/> then this method does nothing.</para>
         /// </summary>
+        /// <remarks>
+        /// If the strategy is configured to be blocking through the switch 'Switch.Microsoft.IdentityModel.UpdateConfigAsBlocking',
+        /// then this method will not update the configuration, instead it will request the next call to <see cref="GetConfigurationAsync()"/>
+        /// should request new configuration.
+        /// </remarks>
         public override void RequestRefresh()
+        {
+            if (AppContextSwitches.UpdateConfigAsBlocking)
+                RequestRefreshBlocking();
+            else
+                RequestRefreshBackgroundThread();
+        }
+
+        private void RequestRefreshBackgroundThread()
         {
             DateTimeOffset now = TimeProvider.GetUtcNow();
 

src/Microsoft.IdentityModel.Protocols/Configuration/ConfigurationManager_Blocking.cs

@@ -0,0 +1,163 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.IdentityModel.Logging;
+using Microsoft.IdentityModel.Protocols.Configuration;
+using Microsoft.IdentityModel.Telemetry;
+using Microsoft.IdentityModel.Tokens;
+
+namespace Microsoft.IdentityModel.Protocols
+{
+    partial class ConfigurationManager<T> where T : class
+    {
+        private readonly SemaphoreSlim _refreshLock = new(1, 1);
+        private TimeSpan _bootstrapRefreshInterval = TimeSpan.FromSeconds(1);
+
+        /// <summary>
+        /// Only used to track the type of request for telemetry.
+        /// </summary>
+        private bool _refreshRequested;
+
+        private async Task<T> GetConfigurationWithBlockingAsync(CancellationToken cancel)
+        {
+            Exception _fetchMetadataFailure = null;
+            await _refreshLock.WaitAsync(cancel).ConfigureAwait(false);
+
+            long startTimestamp = TimeProvider.GetTimestamp();
+
+            try
+            {
+                if (_syncAfter <= TimeProvider.GetUtcNow())
+                {
+                    try
+                    {
+                        // Don't use the individual CT here, this is a shared operation that shouldn't be affected by an individual's cancellation.
+                        // The transport should have it's own timeouts, etc..
+                        var configuration = await _configRetriever.GetConfigurationAsync(MetadataAddress, _docRetriever, CancellationToken.None).ConfigureAwait(false);
+
+                        var elapsedTime = TimeProvider.GetElapsedTime(startTimestamp);
+                        TelemetryClient.LogConfigurationRetrievalDuration(
+                            MetadataAddress,
+                            elapsedTime);
+
+                        if (_configValidator != null)
+                        {
+                            ConfigurationValidationResult result = _configValidator.Validate(configuration);
+                            if (!result.Succeeded)
+                                throw LogHelper.LogExceptionMessage(new InvalidConfigurationException(LogHelper.FormatInvariant(LogMessages.IDX20810, result.ErrorMessage)));
+                        }
+
+                        _lastRequestRefresh = TimeProvider.GetUtcNow().UtcDateTime;
+
+                        TelemetryForUpdateBlocking();
+
+                        if (_refreshRequested)
+                            _refreshRequested = false;
+
+                        UpdateConfiguration(configuration);
+                    }
+                    catch (Exception ex)
+                    {
+                        _fetchMetadataFailure = ex;
+
+                        if (_currentConfiguration == null)
+                        {
+                            if (_bootstrapRefreshInterval < RefreshInterval)
+                            {
+                                // Adopt exponential backoff for bootstrap refresh interval with a decorrelated jitter if it is not longer than the refresh interval.
+                                TimeSpan _bootstrapRefreshIntervalWithJitter = TimeSpan.FromSeconds(new Random().Next((int)_bootstrapRefreshInterval.TotalSeconds));
+                                _bootstrapRefreshInterval += _bootstrapRefreshInterval;
+                                _syncAfter = DateTimeUtil.Add(DateTime.UtcNow, _bootstrapRefreshIntervalWithJitter);
+                            }
+                            else
+                            {
+                                _syncAfter = DateTimeUtil.Add(
+                                    TimeProvider.GetUtcNow().UtcDateTime,
+                                    AutomaticRefreshInterval < RefreshInterval ? AutomaticRefreshInterval : RefreshInterval);
+                            }
+
+                            TelemetryClient.IncrementConfigurationRefreshRequestCounter(
+                                MetadataAddress,
+                                TelemetryConstants.Protocols.FirstRefresh,
+                                ex);
+
+                            throw LogHelper.LogExceptionMessage(
+                                new InvalidOperationException(
+                                    LogHelper.FormatInvariant(LogMessages.IDX20803, LogHelper.MarkAsNonPII(MetadataAddress ?? "null"), LogHelper.MarkAsNonPII(_syncAfter), LogHelper.MarkAsNonPII(ex)), ex));
+                        }
+                        else
+                        {
+                            _syncAfter = DateTimeUtil.Add(
+                                TimeProvider.GetUtcNow().UtcDateTime,
+                                AutomaticRefreshInterval < RefreshInterval ? AutomaticRefreshInterval : RefreshInterval);
+
+                            var elapsedTime = TimeProvider.GetElapsedTime(startTimestamp);
+
+                            TelemetryClient.LogConfigurationRetrievalDuration(
+                                MetadataAddress,
+                                elapsedTime,
+                                ex);
+
+                            LogHelper.LogExceptionMessage(
+                                new InvalidOperationException(
+                                    LogHelper.FormatInvariant(LogMessages.IDX20806, LogHelper.MarkAsNonPII(MetadataAddress ?? "null"), LogHelper.MarkAsNonPII(ex)), ex));
+                        }
+                    }
+                }
+
+                // Stale metadata is better than no metadata
+                if (_currentConfiguration != null)
+                    return _currentConfiguration;
+                else
+                    throw LogHelper.LogExceptionMessage(
+                        new InvalidOperationException(
+                            LogHelper.FormatInvariant(
+                                LogMessages.IDX20803,
+                                LogHelper.MarkAsNonPII(MetadataAddress ?? "null"),
+                                LogHelper.MarkAsNonPII(_syncAfter),
+                                LogHelper.MarkAsNonPII(_fetchMetadataFailure)),
+                            _fetchMetadataFailure));
+            }
+            finally
+            {
+                _refreshLock.Release();
+            }
+        }
+
+        private void RequestRefreshBlocking()
+        {
+            DateTime now = TimeProvider.GetUtcNow().UtcDateTime;
+
+            if (now >= DateTimeUtil.Add(_lastRequestRefresh.UtcDateTime, RefreshInterval) || _isFirstRefreshRequest)
+            {
+                _refreshRequested = true;
+                _syncAfter = now;
+                _isFirstRefreshRequest = false;
+            }
+        }
+
+        private void TelemetryForUpdateBlocking()
+        {
+            string updateMode;
+
+            if (_currentConfiguration is null)
+                updateMode = TelemetryConstants.Protocols.FirstRefresh;
+            else
+                updateMode = _refreshRequested ? TelemetryConstants.Protocols.Manual : TelemetryConstants.Protocols.Automatic;
+
+            try
+            {
+                TelemetryClient.IncrementConfigurationRefreshRequestCounter(
+                    MetadataAddress,
+                    updateMode);
+            }
+#pragma warning disable CA1031 // Do not catch general exception types
+            catch
+            { }
+#pragma warning restore CA1031 // Do not catch general exception types
+        }
+    }
+}

src/Microsoft.IdentityModel.Tokens/AppContextSwitches.cs

@@ -70,7 +70,11 @@ internal static class AppContextSwitches
 
         internal static bool UseRfcDefinitionOfEpkAndKid => _useRfcDefinitionOfEpkAndKid ??= (AppContext.TryGetSwitch(UseRfcDefinitionOfEpkAndKidSwitch, out bool isEnabled) && isEnabled);
 
-
+        /// <summary>
+        /// Enabling this switch will cause the configuration manager to block other requests to GetConfigurationAsync if a request is already in progress.
+        /// The default configuration refresh behavior is if a request is already in progress, the current configuration will be returned until the ongoing request is completed on
+        /// a background thread.
+        /// </summary>
         internal const string UpdateConfigAsBlockingSwitch = "Switch.Microsoft.IdentityModel.UpdateConfigAsBlocking";
 
         private static bool? _updateConfigAsBlockingCall;
@@ -96,6 +100,9 @@ internal static void ResetAllSwitches()
 
             _useRfcDefinitionOfEpkAndKid = null;
             AppContext.SetSwitch(UseRfcDefinitionOfEpkAndKidSwitch, false);
+
+            _updateConfigAsBlockingCall = null;
+            AppContext.SetSwitch(UpdateConfigAsBlockingSwitch, false);
         }
     }
 }

src/Microsoft.IdentityModel.Tokens/Telemetry/ITelemetryClient.cs

@@ -25,8 +25,6 @@ internal void IncrementConfigurationRefreshRequestCounter(
             string operationStatus,
             Exception exception);
 
-        // Unused, this was part of a previous release, since it is a friend,
-        // it cannot be removed.
         internal void LogBackgroundConfigurationRefreshFailure(
             string metadataAddress,
             Exception exception);

src/Microsoft.IdentityModel.Tokens/Telemetry/TelemetryClient.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using Microsoft.IdentityModel.Logging;
 using Microsoft.IdentityModel.Tokens;
@@ -15,13 +16,19 @@ internal class TelemetryClient : ITelemetryClient
     {
         public string ClientVer = IdentityModelTelemetryUtil.ClientVer;
 
+        private KeyValuePair<string, object> _blockingTagValue = new(
+            TelemetryConstants.BlockingTypeTag,
+            AppContextSwitches.UpdateConfigAsBlocking.ToString()
+        );
+
         public void IncrementConfigurationRefreshRequestCounter(string metadataAddress, string operationStatus)
         {
             var tagList = new TagList()
             {
                 { TelemetryConstants.IdentityModelVersionTag, ClientVer },
                 { TelemetryConstants.MetadataAddressTag, metadataAddress },
-                { TelemetryConstants.OperationStatusTag, operationStatus }
+                { TelemetryConstants.OperationStatusTag, operationStatus },
+                _blockingTagValue
             };
 
             TelemetryDataRecorder.IncrementConfigurationRefreshRequestCounter(tagList);
@@ -34,7 +41,8 @@ public void IncrementConfigurationRefreshRequestCounter(string metadataAddress,
                 { TelemetryConstants.IdentityModelVersionTag, ClientVer },
                 { TelemetryConstants.MetadataAddressTag, metadataAddress },
                 { TelemetryConstants.OperationStatusTag, operationStatus },
-                { TelemetryConstants.ExceptionTypeTag, exception.GetType().ToString() }
+                { TelemetryConstants.ExceptionTypeTag, exception.GetType().ToString() },
+                _blockingTagValue
             };
 
             TelemetryDataRecorder.IncrementConfigurationRefreshRequestCounter(tagList);
@@ -58,7 +66,8 @@ public void LogConfigurationRetrievalDuration(string metadataAddress, TimeSpan o
             {
                 { TelemetryConstants.IdentityModelVersionTag, ClientVer },
                 { TelemetryConstants.MetadataAddressTag, metadataAddress },
-                { TelemetryConstants.ExceptionTypeTag, exception.GetType().ToString() }
+                { TelemetryConstants.ExceptionTypeTag, exception.GetType().ToString() },
+                _blockingTagValue
             };
 
             long durationInMilliseconds = (long)operationDuration.TotalMilliseconds;
@@ -74,7 +83,7 @@ public void LogBackgroundConfigurationRefreshFailure(
                 { TelemetryConstants.IdentityModelVersionTag, ClientVer },
                 { TelemetryConstants.MetadataAddressTag, metadataAddress },
                 { TelemetryConstants.ExceptionTypeTag, exception.GetType().ToString() },
-                { TelemetryConstants.BlockingTypeTag, AppContextSwitches.UpdateConfigAsBlocking.ToString() }
+                _blockingTagValue
             };
 
             TelemetryDataRecorder.IncrementBackgroundConfigurationRefreshFailureCounter(tagList);