Add unit test for #46950

src/Compilers/CSharp/Test/Semantic/Diagnostics/GetDiagnosticsTests.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Linq;
@@ -664,5 +665,32 @@ void M2()
 
             await Task.WhenAll(tasks);
         }
+
+        [Theory, CombinatorialData]
+        [WorkItem(46950, "https://github.com/dotnet/roslyn/issues/46950")]
+        public async Task TestGetAnalyzerSyntaxDiagnosticsWithCancellation(bool concurrent)
+        {
+            var source = @"class C { }";
+            var compilation = CreateCompilation(source);
+            compilation = compilation.WithOptions(compilation.Options.WithConcurrentBuild(concurrent));
+            var tree = compilation.SyntaxTrees.First();
+
+            var analyzer = new RegisterSyntaxTreeCancellationAnalyzer();
+            var analyzers = ImmutableArray.Create<DiagnosticAnalyzer>(analyzer);
+            var compilationWithAnalyzers = compilation.WithAnalyzers(analyzers,
+                new CompilationWithAnalyzersOptions(
+                    new AnalyzerOptions(ImmutableArray<AdditionalText>.Empty),
+                    onAnalyzerException: null,
+                    concurrentAnalysis: concurrent,
+                    logAnalyzerExecutionTime: false));
+
+            // First call into analyzer mimics cancellation.
+            await Assert.ThrowsAsync<OperationCanceledException>(() => compilationWithAnalyzers.GetAnalyzerSyntaxDiagnosticsAsync(tree, analyzer.CancellationToken));
+
+            // Second call into analyzer reports diagnostic.
+            var diagnostics = await compilationWithAnalyzers.GetAnalyzerSyntaxDiagnosticsAsync(tree, CancellationToken.None);
+            var diagnostic = Assert.Single(diagnostics);
+            Assert.Equal(RegisterSyntaxTreeCancellationAnalyzer.DiagnosticId, diagnostic.Id);
+        }
     }
 }

src/Test/Utilities/Portable/Diagnostics/CommonDiagnosticAnalyzers.cs

@@ -2242,5 +2242,42 @@ private void AnalyzeAdditionalFile(AdditionalFileAnalysisContext context)
                 context.ReportDiagnostic(Diagnostic.Create(Descriptor, location));
             }
         }
+
+        [DiagnosticAnalyzer(LanguageNames.CSharp, LanguageNames.VisualBasic)]
+        public sealed class RegisterSyntaxTreeCancellationAnalyzer : DiagnosticAnalyzer
+        {
+            public const string DiagnosticId = "ID0001";
+            private static readonly DiagnosticDescriptor s_descriptor = new DiagnosticDescriptor(
+                DiagnosticId,
+                "Title",
+                "Message",
+                "Category",
+                defaultSeverity: DiagnosticSeverity.Warning,
+                isEnabledByDefault: true);
+            private readonly CancellationTokenSource _cancellationTokenSource = new CancellationTokenSource();
+            public CancellationToken CancellationToken => _cancellationTokenSource.Token;
+
+            public override ImmutableArray<DiagnosticDescriptor> SupportedDiagnostics => ImmutableArray.Create(s_descriptor);
+            public override void Initialize(AnalysisContext context)
+            {
+                context.RegisterSyntaxTreeAction(context =>
+                {
+                    // Mimic cancellation by throwing an OperationCanceledException in first callback.
+                    if (!_cancellationTokenSource.IsCancellationRequested)
+                    {
+                        _cancellationTokenSource.Cancel();
+
+                        while (true)
+                        {
+                            context.CancellationToken.ThrowIfCancellationRequested();
+                        }
+
+                        throw ExceptionUtilities.Unreachable;
+                    }
+
+                    context.ReportDiagnostic(Diagnostic.Create(s_descriptor, context.Tree.GetRoot().GetLocation()));
+                });
+            }
+        }
     }
 }