hybrid cache with GetOrDefault

Author: Erwinvandervalk

access-token-management/src/AccessTokenManagement/DistributedClientCredentialsTokenCache.cs

@@ -20,6 +20,7 @@ public static class ServiceProviderKeys
 public class DistributedClientCredentialsTokenCache(
     [FromKeyedServices(ServiceProviderKeys.DistributedClientCredentialsTokenCache)]HybridCache cache,
     TimeProvider time,
+    ITokenRequestSynchronization synchronization,
     IOptions<ClientCredentialsTokenManagementOptions> options,
     ILogger<DistributedClientCredentialsTokenCache> logger
     )
@@ -38,9 +39,10 @@ public async Task SetAsync(
 
         try
         {
-            var entryOptions = GetHybridCacheEntryOptions(clientName, clientCredentialsToken);
+            var entryOptions = GetHybridCacheEntryOptions(clientCredentialsToken);
 
             var cacheKey = GenerateCacheKey(_options, clientName, requestParameters);
+            logger.LogTrace("Caching access token for client: {clientName}. Expiration: {expiration}", clientName, entryOptions.Expiration);
             await cache.SetAsync(cacheKey, clientCredentialsToken, entryOptions, cancellationToken: cancellationToken).ConfigureAwait(false);
         }
         catch (Exception e)
@@ -51,17 +53,14 @@ public async Task SetAsync(
         }
     }
 
-    private HybridCacheEntryOptions GetHybridCacheEntryOptions(string clientName,
-        ClientCredentialsToken clientCredentialsToken)
+    private HybridCacheEntryOptions GetHybridCacheEntryOptions(ClientCredentialsToken clientCredentialsToken)
     {
         var absoluteCacheExpiration = clientCredentialsToken.Expiration.AddSeconds(-_options.CacheLifetimeBuffer);
         var relativeCacheExpiration = absoluteCacheExpiration - time.GetUtcNow();
         var entryOptions = new HybridCacheEntryOptions()
         {
             Expiration = relativeCacheExpiration
         };
-
-        logger.LogTrace("Caching access token for client: {clientName}. Expiration: {expiration}", clientName, absoluteCacheExpiration);
         return entryOptions;
     }
 
@@ -78,52 +77,52 @@ public async Task<ClientCredentialsToken> GetOrCreateAsync(
 
         var cacheKey = GenerateCacheKey(_options, clientName, requestParameters);
 
-        ClientCredentialsToken token;
+        ClientCredentialsToken? token;
         if (!requestParameters.ForceRenewal)
         {
-            try
-            {
-                // We don't need the token to be absolutely fresh, so we can get one from the cache. 
-                // The GetOrCreate pattern unfortunately doesn't allow us to create a cache entry
-                // that's only valid for as long as it needs to be. 
-                token = await cache.GetOrCreateAsync(
-                    key: cacheKey,
-                    factory: async (ct) =>
-                    {
-                        var result = await factory(clientName, requestParameters, ct).ConfigureAwait(false);
-                        if (result.IsError)
-                        {
-                            // If the token is an error, we throw an exception to prevent the value from being cached
-                            throw new TokenErrorException(result);
-                        }
-
-                        return result;
-                    },
-                    cancellationToken: cancellationToken);
-            }
-            catch (TokenErrorException ex)
-            {
-                return ex.Token;
-            }
+            // We don't need the token to be absolutely fresh, so we can get one from the cache. 
+            token = await cache.GetOrDefaultAsync<ClientCredentialsToken>(
+                key: cacheKey,
+                cancellationToken: cancellationToken);
 
-            if (token.Expiration != DateTimeOffset.MinValue)
+            if (token?.Expiration > DateTimeOffset.MinValue)
             {
                 var absoluteCacheExpiration = token.Expiration.AddSeconds(-_options.CacheLifetimeBuffer);
 
                 if (absoluteCacheExpiration > time.GetUtcNow())
                 {
+                    // It's possible that we have only read the token from L2 cache, not L1 cache. 
+                    // just to be sure, write the token also into L1 cache (which should be fast)
+                    // https://github.com/dotnet/extensions/issues/5688#issuecomment-2692247434
+                    var defaultWriteOptions = GetHybridCacheEntryOptions(token);
+                    await cache.SetAsync(cacheKey, token, new HybridCacheEntryOptions()
+                    {
+                        Flags = HybridCacheEntryFlags.DisableDistributedCacheWrite,
+                        Expiration = defaultWriteOptions.Expiration,
+                        LocalCacheExpiration = defaultWriteOptions.LocalCacheExpiration
+                    }, cancellationToken: cancellationToken);
+
                     return token;
                 }
             }
         }
-        token = await factory(clientName, requestParameters, cancellationToken).ConfigureAwait(false);
 
-        if (!token.IsError)
+        // Apparently, there's either no value in the cache, or we want a fresh one.
+        // Since we aren't using GetOrCreate, we'll have to protect against cache stampedes ourselves.
+        return await synchronization.SynchronizeAsync(cacheKey, async () =>
         {
-            await SetAsync(clientName, token, requestParameters, cancellationToken);
-        }
+            token = await factory(clientName, requestParameters, cancellationToken).ConfigureAwait(false);
+
+            // Don't cache the token if there's an error. 
+            if (!token.IsError)
+            {
+                await SetAsync(clientName, token, requestParameters, cancellationToken);
+            }
+
+            return token;
+        });
 
-        return token;
+        
     }
 
 

access-token-management/src/AccessTokenManagement/DistributedDPoPNonceStore.cs

@@ -4,6 +4,7 @@
 using Microsoft.Extensions.Caching.Hybrid;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
+using Microsoft.IdentityModel.Tokens;
 
 namespace Duende.AccessTokenManagement;
 
@@ -18,6 +19,11 @@ public class DistributedDPoPNonceStore : IDPoPNonceStore
     private readonly HybridCache _cache;
     private readonly ILogger<DistributedDPoPNonceStore> _logger;
 
+    private static readonly HybridCacheEntryOptions NonceStoreCacheOptions = new HybridCacheEntryOptions()
+    {
+        LocalCacheExpiration = TimeSpan.FromHours(1)
+    };
+
     /// <summary>
     /// ctor
     /// </summary>
@@ -56,16 +62,11 @@ public virtual async Task StoreNonceAsync(DPoPNonceContext context, string nonce
 
         var data = nonce;
 
-        var cacheExpiration = TimeSpan.FromHours(1);
-        var entryOptions = new HybridCacheEntryOptions()
-        {
-            Expiration = cacheExpiration
-        };
 
-        _logger.LogTrace("Caching DPoP nonce for URL: {url}, method: {method}. Expiration: {expiration}", context.Url, context.Method, cacheExpiration);
+        _logger.LogTrace("Caching DPoP nonce for URL: {url}, method: {method}. Expiration: {expiration}", context.Url, context.Method, NonceStoreCacheOptions.Expiration);
 
         var cacheKey = GenerateCacheKey(context);
-        await _cache.SetAsync(cacheKey, data, entryOptions, cancellationToken: cancellationToken).ConfigureAwait(false);
+        await _cache.SetAsync(cacheKey, data, NonceStoreCacheOptions, cancellationToken: cancellationToken).ConfigureAwait(false);
     }
 
 

access-token-management/src/AccessTokenManagement/HybridCacheExtMethods.cs

@@ -17,12 +17,11 @@ public static class HybridCacheExtMethods
                 | HybridCacheEntryFlags.DisableUnderlyingData
     };
 
-    public static async ValueTask<T?> GetOrDefaultAsync<T>(this HybridCache cache, string key)
+    public static async ValueTask<T?> GetOrDefaultAsync<T>(this HybridCache cache, string key, CancellationToken cancellationToken = default)
     {
         return await cache.GetOrCreateAsync<T?>(
             key,
             null!, // Don't return a value if it's not in the cache. Also, don't write it to the cache
-            GetOnlyEntryOptions
-        );
+            GetOnlyEntryOptions, cancellationToken: cancellationToken);
     }
 }

access-token-management/test/AccessTokenManagement.Tests/BackChannelClientTests.cs

@@ -5,50 +5,10 @@
 using System.Net;
 using System.Net.Http.Json;
 using Duende.IdentityModel.Client;
-using Microsoft.Extensions.Caching.Hybrid;
 using Microsoft.Extensions.DependencyInjection;
 using RichardSzalay.MockHttp;
 namespace Duende.AccessTokenManagement.Tests;
 
-public class HybridCacheTests
-{
-    [Fact]
-    public async Task Returning_null()
-    {
-        var services = new ServiceCollection()
-            .AddHybridCache()
-            .Services;
-
-        var cache = services.BuildServiceProvider()
-            .GetService<HybridCache>();
-
-        int count = 0;
-        object item;
-
-        try
-        {
-            item = await cache.GetOrCreateAsync<object>("key", (_) =>
-            {
-                count++;
-                throw new InvalidOperationException();
-                return ValueTask.FromResult<object>(null);
-            });
-        }
-        catch (InvalidOperationException)
-        {
-
-        }
-        item = await cache.GetOrCreateAsync<object>("key", (_) =>
-        {
-            count++;
-            return ValueTask.FromResult<object>(null);
-        });
-
-        item.ShouldBeNull();
-        count.ShouldBe(2);
-    }
-}
-
 public class BackChannelClientTests(ITestOutputHelper output)
 {
     [Fact]

access-token-management/test/AccessTokenManagement.Tests/HybridCacheExplorationTests.cs

@@ -0,0 +1,42 @@
+using Microsoft.Extensions.Caching.Hybrid;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Duende.AccessTokenManagement.Tests;
+
+public class HybridCacheExplorationTests
+{
+    [Fact]
+    public async Task Exception_is_not_written_to_cache()
+    {
+        var services = new ServiceCollection()
+            .AddHybridCache()
+            .Services;
+
+        var cache = services.BuildServiceProvider()
+            .GetRequiredService<HybridCache>();
+
+        int count = 0;
+        object item;
+
+        try
+        {
+            item = await cache.GetOrCreateAsync<object>("key", (_) =>
+            {
+                count++;
+                throw new InvalidOperationException();
+            });
+        }
+        catch (InvalidOperationException)
+        {
+
+        }
+        item = await cache.GetOrCreateAsync<object>("key", (_) =>
+        {
+            count++;
+            return ValueTask.FromResult<object>(null);
+        });
+
+        item.ShouldBeNull();
+        count.ShouldBe(2);
+    }
+}

identity-model-oidc-client/test/IdentityModel.OidcClient.Tests/DPoP/Framework/DPoP/DPoPServiceCollectionExtensions.cs

@@ -15,6 +15,8 @@ public static IServiceCollection ConfigureDPoPTokensForScheme(this IServiceColle
 
         services.AddTransient<DPoPJwtBearerEvents>();
         services.AddTransient<DPoPProofValidator>();
+        services.AddDistributedMemoryCache();
+
         services.AddTransient<IReplayCache, DefaultReplayCache>();
 
         services.AddSingleton<IPostConfigureOptions<JwtBearerOptions>>(new ConfigureJwtBearerOptions(scheme));