Removing magic from auto-registration using source generators

At some point in your development career you will encounter a situation where you need to repetitively register objects/types. This is often the case for registering types to the Dependency Injection container. This become tedious, so some library authors have created solutions to make this easier. The solution usually is to specify an assembly (or a type from an assembly). The library will use reflection to 'scan' all types in that assembly, and register them 'automagically' to DI based on implemented interfaces. Some examples:

  • MassTransit, using it to automatically register consumers:
services.AddMassTransit(x =>
    {
        // Registers all consumers that exist in this assembly.
        x.AddConsumers(Assembly.GetExecutingAssembly());

        ...
    });
  • Mediatr, using it to automatically register the Mediatr types (IRequestHandler<>, INotificationHandler<>, ...):
services.AddMediatR(Assembly.GetExecutingAssembly());
  • Scrutor, providing a .Scan method to register arbitrary types to DI:
collection.Scan(scan => scan
    .FromAssemblyOf<ITransientService>()
        .AddClasses(classes => classes.AssignableTo<ITransientService>())
            .AsImplementedInterfaces()
            .WithTransientLifetime()
);
  • Rolling your own solution, scanning assemblies.

This last option requires an example scenario.

Let's say you have an IRegistrationType, and you want to register all types to DI that implement it:

public interface IRegistrationType
{
    public string Id { get; }
}

public class FirstType : IRegistrationType
{
    public string Id => "First";
}

public class SecondType : IRegistrationType
{
    public string Id => "Second";
}

We have a registry that then injects the registered types, and provides a LogTypes method to simply log the registered types:

public class RegistrationTypeRegistry(IEnumerable<IRegistrationType> types, ILogger<RegistrationTypeRegistry> logger)
{
    public void LogTypes()
    {
        logger.LogInformation($"The following types are registered:");
        foreach (var type in types)
        {
            logger.LogInformation($"Id: {type.Id}");
        }
    }
}

To implement assembly scanning to automatically register the types, you can create the following extension method:

public static class ServiceCollectionExtensions
{
    public static void RegisterRegistrationTypes(this IServiceCollection services)
    {
        var registrationTypes = 
            // Get this assembly
            typeof(ServiceCollectionExtensions).Assembly

            // Take all the types
            .GetTypes()

            // Take the ones that implement IRegistrationType, and are not abstract (this excludes the interface itself)
            .Where(x => typeof(IRegistrationType).IsAssignableFrom(x) && !x.IsAbstract);

        // For each of the types, register them to DI.
        foreach (var registrationType in registrationTypes)
        {
            services.AddTransient(typeof(IRegistrationType), registrationType);
        }
    }
}

And then invoke it from the Program.cs.

var builder = Host.CreateApplicationBuilder(args);

// Set up DI
builder.Services.AddSingleton<RegistrationTypeRegistry>();
builder.Services.RegisterRegistrationTypes();

// Build the app
var app = builder.Build();

// Test by invoking LogTypes.
var registry = app.Services.GetRequiredService<RegistrationTypeRegistry>();
registry.LogTypes();

And sure enough, running the program shows the types were registered:

info: AutoRegistration.SourceGenerator.Sample.RegistrationTypeRegistry[0]
      The following types are registered:
info: AutoRegistration.SourceGenerator.Sample.RegistrationTypeRegistry[0]
      Id: First
info: AutoRegistration.SourceGenerator.Sample.RegistrationTypeRegistry[0]
      Id: Second

You might be tempted to use this, or the methods provided by the aforementioned libraries, when you have a large solution requiring many types to be registered to DI. I am arguing however that, especially with larger solutions, it is a pitfall to use this kind of auto-registration. Please don't do this!

The main reason why I emphasize this is because of discoverability. Suppose you are a new team member on a large solution. You are debugging something related to FirstType, and to do so, you want to know how it is used. From the FirstType, you would normally perform a 'find all references' (usually SHIFT + F12) in your favorite IDE. When you do so, you are greeted with:

No references

In addition, the type name is greyed out by the IDE, indicating it is not used. However, it is used, but only at runtime by assembly scanning!

Another reason to avoid this approach is transparency. Transparency regarding which types are registered, in what order, and with which lifetime. For the a simple 'roll your own' solution, it may still look clear. But requirements change. We might want to register types with different lifetimes. Or exclude a certain type for some reason. Should we register only public types, or also internal types? Rules to configure this, such as Scrutor allows, can quickly become complex and unclear in larger solutions.

Lastly, if you plan to use AOT, which requires trimming your application code to only what is used, you cannot use this at all, because the trimmer does not know at compile time that FirstType is used, so it will be trimmed from the built application.

Performance can also be considered a downside, but often these extra clock cycles are not significant enough.

So how to register the types instead?

Solution 1: Register types manually

Really? Yes!

They are not that many extra keystrokes when typing! In fact, it is much more important for the reader to easily understand how the code works. An old quote that still holds true:

The ratio of time spent reading versus writing is well over 10 to 1. - Robert C. Martin (Uncle Bob)

By simply registering the types yourself manually:

  • The IDE shows the usages of types, improving discoverability
  • It is clear which types are registered, and how, maintaining transparency.
  • The code is AOT-compatible.

Sure, with large solutions you will end up with a lot of types to register, which may make it harder to read. You might also argue this could cause merge conflicts when multiple team members append registrations with their changes.

I would advocate for placing registrations in the assembly/folder/feature folder where the types resides, as extension methods, like:

namespace MyProject.FeatureA;

static class Services 
{
    public static void AddFeatureA(this IServiceCollection services)
    {
        services.AddSingleton<IRegistrationType, FeatureAType>();
    }
}

And call the registration methods in the Program.cs:

var builder = Host.CreateApplicationBuilder(args);

builder.Services.AddFeatureA();
builder.Services.AddFeatureB();
builder.Services.AddFeatureC();
...

var app = builder.Build();
...

This reduces the likelihood of merge conflicts, is easier to follow, and has the added benefit of improving the code cohesion.

Of course, this is the boring answer! If you still want to use auto-registration, there is another way:

Solution 2: Autoregistration using Source generators

A way to achieve discoverable, transparent and AOT-compatible auto-registration is by using Source Generators. This way we can scan the project at compile-time for types to be registered, and emit code to register those types.

Many articles have been written already about Source Generators and Roslyn Analyzers in general, so I will quickly go over this. The final code can be found here.

A new .NET Standard project is created, which is used by the C# compiler (Roslyn). Inside the project we register a source generator by annotating it with [Generator] (ironically using a reflection-based registration mechanism😉):

[Generator]
public class RegisterTypeGenerator : IIncrementalGenerator
{
    
    private const string ITypeRegistrationNamespace = "AutoRegistration.SourceGenerator.Sample";
    private const string ExtensionMethodNamespace = "Generators";
    private const string AttributeNamespace = "Generators";
    private const string AttributeName = "RegisterTypeAttribute";

    private const string AttributeSourceCode = $@"
using System;

namespace {AttributeNamespace};

[AttributeUsage(AttributeTargets.Class)]
class {AttributeName} : Attribute
{{
}}";
    ...

We first set up some constants. One constant of note is the source code for the RegisterTypeAttribute. We will use this attribute to generate DI registration code for every type annotated with it. For this example we use a marker attribute, but you could also simply register all types implementing IRegistrationType. We will emit the attribute in the source code, using:

    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        // Add the marker attribute to the compilation.
        context.RegisterPostInitializationOutput(ctx => ctx.AddSource(
            $"{AttributeName}.g.cs",
            SourceText.From(AttributeSourceCode, Encoding.UTF8)));
    }

Next up we need to detect the types we would like to register. For this we subscribe to any TypeDeclarationSyntax the compiler encounters during compilation. The TryGetRegistration then tests if the type declaration contains the RegisterType attribute. This is done using the 'semantic model', the stage in compilation where we know about the meaning of the code. In this case, it allows us to get the full type name of the applied RegisterType attribute.

     public void Initialize(IncrementalGeneratorInitializationContext context)
     {
         // Add the marker attribute to the compilation.
         context.RegisterPostInitializationOutput(ctx => ctx.AddSource(
             $"{AttributeName}.g.cs",
             SourceText.From(AttributeSourceCode, Encoding.UTF8)));
 
+        // Filter classes annotated with the [RegisterType] attribute.
+        var provider = context.SyntaxProvider
+            .CreateSyntaxProvider(
+                (s, _) => s is TypeDeclarationSyntax,
+                (ctx, _) => TryGetRegistration(ctx))
+            .Where(t => t is not null)
+            .Select((t, _) => t!);
     }

+    private static TypeDeclarationSyntax? TryGetRegistration(GeneratorSyntaxContext context)
+    {
+        var typeDeclarationSyntax = (TypeDeclarationSyntax) context.Node;
+        
+        var generatorAttribute = Helpers.FindAttribute(
+            typeDeclarationSyntax.AttributeLists, 
+            context.SemanticModel, 
+            $"{AttributeNamespace}.{AttributeName}");
+        
+        return generatorAttribute is not null ? typeDeclarationSyntax : null;
+    }
}

With the types that need to be registered, we can emit an extension method to register these:

     public void Initialize(IncrementalGeneratorInitializationContext context)
     {
         ...
 
         // Filter classes annotated with the [RegisterType] attribute.
         var provider = ...
 
+        // Generate the source code.
+        context.RegisterSourceOutput(context.CompilationProvider.Combine (provider.Collect()),
+            ((ctx, t) => GenerateCode(ctx, t.Left, t.Right)));
     }

     ...

+    private void GenerateCode(
+       SourceProductionContext context, 
+       Compilation compilation, 
+       ImmutableArray<TypeDeclarationSyntax> typeDeclarations)
+    {
+        // Go through all filtered class declarations.
+        var sb = new StringBuilder();
+        sb.AppendLine($@"using System;
+using System.Collections.Generic;
+using Microsoft.Extensions.DependencyInjection;
+using {ITypeRegistrationNamespace};
+
+namespace {ExtensionMethodNamespace};
+
+internal static class ServiceCollectionExtensions
+{{
+    internal static void RegisterTypes(this IServiceCollection services)
+    {{");
+        
+        foreach (var typeDeclaration in typeDeclarations)
+        {
+            var semanticModel = compilation.GetSemanticModel(typeDeclaration.SyntaxTree);
+
+            if (semanticModel.GetDeclaredSymbol(typeDeclaration) is not INamedTypeSymbol classSymbol)
+                continue;
+
+            
+            sb.AppendLine($"        services.AddSingleton<IRegistrationType, {classSymbol}>();");
+        }
+
+        sb.AppendLine("    }");
+        sb.AppendLine("}");
+
+        // Add the source code to the compilation.
+        context.AddSource($"ServiceCollectionExtensions.g.cs", SourceText.From(sb.ToString(), Encoding.UTF8));
+    }

We use the semantic model again to ensure we register the types using their full type names.

And with this, the implementation of the source generator is done! To test it, we use unit tests using Verify to perform snapshot testing. Essentially, the source generator is ran against different pieces of code, and the full generated source is compared to the expected code. See the unit test project for details.

Using the source generator

Let's use the source generators in our example scenario. I added a special project reference to use the source generator.

<ProjectReference Include="..\AutoRegistration.SourceGenerator\AutoRegistration.SourceGenerator.csproj" OutputItemType="Analyzer" ReferenceOutputAssembly="false"/>

We adjust the types we want to register by adding the [RegisterType] attribute:

+[RegisterType]
 public class FirstType : IRegistrationType
 {
     public string Id => "First";
 }

+[RegisterType]
 public class SecondType : IRegistrationType
 {
     public string Id => "Second";
 }

And we need to use the source generated extension method instead of our assembly scanning method:

 var builder = Host.CreateApplicationBuilder(args);
 
 // Set up DI
 builder.Services.AddSingleton<RegistrationTypeRegistry>();
-builder.Services.RegisterRegistrationTypes();
+builder.Services.RegisterTypes();

 // Build the app
 var app = builder.Build();
 
 // Test by invoking LogTypes.
 var registry = app.Services.GetRequiredService<RegistrationTypeRegistry>();
 registry.LogTypes();

When we now inspect FirstType we can already see (this time using Visual Studio) CodeLens found references and the type is not greyed out:

Usages

And when we find all the references, sure enough, we now see:

References

Navigating to the reference reveals the full generated code, providing insights in all types being registered and how they are registered:

internal static class ServiceCollectionExtensions
{
    internal static void RegisterTypes(this IServiceCollection services)
    {
        services.AddSingleton<IRegistrationType, AutoRegistration.SourceGenerator.Sample.RegisteredTypes.SecondType>();
        services.AddSingleton<IRegistrationType, AutoRegistration.SourceGenerator.Sample.RegisteredTypes.FirstType>();
    }
}

This solution offers all the benefits of auto-registration without the drawbacks of assembly scanning.

Conclusion

Source Generators are a great to reduce magic introduced by reflection, improving the readability of the code.

Additionally by using Source Generators, code can become AOT-compatible. This is something you see being used by Microsoft too to support AOT with ASP.NET Core. They have created source generators for various features that normally require reflection, like JSON serialization, configuration binding and Minimal API's.

I plan to create a blog post at some point about the challenges I encountered while adding AOT-compatibility to an existing app. Stay tuned for that!

← Back to the blog