Resolver scoping middleware (MediatR) for HotChocolate

HC_ResolverScopedService

A middleware to provide resolvers with their own, scoped instance, of a service. Example here is IMediator, but this could be made generic.

HCv10_Usage.cs

// Note, this is only for v10, v11's usage is in the xml comment of the attribute file

[UseResolverScopedMediator]
public async Task<string> GetAccountsAsync(
    IResolverContext context,
    [ScopedState] IMediator mediator)
    => await mediator.Send(new SomeRequest());
    
// Note: Here mediator is scoped, and when the handler is instantiated, all its dependencies are generated in the same scope,
// including DbContext instances, so no multi-thread usage :) 

HCv10_UseResolverScopedMediatorAttribute.cs

using System;
using System.Collections.Generic;
using System.Reflection;
using HotChocolate.Types;
using HotChocolate.Types.Descriptors;
using MediatR;
using Microsoft.Extensions.DependencyInjection;

namespace MyCompany.GraphQL
{
    /// <summary>
    /// This attribute allows us to tap into the field resolution middleware
    /// so that we can dynamically create an <see cref="IServiceScope"/> just for
    /// this field resolver method, and create an <see cref="IMediator"/> instance inside it.
    ///
    /// This is used to solve a problem with using EF Core's DbContextPool, which scopes
    /// DbContext instances per http request and will cause HotChocolate resolvers running
    /// in parallel to execute on the same DbContext, which isn't allowed.
    ///
    /// By resolving our <see cref="IMediator"/> instance inside an isolated scope,
    /// we ensure that if it depends itself on a DbContext that it will be unique in that scope
    /// and there won't be multiple threads using it simultaneously.
    /// </summary>
    [AttributeUsage(AttributeTargets.Method)]
    public class UseResolverScopedMediatorAttribute : ObjectFieldDescriptorAttribute
    {
        private static readonly string _injectedArgumentName = "mediator";

        private static readonly HashSet<string> _localSchemaNames = new HashSet<string>
        {
            SchemaNames.Content,
            SchemaNames.ShopTyre,
            SchemaNames.Task
        };

        public override void OnConfigure(
            IDescriptorContext descriptorContext,
            IObjectFieldDescriptor descriptor,
            MemberInfo member)
        {
            descriptor.Use(next => async context =>
            {
                // Temp workaround for: https://github.com/ChilliCream/hotchocolate/issues/2246
                using var scope = _localSchemaNames.Contains(context.Schema.Name)
                    ? context.Service<Microsoft.AspNetCore.Http.IHttpContextAccessor>().HttpContext.RequestServices.CreateScope()
                    : context.Service<IServiceProvider>().CreateScope();
                var mediator = scope.ServiceProvider.GetRequiredService<IMediator>();
                context.ModifyScopedContext(c => c.SetItem(_injectedArgumentName, mediator));

                await next(context);
            });
        }
    }
}

HCv11_UseResolverScopedMediatorAttribute.cs

using System;
using System.Reflection;
using HotChocolate;
using HotChocolate.Types;
using HotChocolate.Types.Descriptors;
using MediatR;
using Microsoft.Extensions.DependencyInjection;

namespace AutoGuru.GraphQL
{
    /// <summary>
    /// This attribute allows us to tap into the field resolution middleware
    /// so that we can dynamically create an <see cref="IServiceScope"/> just for
    /// this field resolver method, and create an <see cref="IMediator"/> instance inside it.
    ///
    /// This is used to solve a problem with using EF Core's DbContextPool, which scopes
    /// DbContext instances per http request and will cause HotChocolate resolvers running
    /// in parallel to execute on the same DbContext, which isn't allowed.
    ///
    /// By resolving our <see cref="IMediator"/> instance inside an isolated scope,
    /// we ensure that if it depends itself on a DbContext that it will be unique in that scope
    /// and there won't be multiple threads using it simultaneously.
    /// </summary>
    /// <example>
    ///     [UseResolverScopedMediator]
    ///     public async Task&lt;string&gt; GetAccountsAsync(
    ///         IResolverContext context,
    ///         [ScopedService] IMediator mediator)
    ///         =&gt; await mediator.Send(new SomeRequest());
    /// </example>
    /// <remarks>
    /// Similar to <c>HotChocolate.Types.EntityFrameworkObjectFieldDescriptorExtensions</c>.
    /// https://github.com/ChilliCream/hotchocolate/blob/fa9423ea2cb1fca6fe0fbd79822b4e06d5341897/src/HotChocolate/Data/src/EntityFramework/Extensions/EntityFrameworkObjectFieldDescriptorExtensions.cs#L26
    /// </remarks>
    [AttributeUsage(AttributeTargets.Method)]
    public class UseResolverScopedMediatorAttribute : ObjectFieldDescriptorAttribute
    {
        public override void OnConfigure(
            IDescriptorContext descriptorContext,
            IObjectFieldDescriptor descriptor,
            MemberInfo member)
        {
            var scopedServiceName = typeof(IMediator).FullName ?? typeof(IMediator).Name;
            
            descriptor.Use(next => async context =>
            {
                using var scope = context.Service<IServiceProvider>().CreateScope();

                try
                {
                    var mediator = scope.ServiceProvider.GetRequiredService<IMediator>();
                    context.SetLocalValue(scopedServiceName, mediator);

                    await next(context);
                }
                finally
                {
                    context.RemoveLocalValue(scopedServiceName);
                }
            });
        }
    }
}

HCv11_UseResolverScopedMediatorAttributeUsageTests.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using HotChocolate;
using HotChocolate.Types;
using MediatR;
using Shouldly;
using Xunit;

namespace MyCompany.GraphQL
{
    public abstract class ScopedMediatorTestsBase
    {
        protected abstract Type MarkerTypeInAssemblyToTest { get; }

        [Fact]
        protected void GraphResolverMethodsUseScopedMediator()
        {
            var extendObjectAttributeType = typeof(ExtendObjectTypeAttribute);

            var scopedServiceAttributeType = typeof(ScopedServiceAttribute);
            var scopedServiceExample = $"[{nameof(ScopedServiceAttribute).Replace("Attribute", "")}]";

            var useResolverScopedMediatorAttributeType = typeof(UseResolverScopedMediatorAttribute);
            var useResolverScopedMediatorExample = $"[{nameof(UseResolverScopedMediatorAttribute).Replace("Attribute", "")}]";

            var mediatorType = typeof(IMediator);

            var assemblies = new Assembly[] { typeof(Image).Assembly, MarkerTypeInAssemblyToTest.Assembly };
            var graphTypeExtensions = assemblies
                .SelectMany(a => a.GetTypes())
                .Where(t => t.CustomAttributes.Any(cad => cad.AttributeType == extendObjectAttributeType))
                .ToArray();

            var failures = new List<string>();

            foreach (var graphTypeExtension in graphTypeExtensions)
            {
                var methods = graphTypeExtension.GetMethods(BindingFlags.Instance | BindingFlags.Public | BindingFlags.DeclaredOnly);
                foreach (var method in methods)
                {
                    var hasOneUseResolverScopedMediatorAttr = method.GetCustomAttributes(useResolverScopedMediatorAttributeType).Count() == 1;

                    var parameters = method.GetParameters();
                    foreach (var parameter in parameters)
                    {
                        if (parameter.ParameterType == mediatorType)
                        {
                            var hasOneScopedServiceAttr = parameter.GetCustomAttributes(scopedServiceAttributeType).Count() == 1;
                            if (!hasOneUseResolverScopedMediatorAttr || !hasOneScopedServiceAttr)
                            {
                                var msg = $"Resolver method {method.Name} in class {graphTypeExtension.FullName} " +
                                    $"injects IMediator, but doesn't do it scoped. If an operation is fired with it" +
                                    $"that uses a handler that injects a DbContext, this could cause an issue with DbContext usage" +
                                    $"across multiple threads during parallel resolver execution. You need to:\r\n" +
                                    (hasOneUseResolverScopedMediatorAttr ? "" : $" - Annotate the method with {useResolverScopedMediatorExample}\r\n") +
                                    (hasOneScopedServiceAttr ? "" : $" - Annotate the IMediator parameter with {scopedServiceExample}\r\n");

                                failures.Add(msg);
                            }
                        }
                    }
                }
            }

            // Assert
            failures.ShouldBeEmpty(string.Join("\r\n ", failures));
        }
    }
}