Generator adding logging to locks

Code.cs

using System;
using System.Threading;
using System.Threading.Tasks;
using DeadlockDemo;
using JetBrains.Annotations;

var test = new MyDeadlockDemoWithLogging();
    
Console.WriteLine("Running without deadlock...");
test.MyFirstMethod(false);
    
Console.WriteLine("\r\n\r\n\r\nRunning with deadlock...");
test.MyFirstMethod(true);

namespace DeadlockDemo
{
  partial class MyDeadlockDemoWithLogging
  {
    private readonly object _firstObj = new();
    private readonly object _secondObj = new();
    
    [UsedImplicitly]
    public void MyFirstMethodWithDeadlockDetection(bool initiateDeadlock)
    {
      lock (_firstObj)
      {
        Task.Delay(500).ContinueWith(_ => MySecondMethod());
        Thread.Sleep(initiateDeadlock ? 2000 : 0);
        
        lock (_secondObj)
          Console.WriteLine("Doing some logic under two locks in MyFirstMethod");
      }
      
      Thread.Sleep(1000); // wait for the second task
    }
    
    [UsedImplicitly]
    public void MySecondMethodWithDeadlockDetection()
    {
      lock (_secondObj)
      lock (_firstObj)
        Console.WriteLine("Doing some logic under two locks in MySecondMethod");
    }
  }
}

Output.txt

Running without deadlock...
ThreadId #1 Acquiring lock for '_firstObj' in 'MyDeadlockDemoWithLogging.MyFirstMethod'
ThreadId #1 Lock taken for  '_firstObj' in 'MyDeadlockDemoWithLogging.MyFirstMethod'
ThreadId #1 Acquiring lock for '_secondObj' in 'MyDeadlockDemoWithLogging.MyFirstMethod'
ThreadId #1 Lock taken for  '_secondObj' in 'MyDeadlockDemoWithLogging.MyFirstMethod'
Doing some logic under two locks in MyFirstMethod
ThreadId #1 Releasing lock for '_secondObj' in 'MyDeadlockDemoWithLogging.MyFirstMethod'
ThreadId #1 Releasing lock for '_firstObj' in 'MyDeadlockDemoWithLogging.MyFirstMethod'
ThreadId #5 Acquiring lock for '_secondObj' in 'MyDeadlockDemoWithLogging.MySecondMethod'
ThreadId #5 Lock taken for  '_secondObj' in 'MyDeadlockDemoWithLogging.MySecondMethod'
ThreadId #5 Acquiring lock for '_firstObj' in 'MyDeadlockDemoWithLogging.MySecondMethod'
ThreadId #5 Lock taken for  '_firstObj' in 'MyDeadlockDemoWithLogging.MySecondMethod'
Doing some logic under two locks in MySecondMethod
ThreadId #5 Releasing lock for '_firstObj' in 'MyDeadlockDemoWithLogging.MySecondMethod'
ThreadId #5 Releasing lock for '_secondObj' in 'MyDeadlockDemoWithLogging.MySecondMethod'



Running with deadlock...
ThreadId #1 Acquiring lock for '_firstObj' in 'MyDeadlockDemoWithLogging.MyFirstMethod'
ThreadId #1 Lock taken for  '_firstObj' in 'MyDeadlockDemoWithLogging.MyFirstMethod'
ThreadId #5 Acquiring lock for '_secondObj' in 'MyDeadlockDemoWithLogging.MySecondMethod'
ThreadId #5 Lock taken for  '_secondObj' in 'MyDeadlockDemoWithLogging.MySecondMethod'
ThreadId #5 Acquiring lock for '_firstObj' in 'MyDeadlockDemoWithLogging.MySecondMethod'
ThreadId #1 Acquiring lock for '_secondObj' in 'MyDeadlockDemoWithLogging.MyFirstMethod'

SourceGenerator.cs

using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Text;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace Generator
{
  [Generator]
  public class DeadlockDetectorGenerator : ISourceGenerator
  {
    public void Initialize(GeneratorInitializationContext context)
    {
      context.RegisterForSyntaxNotifications(() => new TemplateMethodsCollector());
    }
  
    public void Execute(GeneratorExecutionContext context)
    {
      var templatesCollector = (TemplateMethodsCollector) context.SyntaxReceiver;
      foreach (var classDeclarationSyntax in templatesCollector.Set)
      {
        var sourceBuilder = new StringBuilder();
        var syntaxTree = classDeclarationSyntax.SyntaxTree;
        var semanticModel = context.Compilation.GetSemanticModel(syntaxTree);
        var fileNodes = syntaxTree.GetRoot().DescendantNodesAndSelf();
        foreach (var usingDirectiveSyntax in fileNodes.OfType<UsingDirectiveSyntax>())
          sourceBuilder.AppendLine(usingDirectiveSyntax.ToString());
  
        sourceBuilder.AppendLine("using static System.Console;");
        sourceBuilder.AppendLine("using System.Threading;");
  
        var declaredClass = (ITypeSymbol) ModelExtensions.GetDeclaredSymbol(semanticModel, classDeclarationSyntax);
        sourceBuilder.AppendLine("namespace " + declaredClass.ContainingNamespace);
        sourceBuilder.AppendLine("{");
        sourceBuilder.AppendLine($"  partial class {declaredClass.Name}");
        sourceBuilder.AppendLine("  {");
        
        foreach (var methodDeclarationSyntax in classDeclarationSyntax.Members.OfType<MethodDeclarationSyntax>()
          .Where(x => x.Identifier.ToString().EndsWith("WithDeadlockDetection")))
        {
          var declaration = methodDeclarationSyntax;
          var newName = declaration.Identifier.ToString()[..^"WithDeadlockDetection".Length];
          var lockStatements = methodDeclarationSyntax.DescendantNodesAndSelf().OfType<LockStatementSyntax>().ToImmutableHashSet();
          declaration = declaration.ReplaceNodes(
            lockStatements.Reverse(),
            (_, lockStatementSyntax) =>
            {
              var lockDescription = $"'{lockStatementSyntax.Expression.ToString()}' in '{classDeclarationSyntax.Identifier.ToString()}.{newName}'";
              var acquiringStatement = CreateLogStatement(" Acquiring lock for " + lockDescription);
              var lockTakenStatement = CreateLogStatement(" Lock taken for  " + lockDescription);
              var lockWithUpdatedBody = lockStatementSyntax.WithStatement(SyntaxFactory.Block(lockTakenStatement, lockStatementSyntax.Statement));
              var releasingStatement = CreateLogStatement(" Releasing lock for " + lockDescription);
              return SyntaxFactory.Block(acquiringStatement, lockWithUpdatedBody, releasingStatement);
            });
  
          declaration = declaration.WithIdentifier(SyntaxFactory.Identifier(newName));
          declaration = declaration.NormalizeWhitespace();
          sourceBuilder.AppendLine(declaration.ToString());
        }
  
        sourceBuilder.AppendLine("  }");
        sourceBuilder.AppendLine("}");
  
        context.AddSource($"{classDeclarationSyntax.Identifier}.DeadlockDetection.cs", sourceBuilder.ToString());
      }
    }
  
    private static ExpressionStatementSyntax CreateLogStatement(string line)
    {
      var logStatement = SyntaxFactory.ExpressionStatement(
        SyntaxFactory.InvocationExpression(
            SyntaxFactory.IdentifierName("WriteLine"))
          .WithArgumentList(
            SyntaxFactory.ArgumentList(
              SyntaxFactory.SingletonSeparatedList<ArgumentSyntax>(
                SyntaxFactory.Argument(
                  SyntaxFactory.InterpolatedStringExpression(
                      SyntaxFactory.Token(SyntaxKind.InterpolatedStringStartToken))
                    .WithContents(
                      SyntaxFactory.List<InterpolatedStringContentSyntax>(
                        new InterpolatedStringContentSyntax[]
                        {
                          SyntaxFactory.InterpolatedStringText()
                            .WithTextToken(
                              SyntaxFactory.Token(
                                SyntaxFactory.TriviaList(),
                                SyntaxKind.InterpolatedStringTextToken,
                                "ThreadId #",
                                "ThreadId #",
                                SyntaxFactory.TriviaList())),
                          SyntaxFactory.Interpolation(
                            SyntaxFactory.MemberAccessExpression(
                              SyntaxKind.SimpleMemberAccessExpression,
                              SyntaxFactory.MemberAccessExpression(
                                SyntaxKind.SimpleMemberAccessExpression,
                                SyntaxFactory.IdentifierName("Thread"),
                                SyntaxFactory.IdentifierName("CurrentThread")),
                              SyntaxFactory.IdentifierName("ManagedThreadId"))),
                          SyntaxFactory.InterpolatedStringText()
                            .WithTextToken(
                              SyntaxFactory.Token(
                                SyntaxFactory.TriviaList(),
                                SyntaxKind.InterpolatedStringTextToken,
                                line,
                                line,
                                SyntaxFactory.TriviaList()))
                        })))))));
      return logStatement;
    }
  }

  public class TemplateMethodsCollector : ISyntaxReceiver
  {
    public void OnVisitSyntaxNode(SyntaxNode syntaxNode)
    {
      if (syntaxNode is ClassDeclarationSyntax declarationSyntax
          && declarationSyntax.Identifier.ToString().EndsWith("WithLogging"))
      {
        Set.Add(declarationSyntax);
      }
    }

    public HashSet<ClassDeclarationSyntax> Set { get; } = new();
  }
}