Decorate MediatR IRequestHandler using Microsoft DI

Question

I'm trying to decorate some of my MediatR's IRequestHandlers. To mark handlers I wanted to decorate, I created the interface ICommandHandler that inherits from IRequestHandler and ICommand that inherits from IRequest.

public interface ICommand : IRequest { }
public interface ICommandHandler<in TCommand> : IRequestHandler<TCommand>
    where TCommand : ICommand { }

Now I want every ICommandHandler to be decorated with the UnitOfWorkCommandHandlerDecorator.

internal class UnitOfWorkCommandHandlerDecorator<TCommand> : ICommandHandler<TCommand>
    where TCommand : ICommand
{
    private readonly ICommandHandler<TCommand> _decorated;
    private readonly IUnitOfWork _unitOfWork;

    public UnitOfWorkCommandHandlerDecorator(ICommandHandler<TCommand> decorated, IUnitOfWork unitOfWork)
    {
        _decorated = decorated;
        _unitOfWork = unitOfWork;
    }

    public async Task<Unit> Handle(TCommand request, CancellationToken cancellationToken)
    {
        var res = await _decorated.Handle(request, cancellationToken);
        await _unitOfWork.CommitAsync();
        return res;
    }
}

I tried registering a decorator with Scrutor's decorate method, however the decorator isn't invoked when a request is sent with IMediator.Send. That's probably because it's an ICommandHandler decorator rather than an IRequestHandler decorator.

services.Decorate(typeof(ICommandHandler<>), typeof(UnitOfWorkCommandHandlerDecorator<>));

Is it possible with Microsoft DI or do I need another DI container?

Answer 1

From what I read it is not possible, with Microsofts DI, to implement a decoration of an open generic on an open generic. But you can decorate a closed generic with an open generic. That leads to the idea that you scan your assembly for the generic, find the closed implementation and then decorate that.

services.Scan(scan =>
  scan.FromAssembliesOf(typeof(ICommandHandler<>))
    .AddClasses(classes => 
    classes.AssignableTo(typeof(ICommandHandler<>)).Where(_ => !_.IsGenericType))
        .AsImplementedInterfaces()
        .WithTransientLifetime());

I cannot take any credit for the code, that comes from the GitHub ticket: https://github.com/khellang/Scrutor/issues/68

Answer 2

As I expected, the problem was decorating ICommandHandler when MediatR requests IRequestHandler from Microsoft DI.

I found two approaches to making this work:

1. When MediatR requests an IRequestHandler, have Microsoft DI return a decorated ICommandHandler.

// Register all concrete classes implementing ICommandHandler as ICommandHandlers.
services.Scan(x => x.FromAssembliesOf(typeof(ICommand))
    .AddClasses(classes => classes.AssignableTo(typeof(ICommandHandler<>))
        .Where(t => !t.IsGenericType))
    .AsImplementedInterfaces(t => t.IsGenericType
        && t.GetGenericTypeDefinition() == typeof(ICommandHandler<>))
    .WithTransientLifetime());

// Decorate registered ICommandHandlers with open generic decorator.
services.Decorate(typeof(ICommandHandler<>), typeof(UnitOfWorkCommandHandlerDecorator<>));

// Get registered ICommandHandlers service descriptors from the container.
var registeredCommandHandlers = services
    .Where(x => x.ServiceType.IsGenericType
        && x.ServiceType.GetGenericTypeDefinition() == typeof(ICommandHandler<>))
    .ToArray();
foreach(var serviceDescriptor in registeredCommandHandlers)
{
    // Find the IRequestHandler type of registered ICommandHandler.
    var requestHandlerType = serviceDescriptor.ServiceType.GetInterfaces()
        .First(x => x.IsGenericType && x.GetGenericTypeDefinition() == typeof(IRequestHandler<,>));
    // Register a decorated ICommandHandler to be returned for IRequestHandler.
    services.Add(new ServiceDescriptor(requestHandlerType,
        sp => sp.GetRequiredService(serviceDescriptor.ServiceType), serviceDescriptor.Lifetime));
}

// Invoke AddMediatR method to register IRequestHandlers that are not ICommandHandlers
// and others required by MediatR classes.
services.AddMediatR(typeof(TestCommandHandler));

2. Change MediatR to request ICommandHandler rather than IRequestHandler.

// Register all concrete classes implementing ICommandHandler as ICommandHandlers.
services.Scan(x => x.FromAssembliesOf(typeof(ICommand))
    .AddClasses(classes => classes.AssignableTo(typeof(ICommandHandler<,>))
        .Where(t => !t.IsGenericType))
    .AsImplementedInterfaces(t => t.IsGenericType
        && t.GetGenericTypeDefinition() == typeof(ICommandHandler<,>))
    .WithTransientLifetime());

// Decorate registered ICommandHandlers with open generic decorator.
services.Decorate(typeof(ICommandHandler<,>), typeof(UnitOfWorkCommandHandlerDecorator<,>));

// Invoke AddMediatR method to register IRequestHandlers that are not ICommandHandlers
// and others required by MediatR classes.
services.AddMediatR(typeof(TestCommandHandler));

// Register ServiceFactory delegate, which MediatR uses to resolve all services.
services.AddTransient<ServiceFactory>(sp =>
{
    return t =>
    {
        if (t.IsGenericType
            && t.GetGenericTypeDefinition() == typeof(IRequestHandler<,>))
        {
            var genericArguments = t.GetGenericArguments();
            var commandHandlerType = typeof(ICommandHandler<,>).MakeGenericType(genericArguments);
            // Return ICommandHandler or IRequestHandler if it is not registered.
            return sp.GetService(commandHandlerType)
                ?? sp.GetRequiredService(t);
        }
        return sp.GetRequiredService(t);
    };
});

Personally, I like the first approach better; it is easier to understand, and it does not have as much overhead on runtime.