DI コンテナに登録済みのサービスを装飾する

調査目的で、ASP.NET Core MVC とかが DI コンテナに登録したサービスに処理を挟み込みたいことがある。そんな時は、下記のようなクラスと拡張メソッドを用意。

internal class Decorator<TService>
{
    public TService Instance { get; set; }

    public Decorator(TService instance)
    {
        Instance = instance;
    }
}

internal class Decorator<TService, TImpl> : Decorator<TService>
    where TImpl : class, TService
{
    public Decorator(TImpl instance) : base(instance)
    {

    }
}

internal class DisposableDecorator<TService> : Decorator<TService>, IDisposable
{
    public DisposableDecorator(TService instance) : base(instance)
    {
    }

    public void Dispose()
    {
        (Instance as IDisposable)?.Dispose();
    }
}

static class DecoratorServiceCollectionExtensions
{
    public static IServiceCollection AddTransientDecorator<TService, TImplementation>(this IServiceCollection services)
        where TService : class
        where TImplementation : class, TService
    {
        services.AddDecorator<TService>();
        services.AddTransient<TService, TImplementation>();
        return services;
    }

    public static IServiceCollection AddDecorator<TService>(this IServiceCollection services)
    {
        var registration = services.LastOrDefault(x => x.ServiceType == typeof(TService));

        // デコレート対象が登録されていなかったら NG
        if (registration == null)
        {
            throw new InvalidOperationException(
                $"{typeof(TService).Name} は登録されていません。");
        }

        // すでにデコレーターが登録されていたら NG
        if (services.Any(x => x.ServiceType == typeof(Decorator<TService>)))
        {
            throw new InvalidOperationException(
                $"Decorator<{typeof(TService).Name}> はすでに登録されています。");
        }

        services.Remove(registration);

        if (registration.ImplementationInstance != null)
        {
            var type = registration.ImplementationInstance.GetType();
            var innerType = typeof(Decorator<,>).MakeGenericType(typeof(TService), type);
            services.Add(new ServiceDescriptor(typeof(Decorator<TService>), innerType, ServiceLifetime.Transient));
            services.Add(new ServiceDescriptor(type, registration.ImplementationInstance));
        }
        else if (registration.ImplementationFactory != null)
        {
            services.Add(new ServiceDescriptor(typeof(Decorator<TService>), provider =>
            {
                return new DisposableDecorator<TService>((TService)registration.ImplementationFactory(provider));
            }, registration.Lifetime));
        }
        else
        {
            var type = registration.ImplementationType;
            var innerType = typeof(Decorator<,>).MakeGenericType(typeof(TService), registration.ImplementationType);
            services.Add(new ServiceDescriptor(typeof(Decorator<TService>), innerType, ServiceLifetime.Transient));
            services.Add(new ServiceDescriptor(type, type, registration.Lifetime));
        }

        return services;
    }
}

実際に、登録済みのサービスを装飾するクラスを用意。今回は ASP.NET Core MVC の IActionInvokerFactory を選んでみた。

internal class TelemetryActionInvokerFactory : IActionInvokerFactory
{
    readonly IActionInvokerFactory _inner;

    public TelemetryActionInvokerFactory (Decorator<IActionInvokerFactory> inner)
    {
        _inner = inner.Instance;
    }

    public IActionInvoker CreateInvoker(ActionContext actionContext)
    {
        var telemetryClient = actionContext.HttpContext.RequestServices.GetRequiredService<TelemetryClient>();
        var invoker = _inner.CreateInvoker(actionContext);
        return new TelemetryActionInvoker(invoker, telemetryClient);
    }

    class TelemetryActionInvoker : IActionInvoker
    {
        readonly IActionInvoker _inner;

        readonly TelemetryClient _telemetryClient;

        public TelemetryActionInvoker(
            IActionInvoker inner,
            TelemetryClient telemetryClient)
        {
            _inner = inner;
            _telemetryClient = telemetryClient;
        }

        public async Task InvokeAsync()
        {
            var telemetry = new DependencyTelemetry();
            var operation = _telemetryClient.StartOperation(telemetry);
            try
            {
                await _inner.InvokeAsync().ConfigureAwait(false);

                telemetry.Success = true;
                telemetry.Type = "ASP.NET Core MVC";
                telemetry.Name = name;
            }
            catch (Exception ex)
            {
                telemetry.Success = false;
                _telemetryClient.TrackException(ex);
                throw;
            }
            finally
            {
                _telemetryClient.StopOperation(operation);
            }
        }
    }
}

あとは、IServiceCollection にサービスを登録する箇所の最後の方で、デコレーターを登録してやればいい。

services.AddTransientDecorator<IActionInvokerFactory, TelemetryActionInvokerFactory>();

ASP.NET Core MVC を使って実装した REST API のボトルネック調査で、Application Insights にテレメトリを送信する処理を挿入しまくるときに、デコレーターが大活躍した。