DecoratorパターンをSource Generatorで作成したら便利だった

前提

.NETの汎用ホストでOpenTelemetryを作成した際に、OpenTelemetryのTracingを計測するためには、汎用ホストの場合はASP.NETと違って手動でActivityをスタートさせる必要がありました。

Activity? activity = MyActivitySource.StartActivity("mycommand");

全てのユースケースに手動で追加するのは大変だと思ったのでSource Generatorを使用してユースケースをActivityでデコレートできないかと思い実装してみたため、投稿します。

具体例

試した例は以下に置いています。

MyOpenTelemetryDotnetSample

以下のユースケースが定義されていたとして

IDoA.cs

public interface IDoA
{
    void DoitA(string something);
}

IDoB.cs

public interface IDoB
{
    void DoitB(string something1, int something2);
}

Activityでデコレートしようと以下のような実装をまずは書いてみました。
デコレート元の委譲はMicrosoft.Extensions.DependencyInjectionの名前解決(FromKeyedServices)で行います。(.NET 8から追加されました)

IDoAのデコレータ

public record DoADecorator(ActivitySource activitySource, [FromKeyedServices("DoABase")]IDoA DoABase) : IDoA
{
    public void DoitA(string something)
    {
        Activity? activity = activitySource.StartActivity("DoA.DoitA");

        try
        {
            DoABase.DoitA(something);
            activity?.SetStatus(ActivityStatusCode.Ok);
        }
        catch (Exception ex)
        {
            activity?.SetStatus(ActivityStatusCode.Error, ex.Message);
            throw;
        }
        finally
        {
            activity?.Dispose();
        }
    }
}

IDoBのデコレータ

public record DoBDecorator(ActivitySource activitySource, [FromKeyedServices("DoBBase")]IDoB DoBBase) : IDoB
{
    public void DoitB(string something1, int something2)
    {
        Activity? activity = activitySource.StartActivity("DoB.DoitB");

        try
        {
            DoBBase.DoitB(something1, something2);
            activity?.SetStatus(ActivityStatusCode.Ok);
        }
        catch (Exception ex)
        {
            activity?.SetStatus(ActivityStatusCode.Error, ex.Message);
            throw;
        }
        finally
        {
            activity?.Dispose();
        }
    }
}

これらを見比べると
・クラス名
・引数
・インターフェース名
・委譲されているインスタンスを呼ぶ部分

等々の違いだけで実動作部分は同じにすることがわかりました。

この情報を元に作成したジェネレーターが以下です。

UsecaseDecoratorClassGenerator.cs

[Generator(LanguageNames.CSharp)]
public class UsecaseDecoratorClassGenerator : IIncrementalGenerator
{
    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        // ①
        var source = context.SyntaxProvider.ForAttributeWithMetadataName(
                    "usecase.GenerateDecoratorAttribute",
                    (node, token) => true,
                    (context, token) => context);

        context.RegisterSourceOutput(source, action);
    }

    private void action(SourceProductionContext context, GeneratorAttributeSyntaxContext source)
    {
        // ③
        var typeSymbol = (INamedTypeSymbol)source.TargetSymbol;
        var typeNode = (TypeDeclarationSyntax)source.TargetNode;

        var isymbol = typeSymbol.Interfaces.First();
        var methodName = isymbol.MemberNames.First();
        MethodDeclarationSyntax method = typeNode.Members.OfType<MethodDeclarationSyntax>().First(x => x.Identifier.NormalizeWhitespace().Text == methodName);

        // ②
        var code = $$"""
            // <auto-generated/>
            #nullable enable

            using System.Diagnostics;
            using Microsoft.Extensions.DependencyInjection;
            using {{isymbol.ContainingNamespace.Name}};

            namespace usecase;

            public record {{typeSymbol.Name}}Decorator(ActivitySource activitySource, [FromKeyedServices("{{typeSymbol.Name}}Base")]{{isymbol.Name}} {{typeSymbol.Name}}Base) : {{isymbol.Name}}
            {
                {{String.Join(" ", method.Modifiers.Select(x => x.Text))}} {{method.ReturnType}} {{method.Identifier}}({{method.ParameterList.Parameters}})
                {
                    Activity? activity = activitySource.StartActivity("{{typeSymbol.Name}}.{{method.Identifier}}");

                    try
                    {
                        {{typeSymbol.Name}}Base.{{method.Identifier}}({{String.Join(",", method.ParameterList.Parameters.Select(x => x.Identifier.ValueText))}});
                        activity?.SetStatus(ActivityStatusCode.Ok);
                    }
                    catch (Exception ex)
                    {
                        activity?.SetStatus(ActivityStatusCode.Error, ex.Message);
                        throw;
                    }
                    finally
                    {
                        activity?.Dispose();
                    }
                }
            }
            """;

        context.AddSource($"{typeSymbol.Name}Decorator.cs", code);
    }
}

①IIncrementalGeneratorで実装すると特定の属性だけを引っかけることが容易でした。
引っかかった要素分actionが実行されます。
ひとつのクラス(インターフェースetc)を元にひとつファイルを作成する場合には便利です。

context.CompilationProviderをsourceに指定すれば、ISourceGeneratorで実装するのと同じことが出来るのでジェネレータ作成時は常にIIncrementalGeneratorを使用するで良いと思いました。

②べた書きするコード部分はC# 11 以降で使用できる補間された生文字列リテラルが便利でした。具体的なコードをまず貼り付けて、変化する部分だけを差し替える感じです。

③変化する部分を抽出しています。デバッグで動かしながらでないと欲しい情報を取り出すのは難しい気がしました。

生成した結果は上のデコレータの内容と同じです。

また、DIコンテナへの登録は以下のように行います。

登録例

builder.ConfigureServices((HostBuilderContext ctxt, IServiceCollection services) =>
{
services.AddKeyedTransient<IDoA, DoA>("DoABase");
services.AddTransient<IDoA, DoADecorator>();
}

これもジェネレーターで作成してみました。ServiceCollectionExtensionGenerator.cs

まとめ

ソースジェネレーター自体はあっさり出来てびっくりしました。
活用方法はアイディア次第かと思います。
とりあえず処理の前後に定型文を挟むくらいはソースジェネレーターで簡単に出来ることがわかりました。
ただデバッグ環境がないと正確に作成するのは難しいと思いました。
普段vscodeを使っていますが今回はvsを使用しました。
誰かの参考になれば幸いです。

参考

2022年のC# (Incremental) Source Generator開発手法
Keyed service dependency injection container support