Improve GracefulStop() analyzer and code fix (#60)

* Improve GracefulStop analyzer and code fix

* Fix descriptor typo

src/Akka.Analyzers.Tests/Analyzers/AK1000/MustNotAwaitGracefulStopInsideReceiveAsyncAnalyzerSpec.cs

@@ -116,6 +116,34 @@ public Task ReceiveAnyAsync(Func<object, Task> func)
         return func("test");
     }
 }
+""",
+
+        // User defined `GracefulStop()` method, we're not responsible for this.
+"""
+using System;
+using Akka.Actor;
+using System.Threading.Tasks;
+
+public class MyActor: ReceiveActor
+{
+    public MyActor()
+    {
+        ReceiveAsync<string>(async str =>
+        {
+            await GracefulStop(TimeSpan.FromSeconds(3));
+        });
+
+        ReceiveAnyAsync(async _ =>
+        {
+            await GracefulStop(TimeSpan.FromSeconds(3));
+        });
+    }
+    
+    public Task GracefulStop(TimeSpan timeout)
+    {
+        return Task.CompletedTask;
+    }
+}
 """,
     };
 
@@ -191,7 +219,7 @@ public MyActor()
         ReceiveAnyAsync(async obj => await Context.Self.GracefulStop(TimeSpan.FromSeconds(3)));
     }
 }
-""", (9, 38, 9, 94)),            
+""", (9, 38, 9, 94)),
         };
 
     [Theory]

src/Akka.Analyzers/AK1000/MustNotAwaitGracefulStopInsideReceiveAsyncAnalyzer.cs

@@ -4,6 +4,7 @@
 //  </copyright>
 // -----------------------------------------------------------------------
 
+using System.Runtime.CompilerServices;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
@@ -24,36 +25,43 @@ public override void AnalyzeCompilation(CompilationStartAnalysisContext context,
         {
             var invocationExpr = (InvocationExpressionSyntax)ctx.Node;
 
-            if (invocationExpr.Expression is not MemberAccessExpressionSyntax memberAccess || memberAccess.Name.Identifier.Text != "GracefulStop") 
+            if(!IsGracefulStopInvocation(invocationExpr, ctx.SemanticModel, akkaContext))
                 return;
 
             // Check 1: GracefulStop() should not be awaited
             if (invocationExpr.Parent is not AwaitExpressionSyntax awaitExpression)
                 return;
 
             // Check 2: Ensure called within ReceiveAsync<T> lambda expression
-            if (!IsInsideReceiveAsyncLambda(invocationExpr, ctx.SemanticModel, akkaContext))
+            if (!invocationExpr.IsInsideReceiveAsyncLambda(ctx.SemanticModel, akkaContext))
                 return;
 
             var diagnostic = Diagnostic.Create(RuleDescriptors.Ak1002DoNotAwaitOnGracefulStop, awaitExpression.GetLocation());
             ctx.ReportDiagnostic(diagnostic);
         }, SyntaxKind.InvocationExpression);
     }
 
-    private static bool IsInsideReceiveAsyncLambda(SyntaxNode node, SemanticModel semanticModel, AkkaContext akkaContext)
-    {
-        // Traverse up the syntax tree to find the first lambda expression ancestor
-        var lambdaExpression = node.FirstAncestorOrSelf<LambdaExpressionSyntax>();
 
-        // Check if this lambda expression is an argument to an invocation expression
-        if (lambdaExpression?.Parent is not ArgumentSyntax { Parent: ArgumentListSyntax { Parent: InvocationExpressionSyntax invocationExpression } }) 
-            return false;
-
+    /// <summary>
+    /// Check if the invocation expression is a valid `GracefulStop` method invocation
+    /// </summary>
+    /// <param name="invocationExpression">The invocation expression being analyzed</param>
+    /// <param name="semanticModel">The semantic model</param>
+    /// <param name="akkaContext">The Akka context</param>
+    /// <returns>true if the invocation expression is a valid `GracefulStop` method</returns>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private static bool IsGracefulStopInvocation(
+        InvocationExpressionSyntax invocationExpression,
+        SemanticModel semanticModel,
+        AkkaContext akkaContext)
+    {
         // Get the method symbol from the invocation expression
         if (semanticModel.GetSymbolInfo(invocationExpression).Symbol is not IMethodSymbol methodSymbol)
             return false;
 
-        // Check if the method name is 'ReceiveAsync' and it is defined inside the ReceiveActor class
-        return methodSymbol.Name is "ReceiveAsync" or "ReceiveAnyAsync" && SymbolEqualityComparer.Default.Equals(methodSymbol.ContainingType, akkaContext.AkkaCore.ReceiveActorType);
+        // Check if the method name is 'GracefulStop', that it is an extension method,
+        // and it is defined inside the GracefulStopSupport static class
+        return methodSymbol is { Name: "GracefulStop", IsExtensionMethod: true }
+               && SymbolEqualityComparer.Default.Equals(methodSymbol.ContainingType, akkaContext.AkkaCore.GracefulStopSupportType);
     }
 }

src/Akka.Analyzers/Utility/AkkaCoreContext.cs

@@ -26,6 +26,7 @@ public interface IAkkaCoreContext
 
     public INamedTypeSymbol? IndirectActorProducerType { get; }
     public INamedTypeSymbol? ReceiveActorType { get; }
+    public INamedTypeSymbol? GracefulStopSupportType { get; }
 }
 
 /// <summary>
@@ -49,6 +50,8 @@ private EmptyCoreContext()
     public INamedTypeSymbol? IndirectActorProducerType => null;
 
     public INamedTypeSymbol? ReceiveActorType => null;
+
+    public INamedTypeSymbol? GracefulStopSupportType => null;
 }
 
 /// <summary>
@@ -66,6 +69,7 @@ public class AkkaCoreContext : IAkkaCoreContext
     private readonly Lazy<INamedTypeSymbol?> _lazyActorContextType;
     private readonly Lazy<INamedTypeSymbol?> _lazyIIndirectActorProducerType;
     private readonly Lazy<INamedTypeSymbol?> _lazyReceiveActorType;
+    private readonly Lazy<INamedTypeSymbol?> _lazyGracefulStopSupport;
 
     private AkkaCoreContext(Compilation compilation, Version version)
     {
@@ -76,6 +80,7 @@ private AkkaCoreContext(Compilation compilation, Version version)
         _lazyActorContextType = new Lazy<INamedTypeSymbol?>(() => TypeSymbolFactory.ActorContext(compilation));
         _lazyIIndirectActorProducerType = new Lazy<INamedTypeSymbol?>(() => TypeSymbolFactory.IndirectActorProducer(compilation));
         _lazyReceiveActorType = new Lazy<INamedTypeSymbol?>(() => TypeSymbolFactory.ReceiveActor(compilation));
+        _lazyGracefulStopSupport = new Lazy<INamedTypeSymbol?>(() => TypeSymbolFactory.GracefulStopSupport(compilation));
     }
 
     /// <inheritdoc />
@@ -90,6 +95,8 @@ private AkkaCoreContext(Compilation compilation, Version version)
     public INamedTypeSymbol? IndirectActorProducerType => _lazyIIndirectActorProducerType.Value;
 
     public INamedTypeSymbol? ReceiveActorType => _lazyReceiveActorType.Value;
+
+    public INamedTypeSymbol? GracefulStopSupportType => _lazyGracefulStopSupport.Value;
 
     public static AkkaCoreContext? Get(
         Compilation compilation,

src/Akka.Analyzers/Utility/CodeAnalysisExtensions.cs

@@ -4,7 +4,9 @@
 // </copyright>
 // -----------------------------------------------------------------------
 
+using System.Runtime.CompilerServices;
 using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
 
 namespace Akka.Analyzers;
 
@@ -41,4 +43,47 @@ public static bool IsActorBaseSubclass(this INamedTypeSymbol typeSymbol, AkkaCon
 
         return false;
     }
+
+    /// <summary>
+    /// Check if a syntax node is within a lambda expression that is an argument for either 
+    /// `ReceiveAsync` or `ReceiveAnyAsync` method invocation 
+    /// </summary>
+    /// <param name="node">The syntax node being analyzed</param>
+    /// <param name="semanticModel">The semantic model</param>
+    /// <param name="akkaContext">The Akka context</param>
+    /// <returns>true if the syntax node is inside a valid `ReceiveAsync` or `ReceiveAnyAsync` method</returns>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static bool IsInsideReceiveAsyncLambda(
+        this SyntaxNode node,
+        SemanticModel semanticModel,
+        AkkaContext akkaContext)
+    {
+        // Traverse up the syntax tree to find the first lambda expression ancestor
+        var lambdaExpression = node.FirstAncestorOrSelf<LambdaExpressionSyntax>();
+
+        // Check if this lambda expression is an argument to an invocation expression
+        if (lambdaExpression?.Parent is not ArgumentSyntax { Parent: ArgumentListSyntax { Parent: InvocationExpressionSyntax invocationExpression } }) 
+            return false;
+
+        return invocationExpression.IsReceiveAsyncInvocation(semanticModel, akkaContext);
+    }
+
+    /// <summary>
+    /// Check if the invocation expression is a valid `ReceiveAsync` or `ReceiveAnyAsync` method invocation
+    /// </summary>
+    /// <param name="invocationExpression">The invocation expression being analyzed</param>
+    /// <param name="semanticModel">The semantic model</param>
+    /// <param name="akkaContext">The Akka context</param>
+    /// <returns>true if the invocation expression is a valid `ReceiveAsync` or `ReceiveAnyAsync` method</returns>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static bool IsReceiveAsyncInvocation(this InvocationExpressionSyntax invocationExpression, SemanticModel semanticModel, AkkaContext akkaContext)
+    {
+        // Get the method symbol from the invocation expression
+        if (semanticModel.GetSymbolInfo(invocationExpression).Symbol is not IMethodSymbol methodSymbol)
+            return false;
+
+        // Check if the method name is `ReceiveAsync` or `ReceiveAnyAsync` and it is defined inside the ReceiveActor class
+        return methodSymbol.Name is "ReceiveAsync" or "ReceiveAnyAsync" 
+               && SymbolEqualityComparer.Default.Equals(methodSymbol.ContainingType, akkaContext.AkkaCore.ReceiveActorType);
+    }
 }

src/Akka.Analyzers/Utility/RuleDescriptors.cs

@@ -40,7 +40,7 @@ private static DiagnosticDescriptor Rule(
         category: AnalysisCategory.ActorDesign, 
         defaultSeverity: DiagnosticSeverity.Error,
         messageFormat: "Do not await on `Self.GracefulStop()` inside `ReceiveAsync()` because this will lead into " +
-                       "a deadlock inside the `ReceiveAsync()` and the actor will never receive the `PoisonPill` message sent by `GracefuLStop` while it's `await`-ing.");
+                       "a deadlock inside the `ReceiveAsync()` and the actor will never receive the `PoisonPill` message sent by `GracefulStop` while it's `await`-ing.");
 
     #endregion
 

src/Akka.Analyzers/Utility/TypeSymbolFactory.cs

@@ -93,4 +93,8 @@ public static class TypeSymbolFactory
         return Guard.AssertIsNotNull(compilation)
             .GetTypeByMetadataName("Akka.Actor.ReceiveActor");
     }
+
+    public static INamedTypeSymbol? GracefulStopSupport(Compilation compilation)
+        => Guard.AssertIsNotNull(compilation)
+            .GetTypeByMetadataName("Akka.Actor.GracefulStopSupport");
 }