Resolve key vault references concurrently

src/Microsoft.Extensions.Configuration.AzureAppConfiguration/AzureAppConfigurationKeyVaultOptions.cs

@@ -31,6 +31,11 @@ public class AzureAppConfigurationKeyVaultOptions
         internal TimeSpan? DefaultSecretRefreshInterval = null;
         internal bool IsKeyVaultRefreshConfigured = false;
 
+        /// <summary>
+        /// Flag to indicate whether Key Vault references should be resolved in parallel. Disabled by default.
+        /// </summary>
+        public bool ParallelSecretResolutionEnabled { get; set; }
+
         /// <summary>
         /// Sets the credentials used to authenticate to key vaults that have no registered <see cref="SecretClient"/>.
         /// </summary>

src/Microsoft.Extensions.Configuration.AzureAppConfiguration/AzureAppConfigurationOptions.cs

@@ -137,12 +137,17 @@ internal IEnumerable<IKeyValueAdapter> Adapters
         /// <summary>
         /// Flag to indicate whether Key Vault options have been configured.
         /// </summary>
-        internal bool IsKeyVaultConfigured { get; private set; } = false;
+        internal bool IsKeyVaultConfigured { get; private set; }
 
         /// <summary>
         /// Flag to indicate whether Key Vault secret values will be refreshed automatically.
         /// </summary>
-        internal bool IsKeyVaultRefreshConfigured { get; private set; } = false;
+        internal bool IsKeyVaultRefreshConfigured { get; private set; }
+
+        /// <summary>
+        /// Flag to indicate whether Key Vault references should be resolved in parallel.
+        /// </summary>
+        internal bool IsParallelSecretResolutionEnabled { get; private set; }
 
         /// <summary>
         /// Indicates all feature flag features used by the application.
@@ -520,6 +525,7 @@ public AzureAppConfigurationOptions ConfigureKeyVault(Action<AzureAppConfigurati
             _adapters.Add(new AzureKeyVaultKeyValueAdapter(new AzureKeyVaultSecretProvider(keyVaultOptions)));
 
             IsKeyVaultRefreshConfigured = keyVaultOptions.IsKeyVaultRefreshConfigured;
+            IsParallelSecretResolutionEnabled = keyVaultOptions.ParallelSecretResolutionEnabled;
             IsKeyVaultConfigured = true;
             return this;
         }

src/Microsoft.Extensions.Configuration.AzureAppConfiguration/AzureAppConfigurationProvider.cs

@@ -625,16 +625,19 @@ private async Task<Dictionary<string, string>> PrepareData(Dictionary<string, Co
                 _requestTracingOptions.ResetAiConfigurationTracing();
             }
 
-            foreach (KeyValuePair<string, ConfigurationSetting> kvp in data)
+            foreach (IKeyValueAdapter adapter in _options.Adapters)
             {
-                IEnumerable<KeyValuePair<string, string>> keyValuePairs = null;
+                await adapter.PreloadAsync(data.Values, _logger, cancellationToken).ConfigureAwait(false);
+            }
 
+            foreach (KeyValuePair<string, ConfigurationSetting> kvp in data)
+            {
                 if (_requestTracingEnabled && _requestTracingOptions != null)
                 {
                     _requestTracingOptions.UpdateAiConfigurationTracing(kvp.Value.ContentType);
                 }
 
-                keyValuePairs = await ProcessAdapters(kvp.Value, cancellationToken).ConfigureAwait(false);
+                IEnumerable<KeyValuePair<string, string>> keyValuePairs = await ProcessAdapters(kvp.Value, cancellationToken).ConfigureAwait(false);
 
                 foreach (KeyValuePair<string, string> kv in keyValuePairs)
                 {

src/Microsoft.Extensions.Configuration.AzureAppConfiguration/AzureKeyVaultReference/AzureKeyVaultKeyValueAdapter.cs

@@ -9,7 +9,6 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
-using System.Net.Mime;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
@@ -75,15 +74,12 @@ KeyVaultReferenceException CreateKeyVaultReferenceException(string message, Conf
 
         public bool CanProcess(ConfigurationSetting setting)
         {
-            if (setting == null ||
-                string.IsNullOrWhiteSpace(setting.Value) ||
-                string.IsNullOrWhiteSpace(setting.ContentType))
+            if (setting == null || string.IsNullOrWhiteSpace(setting.Value))
             {
                 return false;
             }
 
-            return setting.ContentType.TryParseContentType(out ContentType contentType)
-                && contentType.IsKeyVaultReference();
+            return setting.IsKeyVaultReference();
         }
 
         public void OnChangeDetected(ConfigurationSetting setting = null)
@@ -116,6 +112,79 @@ public bool NeedsRefresh()
             return _secretProvider.ShouldRefreshKeyVaultSecrets();
         }
 
+        public async Task PreloadAsync(IEnumerable<ConfigurationSetting> settings, Logger logger, CancellationToken cancellationToken)
+        {
+            if (settings == null || !_secretProvider.IsParallelSecretResolutionEnabled)
+            {
+                return;
+            }
+
+            HashSet<Uri> seen = null;
+            List<(KeyVaultSecretIdentifier Identifier, string Key, string Label)> toFetch = null;
+
+            foreach (ConfigurationSetting setting in settings)
+            {
+                if (!CanProcess(setting))
+                {
+                    continue;
+                }
+
+                string secretRefUri = ParseSecretReferenceUri(setting);
+
+                if (string.IsNullOrEmpty(secretRefUri) ||
+                    !Uri.TryCreate(secretRefUri, UriKind.Absolute, out Uri secretUri) ||
+                    !KeyVaultSecretIdentifier.TryCreate(secretUri, out KeyVaultSecretIdentifier secretIdentifier))
+                {
+                    // Invalid references are surfaced from ProcessKeyValue with full exception context.
+                    continue;
+                }
+
+                seen = seen ?? new HashSet<Uri>();
+
+                if (!seen.Add(secretIdentifier.SourceId))
+                {
+                    continue;
+                }
+
+                toFetch = toFetch ?? new List<(KeyVaultSecretIdentifier, string, string)>();
+                toFetch.Add((secretIdentifier, setting.Key, setting.Label));
+            }
+
+            if (toFetch == null)
+            {
+                return;
+            }
+
+            var tasks = new Task[toFetch.Count];
+
+            for (int i = 0; i < toFetch.Count; i++)
+            {
+                (KeyVaultSecretIdentifier identifier, string key, string label) = toFetch[i];
+                tasks[i] = PreloadSecretAsync(identifier, key, label, logger, cancellationToken);
+            }
+
+            await Task.WhenAll(tasks).ConfigureAwait(false);
+        }
+
+        private async Task PreloadSecretAsync(KeyVaultSecretIdentifier identifier, string key, string label, Logger logger, CancellationToken cancellationToken)
+        {
+            try
+            {
+                await _secretProvider.GetSecretValue(identifier, key, label, logger, cancellationToken).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                throw;
+            }
+            catch
+            {
+                // Per-secret failures are deferred so ProcessKeyValue can throw a properly populated
+                // KeyVaultReferenceException. Evict the negative cache entry written by GetSecretValue
+                // so the retry actually re-fetches instead of returning a cached null within the backoff.
+                _secretProvider.RemoveSecretFromCache(identifier.SourceId);
+            }
+        }
+
         private string ParseSecretReferenceUri(ConfigurationSetting setting)
         {
             string secretRefUri = null;

src/Microsoft.Extensions.Configuration.AzureAppConfiguration/AzureKeyVaultReference/AzureKeyVaultSecretProvider.cs

@@ -4,8 +4,8 @@
 using Azure.Security.KeyVault.Secrets;
 using Microsoft.Extensions.Configuration.AzureAppConfiguration.Extensions;
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -14,16 +14,16 @@ namespace Microsoft.Extensions.Configuration.AzureAppConfiguration.AzureKeyVault
     internal class AzureKeyVaultSecretProvider
     {
         private readonly AzureAppConfigurationKeyVaultOptions _keyVaultOptions;
-        private readonly IDictionary<string, SecretClient> _secretClients;
-        private readonly Dictionary<Uri, CachedKeyVaultSecret> _cachedKeyVaultSecrets;
-        private Uri _nextRefreshSourceId;
-        private DateTimeOffset? _nextRefreshTime;
+        private readonly ConcurrentDictionary<string, SecretClient> _secretClients;
+        private readonly ConcurrentDictionary<Uri, CachedKeyVaultSecret> _cachedKeyVaultSecrets;
+
+        public bool IsParallelSecretResolutionEnabled => _keyVaultOptions.ParallelSecretResolutionEnabled;
 
         public AzureKeyVaultSecretProvider(AzureAppConfigurationKeyVaultOptions keyVaultOptions = null)
         {
             _keyVaultOptions = keyVaultOptions ?? new AzureAppConfigurationKeyVaultOptions();
-            _cachedKeyVaultSecrets = new Dictionary<Uri, CachedKeyVaultSecret>();
-            _secretClients = new Dictionary<string, SecretClient>(StringComparer.OrdinalIgnoreCase);
+            _cachedKeyVaultSecrets = new ConcurrentDictionary<Uri, CachedKeyVaultSecret>();
+            _secretClients = new ConcurrentDictionary<string, SecretClient>(StringComparer.OrdinalIgnoreCase);
 
             if (_keyVaultOptions.SecretClients != null)
             {
@@ -52,6 +52,7 @@ public async Task<string> GetSecretValue(KeyVaultSecretIdentifier secretIdentifi
                 throw new UnauthorizedAccessException("No key vault credential or secret resolver callback configured, and no matching secret client could be found.");
             }
 
+            CachedKeyVaultSecret updatedCachedSecret = null;
             bool success = false;
 
             try
@@ -68,55 +69,44 @@ public async Task<string> GetSecretValue(KeyVaultSecretIdentifier secretIdentifi
                     secretValue = await _keyVaultOptions.SecretResolver(secretIdentifier.SourceId).ConfigureAwait(false);
                 }
 
-                cachedSecret = new CachedKeyVaultSecret(secretValue, secretIdentifier.SourceId);
+                updatedCachedSecret = new CachedKeyVaultSecret(secretValue, secretIdentifier.SourceId);
                 success = true;
             }
             finally
             {
-                SetSecretInCache(secretIdentifier.SourceId, key, cachedSecret, success);
+                SetSecretInCache(secretIdentifier.SourceId, key, updatedCachedSecret, success);
             }
 
             return secretValue;
         }
 
         public bool ShouldRefreshKeyVaultSecrets()
         {
-            return _nextRefreshTime.HasValue && _nextRefreshTime.Value < DateTimeOffset.UtcNow;
-        }
-
-        public void ClearCache()
-        {
-            var sourceIdsToRemove = new List<Uri>();
-
-            var utcNow = DateTimeOffset.UtcNow;
-
             foreach (KeyValuePair<Uri, CachedKeyVaultSecret> secret in _cachedKeyVaultSecrets)
             {
-                if (secret.Value.LastRefreshTime + RefreshConstants.MinimumSecretRefreshInterval < utcNow)
+                if (secret.Value.RefreshAt.HasValue && secret.Value.RefreshAt.Value < DateTimeOffset.UtcNow)
                 {
-                    sourceIdsToRemove.Add(secret.Key);
+                    return true;
                 }
             }
 
-            foreach (Uri sourceId in sourceIdsToRemove)
-            {
-                _cachedKeyVaultSecrets.Remove(sourceId);
-            }
+            return false;
+        }
 
-            if (_cachedKeyVaultSecrets.Any())
+        public void ClearCache()
+        {
+            foreach (KeyValuePair<Uri, CachedKeyVaultSecret> secret in _cachedKeyVaultSecrets)
             {
-                UpdateNextRefreshableSecretFromCache();
+                if (secret.Value.LastRefreshTime + RefreshConstants.MinimumSecretRefreshInterval < DateTimeOffset.UtcNow)
+                {
+                    _cachedKeyVaultSecrets.TryRemove(secret.Key, out _);
+                }
             }
         }
 
         public void RemoveSecretFromCache(Uri sourceId)
         {
-            _cachedKeyVaultSecrets.Remove(sourceId);
-
-            if (sourceId == _nextRefreshSourceId)
-            {
-                UpdateNextRefreshableSecretFromCache();
-            }
+            _cachedKeyVaultSecrets.TryRemove(sourceId, out _);
         }
 
         private SecretClient GetSecretClient(Uri secretUri)
@@ -133,14 +123,12 @@ private SecretClient GetSecretClient(Uri secretUri)
                 return null;
             }
 
-            client = new SecretClient(
-                new Uri(secretUri.GetLeftPart(UriPartial.Authority)),
-                _keyVaultOptions.Credential,
-                _keyVaultOptions.ClientOptions);
-
-            _secretClients.Add(keyVaultId, client);
-
-            return client;
+            return _secretClients.GetOrAdd(
+                keyVaultId,
+                _ => new SecretClient(
+                    new Uri(secretUri.GetLeftPart(UriPartial.Authority)),
+                    _keyVaultOptions.Credential,
+                    _keyVaultOptions.ClientOptions));
         }
 
         private void SetSecretInCache(Uri sourceId, string key, CachedKeyVaultSecret cachedSecret, bool success = true)
@@ -152,37 +140,6 @@ private void SetSecretInCache(Uri sourceId, string key, CachedKeyVaultSecret cac
 
             UpdateCacheExpirationTimeForSecret(key, cachedSecret, success);
             _cachedKeyVaultSecrets[sourceId] = cachedSecret;
-
-            if (sourceId == _nextRefreshSourceId)
-            {
-                UpdateNextRefreshableSecretFromCache();
-            }
-            else if ((cachedSecret.RefreshAt.HasValue && _nextRefreshTime.HasValue && cachedSecret.RefreshAt.Value < _nextRefreshTime.Value)
-                    || (cachedSecret.RefreshAt.HasValue && !_nextRefreshTime.HasValue))
-            {
-                _nextRefreshSourceId = sourceId;
-                _nextRefreshTime = cachedSecret.RefreshAt.Value;
-            }
-        }
-
-        private void UpdateNextRefreshableSecretFromCache()
-        {
-            _nextRefreshSourceId = null;
-            _nextRefreshTime = DateTimeOffset.MaxValue;
-
-            foreach (KeyValuePair<Uri, CachedKeyVaultSecret> secret in _cachedKeyVaultSecrets)
-            {
-                if (secret.Value.RefreshAt.HasValue && secret.Value.RefreshAt.Value < _nextRefreshTime)
-                {
-                    _nextRefreshTime = secret.Value.RefreshAt;
-                    _nextRefreshSourceId = secret.Key;
-                }
-            }
-
-            if (_nextRefreshTime == DateTimeOffset.MaxValue)
-            {
-                _nextRefreshTime = null;
-            }
         }
 
         private void UpdateCacheExpirationTimeForSecret(string key, CachedKeyVaultSecret cachedSecret, bool success)

src/Microsoft.Extensions.Configuration.AzureAppConfiguration/Extensions/ConfigurationSettingExtensions.cs

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+//
+using Azure.Data.AppConfiguration;
+using Microsoft.Extensions.Configuration.AzureAppConfiguration.AzureKeyVault;
+using System.Net.Mime;
+
+namespace Microsoft.Extensions.Configuration.AzureAppConfiguration.Extensions
+{
+    internal static class ConfigurationSettingExtensions
+    {
+        public static bool IsKeyVaultReference(this ConfigurationSetting setting)
+        {
+            return setting != null
+                && setting.ContentType.TryParseContentType(out ContentType contentType)
+                && contentType.IsKeyVaultReference();
+        }
+    }
+}

src/Microsoft.Extensions.Configuration.AzureAppConfiguration/FeatureManagement/FeatureManagementKeyValueAdapter.cs

@@ -87,6 +87,11 @@ public void OnConfigUpdated()
             return;
         }
 
+        public Task PreloadAsync(IEnumerable<ConfigurationSetting> settings, Logger logger, CancellationToken cancellationToken)
+        {
+            return Task.CompletedTask;
+        }
+
         private List<KeyValuePair<string, string>> ProcessDotnetSchemaFeatureFlag(FeatureFlag featureFlag, ConfigurationSetting setting, Uri endpoint)
         {
             var keyValues = new List<KeyValuePair<string, string>>();

src/Microsoft.Extensions.Configuration.AzureAppConfiguration/IKeyValueAdapter.cs

@@ -13,6 +13,8 @@ internal interface IKeyValueAdapter
     {
         Task<IEnumerable<KeyValuePair<string, string>>> ProcessKeyValue(ConfigurationSetting setting, Uri endpoint, Logger logger, CancellationToken cancellationToken);
 
+        Task PreloadAsync(IEnumerable<ConfigurationSetting> settings, Logger logger, CancellationToken cancellationToken);
+
         bool CanProcess(ConfigurationSetting setting);
 
         void OnChangeDetected(ConfigurationSetting setting = null);

src/Microsoft.Extensions.Configuration.AzureAppConfiguration/JsonKeyValueAdapter.cs

@@ -80,5 +80,10 @@ public bool NeedsRefresh()
         {
             return false;
         }
+
+        public Task PreloadAsync(IEnumerable<ConfigurationSetting> settings, Logger logger, CancellationToken cancellationToken)
+        {
+            return Task.CompletedTask;
+        }
     }
 }