MediatR: how to use Decorators to add retry policies

Hi All! Today we’ll see an interesting technique to add retry policies to MediatR. It can actually be used also for other types of policies (fallback, circuit breaker, and so on), but we’ll focusing on retries to keep things simple.

As you might have guessed, this “magic trick” involves the use of the Decorator Pattern. We talked already about it in another article so I’m not going to spend time on it. Let’s jump into the code!

public class RetryDecorator<TNotification> : MediatR.INotificationHandler<TNotification>
        where TNotification : MediatR.INotification
{
	private readonly INotificationHandler<TNotification> _inner;
	private readonly Polly.IAsyncPolicy _retryPolicy;

	public RetryDecorator(MediatR.INotificationHandler<TNotification> inner)
	{
		_inner = inner; 
		_retryPolicy = Polly.Policy.Handle<ArgumentOutOfRangeException>()
			.WaitAndRetryAsync(3,
				i => TimeSpan.FromSeconds(i));
	}

	public Task Handle(TNotification notification, CancellationToken cancellationToken)
	{
		return _retryPolicy.ExecuteAsync(() => _inner.Handle(notification, cancellationToken));
	}
}

I showed this class already in another article of my Event Sourcing series, but without going too much into the details. Today we’ll see how we can register it in our system using Dependency Injection.

This class is decorating an instance of INotificationHandler. In SuperSafeBank, I use Notification Handlers to react to Integration Events and update the Query Models.

However, sometimes an update might fail, maybe because other dependencies are (still) not available or there’s a network issue with the DB server or whatever.

Now, MediatR has already built-in support for decorators using Behaviors, but as of today ( v8 ) it doesn’t cover Notification Handlers .

In those cases, we can make use of Polly and add a policy wrapping the “unstable” code. But how can we decorate all our Handlers?

Most of the IoC libraries available have already commodity functions to register decorators. But to be honest, I like the built-in .NET Core IoC Container, it covers all the basic scenarios and lifetimes. Moreover, if you’re in need of something more exotic, most of the time you’d have to review the design.

Unfortunately though, Decorators are not part of the library. But we can use another interesting NuGet: Scrutor. It’s a tiny library, created specifically for this: “Assembly scanning and decoration extensions for Microsoft.Extensions.DependencyInjection“. Perfect for us.

So, the first thing is to look for all our handlers and register them:

public void ConfigureServices(IServiceCollection services) {
	services.Scan(scan => {
		scan.FromAssembliesOf(typeof(Startup))
			.RegisterHandlers(typeof(INotificationHandler<>));
	});
}

Here we use the .Scan() method to go through the assembly containing a specific type and retrieving all the classes implementing INotificationHandler<>. For the sake of the example, I’m assuming that our handlers are in the same assembly as the Startup class.


We have also to take into account that RetryDecorator implements INotificationHandler<> as well, and we don’t want any cyclic registration.

We can handle this with .RegisterHandlers() :

public static class IImplementationTypeSelectorExtensions
{
	public static IImplementationTypeSelector RegisterHandlers(this IImplementationTypeSelector selector, Type type)
	{
		return selector.AddClasses(c =>
				c.AssignableTo(type)
					.Where(t => t != typeof(RetryDecorator<>))
			)
			.UsingRegistrationStrategy(RegistrationStrategy.Append)
			.AsImplementedInterfaces()
			.WithScopedLifetime();
	}
}

Let’s see what’s happening here. AddClasses() goes through the classes picked by our selector and adds them to the IServicesCollection. But with a twist: we only want classes that are assignable to a specific type (in this case INotificationHandler<> ) and we skip the current class if it’s RetryDecorator.

The rest of the code makes sure that the selected classes are appended to the container, without replacing existing ones. We also want them to be registered and discoverable using their implemented interfaces. And finally, we want them to use the Scoped lifetime.

We have one last small step to make: registering the decorator itself 😀

public void ConfigureServices(IServiceCollection services){
    services.Decorate(typeof(INotificationHandler<>), typeof(RetryDecorator<>));
}

Piece of cake!

MediatR: how to use Decorators to add retry policies
Scroll to top