Added DataLoader Scoping Improvements (#6943)

Author: michaelstaib

src/GreenDonut/src/Core/DependencyInjection/DataLoaderFactory.cs

@@ -0,0 +1,14 @@
+using System;
+
+namespace GreenDonut.DependencyInjection;
+
+/// <summary>
+/// Represents a factory that creates a DataLoader instance.
+/// </summary>
+public delegate IDataLoader DataLoaderFactory(IServiceProvider serviceProvider);
+
+/// <summary>
+/// Represents a factory that creates a DataLoader instance.
+/// </summary>
+public delegate T DataLoaderFactory<out T>(IServiceProvider serviceProvider) 
+    where T : IDataLoader;

src/HotChocolate/Core/src/Fetching/DataLoaderRegistration.cs renamed to src/GreenDonut/src/Core/DependencyInjection/DataLoaderRegistration.cs

@@ -1,31 +1,31 @@
 using System;
-using GreenDonut;
 using Microsoft.Extensions.DependencyInjection;
 
-namespace HotChocolate.Fetching;
+namespace GreenDonut.DependencyInjection;
 
 /// <summary>
 /// Represents a registration for a DataLoader.
 /// </summary>
 public sealed class DataLoaderRegistration
 {
-    private readonly Func<IServiceProvider, object> _factory;
+    private readonly DataLoaderFactory _factory;
 
-    public DataLoaderRegistration(Type instanceType) : this(instanceType, instanceType) { }
+    public DataLoaderRegistration(Type instanceType) 
+        : this(instanceType, instanceType) { }
 
     public DataLoaderRegistration(Type serviceType, Type instanceType)
     {
         ServiceType = serviceType;
         InstanceType = instanceType;
 
         var factory = ActivatorUtilities.CreateFactory(instanceType, []);
-        _factory = sp => factory.Invoke(sp, null);
+        _factory = sp => (IDataLoader)factory.Invoke(sp, null);
     }
 
-    public DataLoaderRegistration(Type serviceType, Func<IServiceProvider, object> factory)
+    public DataLoaderRegistration(Type serviceType, DataLoaderFactory factory)
         : this(serviceType, serviceType, factory) { }
 
-    public DataLoaderRegistration(Type serviceType, Type instanceType, Func<IServiceProvider, object> factory)
+    public DataLoaderRegistration(Type serviceType, Type instanceType, DataLoaderFactory factory)
     {
         ServiceType = serviceType;
         InstanceType = instanceType;
@@ -52,5 +52,5 @@ public DataLoaderRegistration(Type serviceType, Type instanceType, Func<IService
     /// Returns the new DataLoader instance.
     /// </returns>
     public IDataLoader CreateDataLoader(IServiceProvider services)
-        => (IDataLoader)_factory(services);
+        => _factory(services);
 }

src/GreenDonut/src/Core/DependencyInjection/DataLoaderServiceCollectionExtension.cs

@@ -0,0 +1,136 @@
+using System;
+using System.Collections.Concurrent;
+#if NET8_0_OR_GREATER
+using System.Collections.Frozen;
+#endif
+using System.Collections.Generic;
+using System.Linq;
+using GreenDonut;
+using GreenDonut.DependencyInjection;
+using Microsoft.Extensions.DependencyInjection.Extensions;
+using Microsoft.Extensions.ObjectPool;
+
+namespace Microsoft.Extensions.DependencyInjection;
+
+public static class DataLoaderServiceCollectionExtension
+{
+    public static IServiceCollection AddDataLoader<T>(
+        this IServiceCollection services)
+        where T : class, IDataLoader
+    {
+        services.TryAddDataLoaderCore();
+        services.AddSingleton(new DataLoaderRegistration(typeof(T)));
+        services.TryAddScoped<T>(sp => sp.GetDataLoader<T>());
+        return services;
+    }
+
+    public static IServiceCollection AddDataLoader<TService, TImplementation>(
+        this IServiceCollection services)
+        where TService : class, IDataLoader
+        where TImplementation : class, TService
+    {
+        services.TryAddDataLoaderCore();
+        services.AddSingleton(new DataLoaderRegistration(typeof(TService), typeof(TImplementation)));
+        services.TryAddScoped<TImplementation>(sp => sp.GetDataLoader<TImplementation>());
+        services.TryAddScoped<TService>(sp => sp.GetDataLoader<TService>());
+        return services;
+    }
+
+    public static IServiceCollection AddDataLoader<T>(
+        this IServiceCollection services,
+        Func<IServiceProvider, T> factory)
+        where T : class, IDataLoader
+    {
+        services.TryAddDataLoaderCore();
+        services.AddSingleton(new DataLoaderRegistration(typeof(T), sp => factory(sp)));
+        services.TryAddScoped<T>(sp => sp.GetDataLoader<T>());
+        return services;
+    }
+    
+    public static IServiceCollection TryAddDataLoaderCore(
+        this IServiceCollection services)
+    {
+         services.TryAddScoped<IDataLoaderScope, DefaultDataLoaderScope>();
+         services.TryAddScoped<IBatchScheduler, AutoBatchScheduler>();
+         
+         services.TryAddSingleton(sp => TaskCachePool.Create(sp.GetRequiredService<ObjectPoolProvider>()));
+         services.TryAddScoped(sp => new TaskCacheOwner(sp.GetRequiredService<ObjectPool<TaskCache>>()));
+         
+         services.TryAddSingleton<IDataLoaderDiagnosticEvents>(
+             sp =>
+             {
+                 var listeners = sp.GetServices<IDataLoaderDiagnosticEventListener>().ToArray();
+
+                 return listeners.Length switch
+                 {
+                     0 => new DataLoaderDiagnosticEventListener(),
+                     1 => listeners[0],
+                     _ => new AggregateDataLoaderDiagnosticEventListener(listeners),
+                 };
+             });
+
+         services.TryAddScoped(
+             sp =>
+             {
+                 var cacheOwner = sp.GetRequiredService<TaskCacheOwner>();
+
+                 return new DataLoaderOptions
+                 {
+                     Cache = cacheOwner.Cache,
+                     CancellationToken = cacheOwner.CancellationToken,
+                     DiagnosticEvents = sp.GetService<IDataLoaderDiagnosticEvents>(),
+                     MaxBatchSize = 1024,
+                 };
+             });
+         
+         return services;
+    }
+}
+
+file static class DataLoaderServiceProviderExtensions
+{
+    public static T GetDataLoader<T>(this IServiceProvider services) where T : IDataLoader
+        => services.GetRequiredService<IDataLoaderScope>().GetDataLoader<T>();
+}
+
+file sealed class DefaultDataLoaderScope(
+    IServiceProvider serviceProvider,
+#if NET8_0_OR_GREATER
+    FrozenDictionary<Type, DataLoaderRegistration> registrations)
+#else
+    Dictionary<Type, DataLoaderRegistration> registrations)
+#endif
+    : IDataLoaderScope
+{
+    private readonly ConcurrentDictionary<string, IDataLoader> _dataLoaders = new();
+
+
+    public T GetDataLoader<T>(DataLoaderFactory<T> createDataLoader, string? name = null) where T : IDataLoader
+    {
+        name ??= CreateKey<T>();
+
+        if (_dataLoaders.GetOrAdd(name, _ => createDataLoader(serviceProvider)) is T dataLoader)
+        {
+            return dataLoader;
+        }
+
+        throw new InvalidOperationException("A with the same name already exists.");
+    }
+
+    public T GetDataLoader<T>() where T : IDataLoader
+        => (T)_dataLoaders.GetOrAdd(CreateKey<T>(), _ => CreateDataLoader<T>());
+
+    private T CreateDataLoader<T>() where T : IDataLoader
+    {
+        if (registrations.TryGetValue(typeof(T), out var registration))
+        {
+            return (T)registration.CreateDataLoader(serviceProvider);
+        }
+
+        var adHocRegistration = new DataLoaderRegistration(typeof(T));
+        return (T)adHocRegistration.CreateDataLoader(serviceProvider);
+    }
+
+    private static string CreateKey<T>()
+        => typeof(T).FullName ?? typeof(T).Name;
+}

src/HotChocolate/Core/src/Fetching/IDataLoaderScope.cs renamed to src/GreenDonut/src/Core/DependencyInjection/IDataLoaderScope.cs

@@ -1,7 +1,4 @@
-using System;
-using GreenDonut;
-
-namespace HotChocolate.Fetching;
+namespace GreenDonut.DependencyInjection;
 
 /// <summary>
 /// The DataLoader scope provides access to the DataLoader bound to the current execution. 
@@ -23,7 +20,7 @@ public interface IDataLoaderScope
     /// <returns>
     /// Returns a <see cref="IDataLoader"/> instance from the current execution scope.
     /// </returns>
-    T GetDataLoader<T>(Func<T> createDataLoader, string? name = null) where T : IDataLoader;
+    T GetDataLoader<T>(DataLoaderFactory<T> createDataLoader, string? name = null) where T : IDataLoader;
 
     /// <summary>
     /// Gets a <see cref="IDataLoader"/> from the current execution scope; or, creates a new instance for this scope.

src/GreenDonut/src/Core/GreenDonut.csproj

@@ -9,20 +9,23 @@
 
   <ItemGroup Condition="'$(TargetFramework)' == 'net8.0'">
     <PackageReference Include="Microsoft.Extensions.ObjectPool" Version="8.0.0" />
+    <PackageReference Include="Microsoft.Extensions.DependencyInjection.Abstractions" Version="8.0.0" />
   </ItemGroup>
 
   <ItemGroup Condition="'$(TargetFramework)' == 'net7.0'">
     <PackageReference Include="Microsoft.Extensions.ObjectPool" Version="7.0.0" />
+    <PackageReference Include="Microsoft.Extensions.DependencyInjection.Abstractions" Version="7.0.0" />
   </ItemGroup>
 
   <ItemGroup Condition="'$(TargetFramework)' == 'net6.0'">
     <PackageReference Include="Microsoft.Extensions.ObjectPool" Version="6.0.0" />
+    <PackageReference Include="Microsoft.Extensions.DependencyInjection.Abstractions" Version="6.0.0" />
   </ItemGroup>
 
   <ItemGroup Condition="'$(TargetFramework)' == 'netstandard2.0'">
     <PackageReference Include="System.Diagnostics.DiagnosticSource" Version="6.0.0" />
-    <PackageReference Include="System.Threading.Tasks.Extensions" Version="4.5.0" />
     <PackageReference Include="Microsoft.Extensions.ObjectPool" Version="6.0.0" />
+    <PackageReference Include="Microsoft.Extensions.DependencyInjection.Abstractions" Version="6.0.0" />
   </ItemGroup>
 
 </Project>

src/GreenDonut/src/Core/Instrumentation/AggregateDataLoaderDiagnosticEventListener.cs

@@ -0,0 +1,76 @@
+using System;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+
+namespace GreenDonut;
+
+internal class AggregateDataLoaderDiagnosticEventListener(
+    IDataLoaderDiagnosticEventListener[] listeners)
+    : DataLoaderDiagnosticEventListener
+{
+    public override void ResolvedTaskFromCache(
+        IDataLoader dataLoader,
+        TaskCacheKey cacheKey,
+        Task task)
+    {
+        for (var i = 0; i < listeners.Length; i++)
+        {
+            listeners[i].ResolvedTaskFromCache(dataLoader, cacheKey, task);
+        }
+    }
+
+    public override IDisposable ExecuteBatch<TKey>(
+        IDataLoader dataLoader,
+        IReadOnlyList<TKey> keys)
+    {
+        var scopes = new IDisposable[listeners.Length];
+
+        for (var i = 0; i < listeners.Length; i++)
+        {
+            scopes[i] = listeners[i].ExecuteBatch(dataLoader, keys);
+        }
+
+        return new AggregateEventScope(scopes);
+    }
+
+    public override void BatchResults<TKey, TValue>(
+        IReadOnlyList<TKey> keys,
+        ReadOnlySpan<Result<TValue>> values)
+    {
+        for (var i = 0; i < listeners.Length; i++)
+        {
+            listeners[i].BatchResults(keys, values);
+        }
+    }
+
+    public override void BatchError<TKey>(
+        IReadOnlyList<TKey> keys,
+        Exception error)
+    {
+        for (var i = 0; i < listeners.Length; i++)
+        {
+            listeners[i].BatchError(keys, error);
+        }
+    }
+
+    public override void BatchItemError<TKey>(
+        TKey key,
+        Exception error)
+    {
+        for (var i = 0; i < listeners.Length; i++)
+        {
+            listeners[i].BatchItemError(key, error);
+        }
+    }
+
+    private sealed class AggregateEventScope(IDisposable[] scopes) : IDisposable
+    {
+        public void Dispose()
+        {
+            for (var i = 0; i < scopes.Length; i++)
+            {
+                scopes[i].Dispose();
+            }
+        }
+    }
+}