Added promise cache interceptor. (#8190)

Author: michaelstaib

src/GreenDonut/src/GreenDonut.Abstractions/IPromiseCache.cs

@@ -1,3 +1,5 @@
+using System.Diagnostics.CodeAnalysis;
+
 namespace GreenDonut;
 
 /// <summary>
@@ -92,7 +94,7 @@ public interface IPromiseCache
     void Publish<T>(T value);
 
     /// <summary>
-    /// Publishes the values to the cache subscribers without adding it to the cache iself.
+    /// Publishes the values to the cache subscribers without adding it to the cache itself.
     /// This allows the subscribers to decide if they want to cache the values.
     /// </summary>
     /// <param name="values">

src/GreenDonut/src/GreenDonut.Abstractions/IPromiseCacheInterceptor.cs

@@ -0,0 +1,44 @@
+namespace GreenDonut;
+
+/// <summary>
+/// Allows to implement a second-level cache for the DataLoader promise cache.
+/// </summary>
+public interface IPromiseCacheInterceptor
+{
+    /// <summary>
+    /// Gets a task from the cache if a task with the specified <paramref name="key"/> already
+    /// exists; otherwise, the <paramref name="createPromise"/> factory is used to create a new
+    /// task and add it to the cache.
+    /// </summary>
+    /// <param name="key">A cache entry key.</param>
+    /// <param name="createPromise">A factory to create the new task.</param>
+    /// <typeparam name="T">The task type.</typeparam>
+    /// <returns>
+    /// Returns either the retrieved or new task from the cache.
+    /// </returns>
+    /// <exception cref="ArgumentNullException">
+    /// Throws if <paramref name="key"/> is <c>null</c>.
+    /// </exception>
+    /// <exception cref="ArgumentNullException">
+    /// Throws if <paramref name="createPromise"/> is <c>null</c>.
+    /// </exception>
+    Promise<T> GetOrAddPromise<T>(PromiseCacheKey key, Func<PromiseCacheKey, Promise<T>> createPromise);
+
+    /// <summary>
+    /// Tries to add a single task to the cache. It does nothing if the
+    /// task exists already.
+    /// </summary>
+    /// <param name="key">A cache entry key.</param>
+    /// <param name="promise">A task.</param>
+    /// <typeparam name="T">The task type.</typeparam>
+    /// <exception cref="ArgumentNullException">
+    /// Throws if <paramref name="key"/> is <c>null</c>.
+    /// </exception>
+    /// <exception cref="ArgumentNullException">
+    /// Throws if <paramref name="promise"/> is <c>null</c>.
+    /// </exception>
+    /// <returns>
+    /// A value indicating whether the add was successful.
+    /// </returns>
+    bool TryAdd<T>(PromiseCacheKey key, Promise<T> promise);
+}

src/GreenDonut/src/GreenDonut/DependencyInjection/DataLoaderServiceCollectionExtensions.cs

@@ -1,10 +1,10 @@
 using System.Collections.Concurrent;
-using System.Collections.Frozen;
 using GreenDonut;
 using GreenDonut.DependencyInjection;
 using Microsoft.Extensions.DependencyInjection.Extensions;
 using Microsoft.Extensions.ObjectPool;
 
+// ReSharper disable once CheckNamespace
 namespace Microsoft.Extensions.DependencyInjection;
 
 public static class DataLoaderServiceCollectionExtensions
@@ -66,7 +66,12 @@ public static IServiceCollection TryAddDataLoaderCore(
         services.TryAddScoped<IBatchScheduler, AutoBatchScheduler>();
 
         services.TryAddSingleton(sp => PromiseCachePool.Create(sp.GetRequiredService<ObjectPoolProvider>()));
-        services.TryAddScoped(sp => new PromiseCacheOwner(sp.GetRequiredService<ObjectPool<PromiseCache>>()));
+        services.TryAddScoped(sp =>
+        {
+            var pool = sp.GetRequiredService<ObjectPool<PromiseCache>>();
+            var interceptor = sp.GetService<IPromiseCacheInterceptor>();
+            return new PromiseCacheOwner(pool, interceptor);
+        });
 
         services.TryAddSingleton<IDataLoaderDiagnosticEvents>(
             sp =>

src/GreenDonut/src/GreenDonut/PromiseCache.cs

@@ -38,6 +38,8 @@ public PromiseCache(int size)
     /// <inheritdoc />
     public int Usage => _usage;
 
+    public IPromiseCacheInterceptor? Interceptor { get; set; }
+
     /// <inheritdoc />
     public Task<T> GetOrAddTask<T>(PromiseCacheKey key, Func<PromiseCacheKey, Promise<T>> createPromise)
     {
@@ -51,12 +53,15 @@ public Task<T> GetOrAddTask<T>(PromiseCacheKey key, Func<PromiseCacheKey, Promis
             throw new ArgumentNullException(nameof(createPromise));
         }
 
+        var interceptor = Interceptor;
         var read = true;
 
         var entry = _map.GetOrAdd(key, k =>
         {
             read = false;
-            return AddNewEntry(k, createPromise(k));
+            return interceptor is null
+                ? AddNewEntry(k, createPromise(k))
+                : new Entry(k, interceptor.GetOrAddPromise(k, createPromise));
         });
 
         if (read)
@@ -95,6 +100,11 @@ public bool TryAdd<T>(PromiseCacheKey key, Promise<T> promise)
             return AddNewEntry(k, promise);
         });
 
+        if(!read)
+        {
+            Interceptor?.TryAdd(key, promise);
+        }
+
         if (!promise.IsClone)
         {
             promise.OnComplete(NotifySubscribers, new CacheAndKey(this, key));
@@ -124,6 +134,11 @@ public bool TryAdd<T>(PromiseCacheKey key, Func<Promise<T>> createPromise)
             return AddNewEntry(k, createPromise());
         });
 
+        if(!read)
+        {
+            Interceptor?.TryAdd(key, (Promise<T>)entry.Value);
+        }
+
         var promise = (Promise<T>)entry.Value;
 
         if (!promise.IsClone)

src/GreenDonut/src/GreenDonut/PromiseCacheOwner.cs

@@ -15,19 +15,21 @@ public sealed class PromiseCacheOwner : IDisposable
     /// <summary>
     /// Rents a new cache from <see cref="PromiseCachePool.Shared"/>.
     /// </summary>
-    public PromiseCacheOwner()
+    public PromiseCacheOwner(IPromiseCacheInterceptor? interceptor = null)
     {
         _pool = PromiseCachePool.Shared;
         _cache = PromiseCachePool.Shared.Get();
+        _cache.Interceptor = interceptor;
     }
 
     /// <summary>
     /// Rents a new cache from the given <paramref name="pool"/>.
     /// </summary>
-    public PromiseCacheOwner(ObjectPool<PromiseCache> pool)
+    public PromiseCacheOwner(ObjectPool<PromiseCache> pool, IPromiseCacheInterceptor? interceptor = null)
     {
         _pool = pool ?? throw new ArgumentNullException(nameof(pool));
         _cache = pool.Get();
+        _cache.Interceptor = interceptor;
     }
 
     /// <summary>

src/GreenDonut/src/GreenDonut/PromiseCachePooledObjectPolicy.cs

@@ -8,6 +8,7 @@ internal sealed class PromiseCachePooledObjectPolicy(int size) : PooledObjectPol
 
     public override bool Return(PromiseCache obj)
     {
+        obj.Interceptor = null;
         obj.Clear();
         return true;
     }

src/GreenDonut/test/GreenDonut.Tests/PromiseCacheTests.cs

@@ -1,3 +1,5 @@
+using System.Collections.Concurrent;
+
 namespace GreenDonut;
 
 public class PromiseCacheTests
@@ -9,6 +11,7 @@ public void ConstructorNoException()
         var cacheSize = 1;
 
         // act
+        // ReSharper disable once ObjectCreationAsStatement
         void Verify() => new PromiseCache(cacheSize);
 
         // assert
@@ -240,4 +243,90 @@ public async Task GetOrAddTaskWhenNothingIsCached_IntegerKey()
         // assert
         Assert.Equal("Quox", await resolved);
     }
+
+    [Fact]
+    public async Task GetOrAddTask_When_SecondLevelEntry_Exists()
+    {
+        // arrange
+        var cacheSize = 10;
+        var secondLevel = new SecondLevelCache();
+        var cache = new PromiseCache(cacheSize) { Interceptor = secondLevel };
+        var key = new PromiseCacheKey("abc", "abc");
+
+        // act
+        var resolved = cache.GetOrAddTask(key, _ => new Promise<string>(Task.FromResult("Quox")));
+
+        // assert
+        Assert.Equal("def", await resolved);
+    }
+
+    [Fact]
+    public async Task GetOrAddTask_When_SecondLevelEntry_Not_Exists()
+    {
+        // arrange
+        var cacheSize = 10;
+        var secondLevel = new SecondLevelCache();
+        var cache = new PromiseCache(cacheSize) { Interceptor = secondLevel };
+        var key = new PromiseCacheKey("abc", "123");
+
+        // act
+        await cache.GetOrAddTask(key, _ => new Promise<string>(Task.FromResult("quox")));
+
+        // assert
+        var secondLevelEntry = (Task<string>)secondLevel.Cache[key];
+        Assert.Equal("quox", await secondLevelEntry);
+    }
+
+    [Fact]
+    public async Task TryAddTask_To_SecondLevelCache_1()
+    {
+        // arrange
+        var cacheSize = 10;
+        var secondLevel = new SecondLevelCache();
+        var cache = new PromiseCache(cacheSize) { Interceptor = secondLevel };
+        var key = new PromiseCacheKey("abc", "123");
+
+        // act
+        cache.TryAdd(key, () => new Promise<string>(Task.FromResult("quox")));
+
+        // assert
+        var secondLevelEntry = (Task<string>)secondLevel.Cache[key];
+        Assert.Equal("quox", await secondLevelEntry);
+    }
+
+    [Fact]
+    public async Task TryAddTask_To_SecondLevelCache_2()
+    {
+        // arrange
+        var cacheSize = 10;
+        var secondLevel = new SecondLevelCache();
+        var cache = new PromiseCache(cacheSize) { Interceptor = secondLevel };
+        var key = new PromiseCacheKey("abc", "123");
+
+        // act
+        cache.TryAdd(key, new Promise<string>(Task.FromResult("quox")));
+
+        // assert
+        var secondLevelEntry = (Task<string>)secondLevel.Cache[key];
+        Assert.Equal("quox", await secondLevelEntry);
+    }
+
+    public class SecondLevelCache : IPromiseCacheInterceptor
+    {
+        private readonly ConcurrentDictionary<PromiseCacheKey, Task> _cache = new()
+        {
+            [new PromiseCacheKey("abc", "abc")] = Task.FromResult("def")
+        };
+
+        public ConcurrentDictionary<PromiseCacheKey, Task> Cache => _cache;
+
+        public Promise<T> GetOrAddPromise<T>(PromiseCacheKey key, Func<PromiseCacheKey, Promise<T>> createPromise)
+        {
+            var entry = _cache.GetOrAdd(key, _ => createPromise(key).Task);
+            return new Promise<T>((Task<T>)entry);
+        }
+
+        public bool TryAdd<T>(PromiseCacheKey key, Promise<T> promise)
+            => _cache.TryAdd(key, promise.Task);
+    }
 }

src/HotChocolate/Data/test/Data.PostgreSQL.Tests/HotChocolate.Data.PostgreSQL.Tests.csproj

@@ -16,6 +16,7 @@
   </ItemGroup>
 
   <ItemGroup>
+    <PackageReference Include="Microsoft.Extensions.Caching.Memory" />
     <PackageReference Include="Npgsql.EntityFrameworkCore.PostgreSQL" />
     <PackageReference Include="Squadron.PostgreSql" />
     <PackageReference Include="System.IO.Pipelines" />