Filter cancellation exceptions in generator driver

src/Compilers/CSharp/Test/Semantic/SourceGeneration/SyntaxAwareGeneratorTests.cs

@@ -1960,6 +1960,41 @@ class C { }
             Assert.True(generatorCancelled);
         }
 
+        [Fact]
+        public void Syntax_Receiver_Cancellation_During_Visit()
+        {
+            var source = @"
+class C 
+{
+    int Property { get; set; }
+
+    void Function()
+    {
+        var x = 5;
+        x += 4;
+    }
+}
+";
+            var parseOptions = TestOptions.Regular;
+            Compilation compilation = CreateCompilation(source, options: TestOptions.DebugDll, parseOptions: parseOptions);
+            compilation.VerifyDiagnostics();
+
+            Assert.Single(compilation.SyntaxTrees);
+
+            var testGenerator = new CallbackGenerator(
+                onInit: (i) => i.RegisterForSyntaxNotifications(() => new TestSyntaxReceiver(tag: 0, callback: (a) => { if (a is AssignmentExpressionSyntax) { throw new OperationCanceledException("Simulated cancellation from external source"); } })),
+                onExecute: (e) => { e.AddSource("test", SourceText.From("public class D{}", Encoding.UTF8)); }
+                );
+
+            GeneratorDriver driver = CSharpGeneratorDriver.Create(new[] { testGenerator }, parseOptions: parseOptions);
+            driver = driver.RunGenerators(compilation, CancellationToken.None);
+            var results = driver.GetRunResult();
+
+            Assert.Single(results.Results);
+            Assert.IsType<OperationCanceledException>(results.Results[0].Exception);
+            Assert.Equal("Simulated cancellation from external source", results.Results[0].Exception!.Message);
+        }
+
         private class TestReceiverBase<T>
         {
             private readonly Action<T>? _callback;

src/Compilers/Core/Portable/InternalUtilities/ExceptionUtilities.cs

@@ -4,6 +4,7 @@
 
 using System;
 using System.Diagnostics;
+using System.Threading;
 
 namespace Roslyn.Utilities
 {
@@ -27,5 +28,14 @@ internal static Exception Unreachable
         {
             get { return new InvalidOperationException("This program location is thought to be unreachable."); }
         }
+
+        /// <summary>
+        /// Determine if an exception was an <see cref="OperationCanceledException"/>, and that the provided token caused the cancellation.
+        /// </summary>
+        /// <param name="exception">The exception to test.</param>
+        /// <param name="cancellationToken">Checked to see if the provided token was cancelled.</param>
+        /// <returns><see langword="true"/> if the exception was an <see cref="OperationCanceledException" /> and the token was canceled.</returns>
+        internal static bool IsCurrentOperationBeingCancelled(Exception exception, CancellationToken cancellationToken)
+            => exception is OperationCanceledException && cancellationToken.IsCancellationRequested;
     }
 }

src/Compilers/Core/Portable/InternalUtilities/FatalError.cs

@@ -7,6 +7,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Reflection;
 using System.Threading;
+using Roslyn.Utilities;
 
 #if NET20
 // Some APIs referenced by documentation comments are not available on .NET Framework 2.0.
@@ -91,9 +92,6 @@ public static void CopyHandlerTo(Assembly assembly)
 
 #endif
 
-        private static bool IsCurrentOperationBeingCancelled(Exception exception, CancellationToken cancellationToken)
-            => exception is OperationCanceledException && cancellationToken.IsCancellationRequested;
-
         /// <summary>
         /// Use in an exception filter to report an error without catching the exception.
         /// The error is reported by calling <see cref="Handler"/>.
@@ -143,7 +141,7 @@ public static bool ReportAndPropagateUnlessCanceled(Exception exception, ErrorSe
         [DebuggerHidden]
         public static bool ReportAndPropagateUnlessCanceled(Exception exception, CancellationToken contextCancellationToken, ErrorSeverity severity = ErrorSeverity.Uncategorized)
         {
-            if (IsCurrentOperationBeingCancelled(exception, contextCancellationToken))
+            if (ExceptionUtilities.IsCurrentOperationBeingCancelled(exception, contextCancellationToken))
             {
                 return false;
             }
@@ -258,7 +256,7 @@ public static bool ReportIfNonFatalAndCatchUnlessCanceled(Exception exception, E
 #endif
             static bool ReportAndCatchUnlessCanceled(Exception exception, CancellationToken contextCancellationToken, ErrorSeverity severity = ErrorSeverity.Uncategorized)
         {
-            if (IsCurrentOperationBeingCancelled(exception, contextCancellationToken))
+            if (ExceptionUtilities.IsCurrentOperationBeingCancelled(exception, contextCancellationToken))
             {
                 return false;
             }
@@ -290,7 +288,7 @@ static bool ReportAndCatchUnlessCanceled(Exception exception, CancellationToken
         [Obsolete("This is only to support places the compiler is catching and swallowing exceptions on the command line; do not use in new code.")]
         public static bool ReportIfNonFatalAndCatchUnlessCanceled(Exception exception, CancellationToken contextCancellationToken, ErrorSeverity severity = ErrorSeverity.Uncategorized)
         {
-            if (IsCurrentOperationBeingCancelled(exception, contextCancellationToken))
+            if (ExceptionUtilities.IsCurrentOperationBeingCancelled(exception, contextCancellationToken))
             {
                 return false;
             }

src/Compilers/Core/Portable/SourceGeneration/GeneratorDriver.cs

@@ -10,7 +10,6 @@
 using System.Linq;
 using System.Threading;
 using Microsoft.CodeAnalysis.Diagnostics;
-using Microsoft.CodeAnalysis.ErrorReporting;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Roslyn.Utilities;
 
@@ -193,9 +192,7 @@ internal GeneratorDriverState RunGeneratorsCore(Compilation compilation, Diagnos
                     {
                         generator.Initialize(pipelineContext);
                     }
-#pragma warning disable CS0618 // ReportIfNonFatalAndCatchUnlessCanceled is obsolete; tracked by https://github.com/dotnet/roslyn/issues/58375
-                    catch (Exception e) when (FatalError.ReportIfNonFatalAndCatchUnlessCanceled(e, cancellationToken))
-#pragma warning restore CS0618 // ReportIfNonFatalAndCatchUnlessCanceled is obsolete
+                    catch (Exception e)
                     {
                         ex = e;
                     }

src/Compilers/Core/Portable/SourceGeneration/Nodes/PostInitOutputNode.cs

@@ -9,9 +9,9 @@ namespace Microsoft.CodeAnalysis
 {
     internal sealed class PostInitOutputNode : IIncrementalGeneratorOutputNode
     {
-        private readonly Action<IncrementalGeneratorPostInitializationContext> _callback;
+        private readonly Action<IncrementalGeneratorPostInitializationContext, CancellationToken> _callback;
 
-        public PostInitOutputNode(Action<IncrementalGeneratorPostInitializationContext> callback)
+        public PostInitOutputNode(Action<IncrementalGeneratorPostInitializationContext, CancellationToken> callback)
         {
             _callback = callback;
         }
@@ -20,7 +20,7 @@ public PostInitOutputNode(Action<IncrementalGeneratorPostInitializationContext>
 
         public void AppendOutputs(IncrementalExecutionContext context, CancellationToken cancellationToken)
         {
-            _callback(new IncrementalGeneratorPostInitializationContext(context.Sources, cancellationToken));
+            _callback(new IncrementalGeneratorPostInitializationContext(context.Sources, cancellationToken), cancellationToken);
         }
     }
 }

src/Compilers/Core/Portable/SourceGeneration/Nodes/SourceOutputNode.cs

@@ -17,13 +17,13 @@ internal sealed class SourceOutputNode<TInput> : IIncrementalGeneratorOutputNode
     {
         private readonly IIncrementalGeneratorNode<TInput> _source;
 
-        private readonly Action<SourceProductionContext, TInput> _action;
+        private readonly Action<SourceProductionContext, TInput, CancellationToken> _action;
 
         private readonly IncrementalGeneratorOutputKind _outputKind;
 
         private readonly string _sourceExtension;
 
-        public SourceOutputNode(IIncrementalGeneratorNode<TInput> source, Action<SourceProductionContext, TInput> action, IncrementalGeneratorOutputKind outputKind, string sourceExtension)
+        public SourceOutputNode(IIncrementalGeneratorNode<TInput> source, Action<SourceProductionContext, TInput, CancellationToken> action, IncrementalGeneratorOutputKind outputKind, string sourceExtension)
         {
             _source = source;
             _action = action;
@@ -70,7 +70,7 @@ public NodeStateTable<TOutput> UpdateStateTable(DriverStateTable.Builder graphSt
                     try
                     {
                         var stopwatch = SharedStopwatch.StartNew();
-                        _action(context, entry.Item);
+                        _action(context, entry.Item, cancellationToken);
                         var sourcesAndDiagnostics = (sourcesBuilder.ToImmutable(), diagnostics.ToReadOnly());
                         nodeTable.AddEntry(sourcesAndDiagnostics, EntryState.Added, stopwatch.Elapsed, inputs, EntryState.Added);
                     }

src/Compilers/Core/Portable/SourceGeneration/Nodes/SyntaxReceiverInputNode.cs

@@ -96,7 +96,7 @@ public void VisitTree(Lazy<SyntaxNode> root, EntryState state, SemanticModel? mo
                             lastElapsedTime = stopwatch.Elapsed;
                         }
                     }
-                    catch (Exception e)
+                    catch (Exception e) when (!ExceptionUtilities.IsCurrentOperationBeingCancelled(e, cancellationToken))
                     {
                         throw new UserFunctionException(e);
                     }

src/Compilers/Core/Portable/SourceGeneration/UserFunction.cs

@@ -6,7 +6,7 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Threading;
-using Microsoft.CodeAnalysis.ErrorReporting;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis
 {
@@ -30,9 +30,7 @@ internal static Func<TInput, CancellationToken, TOutput> WrapUserFunction<TInput
                 {
                     return userFunction(input, token);
                 }
-#pragma warning disable CS0618 // ReportIfNonFatalAndCatchUnlessCanceled is obsolete; tracked by https://github.com/dotnet/roslyn/issues/58375
-                catch (Exception e) when (FatalError.ReportIfNonFatalAndCatchUnlessCanceled(e, token))
-#pragma warning restore CS0618 // ReportIfNonFatalAndCatchUnlessCanceled is obsolete
+                catch (Exception e) when (!ExceptionUtilities.IsCurrentOperationBeingCancelled(e, token))
                 {
                     throw new UserFunctionException(e);
                 }
@@ -44,34 +42,30 @@ internal static Func<TInput, CancellationToken, ImmutableArray<TOutput>> WrapUse
             return (input, token) => userFunction.WrapUserFunction()(input, token).ToImmutableArray();
         }
 
-        internal static Action<TInput> WrapUserAction<TInput>(this Action<TInput> userAction)
+        internal static Action<TInput, CancellationToken> WrapUserAction<TInput>(this Action<TInput> userAction)
         {
-            return input =>
+            return (input, token) =>
             {
                 try
                 {
                     userAction(input);
                 }
-#pragma warning disable CS0618 // ReportIfNonFatalAndCatchUnlessCanceled is obsolete; tracked by https://github.com/dotnet/roslyn/issues/58375
-                catch (Exception e) when (FatalError.ReportIfNonFatalAndCatchUnlessCanceled(e))
-#pragma warning restore CS0618 // ReportIfNonFatalAndCatchUnlessCanceled is obsolete
+                catch (Exception e) when (!ExceptionUtilities.IsCurrentOperationBeingCancelled(e, token))
                 {
                     throw new UserFunctionException(e);
                 }
             };
         }
 
-        internal static Action<TInput1, TInput2> WrapUserAction<TInput1, TInput2>(this Action<TInput1, TInput2> userAction)
+        internal static Action<TInput1, TInput2, CancellationToken> WrapUserAction<TInput1, TInput2>(this Action<TInput1, TInput2> userAction)
         {
-            return (input1, input2) =>
+            return (input1, input2, token) =>
             {
                 try
                 {
                     userAction(input1, input2);
                 }
-#pragma warning disable CS0618 // ReportIfNonFatalAndCatchUnlessCanceled is obsolete; tracked by https://github.com/dotnet/roslyn/issues/58375
-                catch (Exception e) when (FatalError.ReportIfNonFatalAndCatchUnlessCanceled(e))
-#pragma warning restore CS0618 // ReportIfNonFatalAndCatchUnlessCanceled is obsolete
+                catch (Exception e) when (!ExceptionUtilities.IsCurrentOperationBeingCancelled(e, token))
                 {
                     throw new UserFunctionException(e);
                 }