Support for .NET CancellationTokens

src/Orleans.CodeGenerator/AnalyzerReleases.Unshipped.md

@@ -1,3 +1,8 @@
-; Unshipped analyzer release
+; Unshipped analyzer release
 ; https://github.com/dotnet/roslyn-analyzers/blob/main/src/Microsoft.CodeAnalysis.Analyzers/ReleaseTrackingAnalyzers.Help.md
 
+### New Rules
+
+Rule ID | Category | Severity | Notes
+--------|----------|----------|--------------------
+ORLEANS0109 | Usage | Error | Method has multiple CancellationToken parameters

src/Orleans.CodeGenerator/CopierGenerator.cs

@@ -40,6 +40,11 @@ public ClassDeclarationSyntax GenerateCopier(
             var members = new List<ISerializableMember>();
             foreach (var member in type.Members)
             {
+                if (!member.IsCopyable)
+                {
+                    continue;
+                }
+
                 if (member is ISerializableMember serializable)
                 {
                     members.Add(serializable);
@@ -283,6 +288,11 @@ public void GetCopierFieldDescriptions(IEnumerable<IMemberDescription> members,
             var uniqueTypes = new HashSet<IMemberDescription>(MemberDescriptionTypeComparer.Default);
             foreach (var member in members)
             {
+                if (!member.IsCopyable)
+                {
+                    continue;
+                }
+
                 var t = member.Type;
 
                 if (LibraryTypes.IsShallowCopyable(t))

src/Orleans.CodeGenerator/Diagnostics/MultipleCancellationTokenParametersDiagnostic.cs

@@ -0,0 +1,16 @@
+using System.Linq;
+using Microsoft.CodeAnalysis;
+
+namespace Orleans.CodeGenerator.Diagnostics;
+
+public static class MultipleCancellationTokenParametersDiagnostic
+{
+    public const string DiagnosticId = "ORLEANS0109";
+    public const string Title = "Grain method has multiple parameters of type CancellationToken";
+    public const string MessageFormat = "The type {0} contains method {1} which has multiple CancellationToken parameters. Only a single CancellationToken parameter is supported.";
+    public const string Category = "Usage";
+
+    private static readonly DiagnosticDescriptor Rule = new DiagnosticDescriptor(DiagnosticId, Title, MessageFormat, Category, DiagnosticSeverity.Error, isEnabledByDefault: true);
+
+    internal static Diagnostic CreateDiagnostic(IMethodSymbol symbol) => Diagnostic.Create(Rule, symbol.Locations.First(), symbol.ContainingType.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat), symbol.Name);
+}

src/Orleans.CodeGenerator/InvokableGenerator.cs

@@ -7,6 +7,7 @@
 using System.Linq;
 using Orleans.CodeGenerator.Diagnostics;
 using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;
+using System.Linq.Expressions;
 
 namespace Orleans.CodeGenerator
 {
@@ -148,12 +149,14 @@ private ClassDeclarationSyntax GetClassDeclarationSyntax(
                     GenerateGetActivityName(method),
                     GenerateGetInterfaceType(method),
                     GenerateGetMethod(),
-                    GenerateSetTargetMethod(method, targetField),
+                    GenerateSetTargetMethod(method, targetField, fieldDescriptions),
                     GenerateGetTargetMethod(targetField),
                     GenerateDisposeMethod(fieldDescriptions, baseClassType),
                     GenerateGetArgumentMethod(method, fieldDescriptions),
                     GenerateSetArgumentMethod(method, fieldDescriptions),
-                    GenerateInvokeInnerMethod(method, fieldDescriptions, targetField));
+                    GenerateInvokeInnerMethod(method, fieldDescriptions, targetField),
+                    GenerateGetCancellationTokenMethod(method, fieldDescriptions),
+                    GenerateTryCancelMethod(method, fieldDescriptions));
 
             if (method.AllTypeParameters.Count > 0)
             {
@@ -182,7 +185,6 @@ private MemberDeclarationSyntax[] GenerateResponseTimeoutPropertyMembers(long va
                 .WithExpressionBody(ArrowExpressionClause(IdentifierName("_responseTimeoutValue")))
                 .WithSemicolonToken(Token(SyntaxKind.SemicolonToken))
                 .AddModifiers(Token(SyntaxKind.PublicKeyword), Token(SyntaxKind.OverrideKeyword));
-;
             return new MemberDeclarationSyntax[] { timespanField, responseTimeoutProperty };
         }
 
@@ -276,7 +278,8 @@ private INamedTypeSymbol GetBaseClassType(InvokableMethodDescription method)
 
         private MemberDeclarationSyntax GenerateSetTargetMethod(
             InvokableMethodDescription methodDescription,
-            TargetFieldDescription targetField)
+            TargetFieldDescription targetField,
+            List<InvokerFieldDescription> fieldDescriptions)
         {
             var holder = IdentifierName("holder");
             var holderParameter = holder.Identifier;
@@ -293,16 +296,29 @@ private MemberDeclarationSyntax GenerateSetTargetMethod(
                                     SingletonSeparatedList(containingInterface.ToTypeSyntax())))))
                 .WithArgumentList(ArgumentList());
 
-            var body =
-                AssignmentExpression(
-                    SyntaxKind.SimpleAssignmentExpression,
-                    IdentifierName(targetField.FieldName),
-                    getTarget);
-            return MethodDeclaration(PredefinedType(Token(SyntaxKind.VoidKeyword)), "SetTarget")
+            var member =
+            MethodDeclaration(PredefinedType(Token(SyntaxKind.VoidKeyword)), "SetTarget")
                 .WithParameterList(ParameterList(SingletonSeparatedList(Parameter(holderParameter).WithType(LibraryTypes.ITargetHolder.ToTypeSyntax()))))
-                .WithExpressionBody(ArrowExpressionClause(body))
-                .WithSemicolonToken(Token(SyntaxKind.SemicolonToken))
                 .WithModifiers(TokenList(Token(SyntaxKind.PublicKeyword), Token(SyntaxKind.OverrideKeyword)));
+
+            var assignmentExpression = AssignmentExpression(SyntaxKind.SimpleAssignmentExpression, IdentifierName(targetField.FieldName), getTarget);
+            if (!methodDescription.IsCancellable)
+            {
+                return member.WithExpressionBody(ArrowExpressionClause(assignmentExpression)).WithSemicolonToken(Token(SyntaxKind.SemicolonToken));
+            }
+            else
+            {
+                var ctsField = fieldDescriptions.First(f => f is CancellationTokenSourceFieldDescription);
+                var cancellationTokenType = LibraryTypes.CancellationToken.ToTypeSyntax();
+                var ctField = fieldDescriptions.First(f => SymbolEqualityComparer.Default.Equals(LibraryTypes.CancellationToken, f.FieldType));
+                return member.WithBody(Block(
+                    new List<StatementSyntax>()
+                    {
+                        ExpressionStatement(assignmentExpression),
+                        ExpressionStatement(AssignmentExpression(SyntaxKind.SimpleAssignmentExpression, IdentifierName(ctsField.FieldName), ImplicitObjectCreationExpression())),
+                        ExpressionStatement(AssignmentExpression(SyntaxKind.SimpleAssignmentExpression, IdentifierName(ctField.FieldName), IdentifierName(ctsField.FieldName).Member("Token")))
+                    }));
+            }
         }
 
         private MemberDeclarationSyntax GenerateGetTargetMethod(TargetFieldDescription targetField)
@@ -314,6 +330,47 @@ private MemberDeclarationSyntax GenerateGetTargetMethod(TargetFieldDescription t
                 .WithModifiers(TokenList(Token(SyntaxKind.PublicKeyword), Token(SyntaxKind.OverrideKeyword)));
         }
 
+        private MemberDeclarationSyntax GenerateGetCancellationTokenMethod(InvokableMethodDescription method, List<InvokerFieldDescription> fields)
+        {
+            if (!method.IsCancellable)
+            {
+                return null;
+            }
+
+            // Method to get the cancellationToken argument
+            // C#: CancellationToken GetCancellationToken() => <fieldName>
+            var cancellationTokenType = LibraryTypes.CancellationToken.ToTypeSyntax();
+            var cancellationTokenField = fields.First(f => SymbolEqualityComparer.Default.Equals(LibraryTypes.CancellationToken, f.FieldType));
+            var member = MethodDeclaration(cancellationTokenType, "GetCancellationToken")
+                 .WithExpressionBody(ArrowExpressionClause(cancellationTokenField.FieldName.ToIdentifierName()))
+                 .AddModifiers(Token(SyntaxKind.PublicKeyword), Token(SyntaxKind.OverrideKeyword))
+                .WithSemicolonToken(Token(SyntaxKind.SemicolonToken));
+            return member;
+        }
+
+        private MemberDeclarationSyntax GenerateTryCancelMethod(InvokableMethodDescription method, List<InvokerFieldDescription> fields)
+        {
+            if (!method.IsCancellable)
+            {
+                return null;
+            }
+
+            // Method to set the CancellationToken argument.
+            // C#:
+            // TryCancel()
+            // {
+            //   _cts.Cancel(false);
+            //   return true;
+            // }
+            var cancellationTokenField = fields.First(f => f is CancellationTokenSourceFieldDescription);
+            var member = MethodDeclaration(PredefinedType(Token(SyntaxKind.BoolKeyword)), "TryCancel")
+                .AddModifiers(Token(SyntaxKind.PublicKeyword), Token(SyntaxKind.OverrideKeyword))
+                .WithBody(Block(
+                            ExpressionStatement(InvocationExpression(IdentifierName(cancellationTokenField.FieldName).Member("Cancel")).WithArgumentList(ArgumentList(SeparatedList([Argument(LiteralExpression(SyntaxKind.FalseLiteralExpression))])))),
+                    ReturnStatement(LiteralExpression(SyntaxKind.TrueLiteralExpression))));
+            return member;
+        }
+
         private MemberDeclarationSyntax GenerateGetArgumentMethod(
             InvokableMethodDescription methodDescription,
             List<InvokerFieldDescription> fields)
@@ -496,6 +553,18 @@ private MemberDeclarationSyntax GenerateDisposeMethod(
             var body = new List<StatementSyntax>();
             foreach (var field in fields)
             {
+                if (field is CancellationTokenSourceFieldDescription ctsField)
+                {
+                    // C#
+                    // _cts?.Dispose();
+                    body.Add(
+                        ExpressionStatement(
+                            ConditionalAccessExpression(
+                                ctsField.FieldName.ToIdentifierName(),
+                                InvocationExpression(
+                                    MemberBindingExpression(IdentifierName("Dispose"))))));
+                }
+
                 if (field.IsInstanceField)
                 {
                     body.Add(
@@ -708,17 +777,23 @@ private ExpressionSyntax GetTypesArray(InvokableMethodDescription method, IEnume
         private List<InvokerFieldDescription> GetFieldDescriptions(InvokableMethodDescription method)
         {
             var fields = new List<InvokerFieldDescription>();
-
             uint fieldId = 0;
+
             foreach (var parameter in method.Method.Parameters)
             {
-                fields.Add(new MethodParameterFieldDescription(method.CodeGenerator, parameter, $"arg{fieldId}", fieldId, method.TypeParameterSubstitutions));
+                var isSerializable = !SymbolEqualityComparer.Default.Equals(LibraryTypes.CancellationToken, parameter.Type);
+                fields.Add(new MethodParameterFieldDescription(method.CodeGenerator, parameter, $"arg{fieldId}", fieldId, method.TypeParameterSubstitutions, isSerializable));
                 fieldId++;
             }
 
             fields.Add(new TargetFieldDescription(method.Method.ContainingType));
             fields.Add(new MethodInfoFieldDescription(LibraryTypes.MethodInfo, "MethodBackingField"));
 
+            if (method.IsCancellable)
+            {
+                fields.Add(new CancellationTokenSourceFieldDescription(LibraryTypes));
+            }
+
             return fields;
         }
 
@@ -738,8 +813,20 @@ protected InvokerFieldDescription(ITypeSymbol fieldType, string fieldName)
 
         internal sealed class TargetFieldDescription : InvokerFieldDescription
         {
-            public TargetFieldDescription(ITypeSymbol fieldType) : base(fieldType, "target") { }
+            public TargetFieldDescription(ITypeSymbol fieldType) : base(fieldType, "_target") { }
+
+            public override bool IsSerializable => false;
+            public override bool IsInstanceField => true;
+        }
 
+        internal sealed class CancellationTokenSourceFieldDescription(LibraryTypes libraryTypes) : InvokerFieldDescription(libraryTypes.CancellationTokenSource, "_cts")
+        {
+            public override bool IsSerializable => false;
+            public override bool IsInstanceField => true;
+        }
+
+        internal sealed class CancellationTokenFieldDescription(LibraryTypes libraryTypes) : InvokerFieldDescription(libraryTypes.CancellationToken, "_ct")
+        {
             public override bool IsSerializable => false;
             public override bool IsInstanceField => true;
         }
@@ -751,7 +838,8 @@ public MethodParameterFieldDescription(
                 IParameterSymbol parameter,
                 string fieldName,
                 uint fieldId,
-                Dictionary<ITypeParameterSymbol, string> typeParameterSubstitutions)
+                Dictionary<ITypeParameterSymbol, string> typeParameterSubstitutions,
+                bool isSerializable)
                 : base(parameter.Type, fieldName)
             {
                 TypeParameterSubstitutions = typeParameterSubstitutions;
@@ -770,6 +858,7 @@ public MethodParameterFieldDescription(
                 }
 
                 Symbol = parameter;
+                IsSerializable = isSerializable;
             }
 
             public CodeGenerator CodeGenerator { get; }
@@ -782,7 +871,8 @@ public MethodParameterFieldDescription(
             public INamedTypeSymbol ContainingType => Parameter.ContainingType;
             public TypeSyntax TypeSyntax { get; }
             public IParameterSymbol Parameter { get; }
-            public override bool IsSerializable => true;
+            public override bool IsSerializable { get; }
+            public bool IsCopyable => true;
             public override bool IsInstanceField => true;
 
             public string AssemblyName => Parameter.Type.ContainingAssembly.ToDisplayName();

src/Orleans.CodeGenerator/LibraryTypes.cs

@@ -77,11 +77,11 @@ private LibraryTypes(Compilation compilation, CodeGeneratorOptions options)
             _dateOnly = TypeOrDefault("System.DateOnly");
             _dateTimeOffset = Type("System.DateTimeOffset");
             _bitVector32 = Type("System.Collections.Specialized.BitVector32");
-            _guid = Type("System.Guid");
             _compareInfo = Type("System.Globalization.CompareInfo");
             _cultureInfo = Type("System.Globalization.CultureInfo");
             _version = Type("System.Version");
             _timeOnly = TypeOrDefault("System.TimeOnly");
+            Guid = Type("System.Guid");
             ICodecProvider = Type("Orleans.Serialization.Serializers.ICodecProvider");
             ValueSerializer = Type("Orleans.Serialization.Serializers.IValueSerializer`1");
             ValueTask = Type("System.Threading.Tasks.ValueTask");
@@ -153,7 +153,8 @@ private LibraryTypes(Compilation compilation, CodeGeneratorOptions options)
             TimeSpan = Type("System.TimeSpan");
             _ipAddress = Type("System.Net.IPAddress");
             _ipEndPoint = Type("System.Net.IPEndPoint");
-            _cancellationToken = Type("System.Threading.CancellationToken");
+            CancellationToken = Type("System.Threading.CancellationToken");
+            CancellationTokenSource = Type("System.Threading.CancellationTokenSource");
             _immutableContainerTypes = new[]
             {
                     compilation.GetSpecialType(SpecialType.System_Nullable_T),
@@ -259,13 +260,14 @@ INamedTypeSymbol Type(string metadataName)
         public INamedTypeSymbol SuppressReferenceTrackingAttribute { get; private set; }
         public INamedTypeSymbol OmitDefaultMemberValuesAttribute { get; private set; }
         public INamedTypeSymbol CopyContext { get; private set; }
+        public INamedTypeSymbol CancellationToken { get; private set; }
+        public INamedTypeSymbol CancellationTokenSource { get; }
+        public INamedTypeSymbol Guid { get; private set; }
         public Compilation Compilation { get; private set; }
         public INamedTypeSymbol TimeSpan { get; private set; }
         private INamedTypeSymbol _ipAddress;
         private INamedTypeSymbol _ipEndPoint;
-        private INamedTypeSymbol _cancellationToken;
         private INamedTypeSymbol[] _immutableContainerTypes;
-        private INamedTypeSymbol _guid;
         private INamedTypeSymbol _bitVector32;
         private INamedTypeSymbol _compareInfo;
         private INamedTypeSymbol _cultureInfo;
@@ -280,14 +282,14 @@ INamedTypeSymbol Type(string metadataName)
             _dateOnly,
             _timeOnly,
             _dateTimeOffset,
-            _guid,
+            Guid,
             _bitVector32,
             _compareInfo,
             _cultureInfo,
             _version,
             _ipAddress,
             _ipEndPoint,
-            _cancellationToken,
+            CancellationToken,
             Type,
             _uri,
             _uInt128,

src/Orleans.CodeGenerator/Model/FieldDescription.cs

@@ -37,6 +37,8 @@ public FieldDescription(uint fieldId, bool isPrimaryConstructorParameter, IField
         public string TypeName => Type.ToDisplayName();
         public string TypeNameIdentifier => Type.GetValidIdentifier();
         public bool IsPrimaryConstructorParameter { get; set; }
+        public bool IsSerializable => true;
+        public bool IsCopyable => true;
 
         public TypeSyntax GetTypeSyntax(ITypeSymbol typeSymbol) => typeSymbol.ToTypeSyntax();
     }

src/Orleans.CodeGenerator/Model/IMemberDescription.cs

@@ -17,6 +17,8 @@ internal interface IMemberDescription
         string TypeNameIdentifier { get; }
         TypeSyntax GetTypeSyntax(ITypeSymbol typeSymbol);
         bool IsPrimaryConstructorParameter { get; } 
+        bool IsSerializable { get; }
+        bool IsCopyable { get; }
     }
 
     internal sealed class MemberDescriptionTypeComparer : IEqualityComparer<IMemberDescription>

src/Orleans.CodeGenerator/Model/InvokableMethodDescription.cs

@@ -4,6 +4,7 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Globalization;
+using System.Linq;
 
 namespace Orleans.CodeGenerator
 {
@@ -206,6 +207,11 @@ static bool TryGetNamedArgument(ImmutableArray<KeyValuePair<string, TypedConstan
         /// </summary>
         public INamedTypeSymbol ContainingInterface { get; }
 
+        /// <summary>
+        /// Gets a value indicating whether this method is cancellable.
+        /// </summary>
+        public bool IsCancellable => Method.Parameters.Any(parameterSymbol => SymbolEqualityComparer.Default.Equals(CodeGenerator.LibraryTypes.CancellationToken, parameterSymbol.Type));
+
         public bool Equals(InvokableMethodDescription other) => Key.Equals(other.Key);
         public override bool Equals(object obj) => obj is InvokableMethodDescription imd && Equals(imd);
         public override int GetHashCode() => Key.GetHashCode();

src/Orleans.CodeGenerator/Model/PropertyDescription.cs

@@ -40,6 +40,8 @@ public PropertyDescription(uint fieldId, bool isPrimaryConstructorParameter, IPr
         public string TypeName => Type.ToDisplayName();
         public string TypeNameIdentifier => Type.GetValidIdentifier();
         public bool IsPrimaryConstructorParameter { get; set; }
+        public bool IsSerializable => true;
+        public bool IsCopyable => true;
 
         public TypeSyntax GetTypeSyntax(ITypeSymbol typeSymbol) => typeSymbol.ToTypeSyntax();
     }

src/Orleans.CodeGenerator/Model/ProxyMethodDescription.cs

@@ -121,7 +121,8 @@ public ConstructedGeneratedInvokableDescription(GeneratedInvokableDescription in
                         proxyMethodParameters[member.ParameterOrdinal],
                         member.FieldName,
                         member.FieldId,
-                        proxyMethod.TypeParameterSubstitutions));
+                        proxyMethod.TypeParameterSubstitutions,
+                        member.IsSerializable));
                 }
             }