Create a derived type for cancelled run results

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

@@ -2121,5 +2121,138 @@ class C { }
             driver = driver.RunGenerators(compilation);
             Assert.Single(referenceList, modifiedRef.Display);
         }
+
+        [Fact]
+        public void Generator_Driver_Supports_Graceful_Cancellation()
+        {
+            var source = @"
+class C { }
+";
+            var parseOptions = TestOptions.RegularPreview;
+            Compilation compilation = CreateCompilation(source, options: TestOptions.DebugDll, parseOptions: parseOptions);
+            compilation.VerifyDiagnostics();
+            Assert.Single(compilation.SyntaxTrees);
+
+            CancellationTokenSource cts = new CancellationTokenSource();
+
+            var generator = new PipelineCallbackGenerator(ctx =>
+            {
+                ctx.RegisterSourceOutput(ctx.CompilationProvider, (context, ct) => cts.Cancel());
+            });
+
+            GeneratorDriver driver = CSharpGeneratorDriver.Create(new ISourceGenerator[] { generator.AsSourceGenerator() }, parseOptions: parseOptions, driverOptions: new GeneratorDriverOptions(IncrementalGeneratorOutputKind.None, true));
+            driver = driver.RunGenerators(compilation, cts.Token);
+
+            Assert.True(cts.Token.IsCancellationRequested);
+
+            var runResult = driver.GetRunResult();
+            Assert.IsType<GeneratorDriverRunResult.CancelledResult>(runResult);
+            Assert.Empty(runResult.Results);
+        }
+
+        [Fact]
+        public void Generator_Driver_Ignores_Cancelled_Results()
+        {
+            var source = @"
+class C { }
+";
+            var parseOptions = TestOptions.RegularPreview;
+            Compilation compilation = CreateCompilation(source, options: TestOptions.DebugDll, parseOptions: parseOptions);
+            compilation.VerifyDiagnostics();
+            Assert.Single(compilation.SyntaxTrees);
+
+            CancellationTokenSource cts = new CancellationTokenSource();
+
+            var generator = new PipelineCallbackGenerator(ctx =>
+            {
+                ctx.RegisterSourceOutput(ctx.CompilationProvider, (context, ct) => { context.AddSource("gen", ""); cts.Cancel(); });
+            });
+
+            GeneratorDriver driver = CSharpGeneratorDriver.Create(new ISourceGenerator[] { generator.AsSourceGenerator() }, parseOptions: parseOptions, driverOptions: new GeneratorDriverOptions(IncrementalGeneratorOutputKind.None, true));
+            driver = driver.RunGenerators(compilation, cts.Token);
+
+            Assert.True(cts.Token.IsCancellationRequested);
+
+            var runResult = driver.GetRunResult();
+            Assert.IsType<GeneratorDriverRunResult.CancelledResult>(runResult);
+            Assert.Empty(runResult.Results);
+        }
+
+        [Fact]
+        public void Generator_Driver_Does_Not_Run_Generators_After_Cancellation()
+        {
+            var source = @"
+class C { }
+";
+            var parseOptions = TestOptions.RegularPreview;
+            Compilation compilation = CreateCompilation(source, options: TestOptions.DebugDll, parseOptions: parseOptions);
+            compilation.VerifyDiagnostics();
+            Assert.Single(compilation.SyntaxTrees);
+
+            CancellationTokenSource cts = new CancellationTokenSource();
+
+            bool invoked = false;
+            var generator = new PipelineCallbackGenerator(ctx =>
+            {
+                ctx.RegisterSourceOutput(ctx.CompilationProvider, (context, ct) => cts.Cancel());
+            });
+
+            var generator2 = new PipelineCallbackGenerator2(ctx =>
+            {
+                ctx.RegisterSourceOutput(ctx.CompilationProvider, (context, ct) => invoked = true);
+            });
+
+            GeneratorDriver driver = CSharpGeneratorDriver.Create(new ISourceGenerator[] { generator.AsSourceGenerator(), generator2.AsSourceGenerator() }, parseOptions: parseOptions, driverOptions: new GeneratorDriverOptions(IncrementalGeneratorOutputKind.None, true));
+            driver = driver.RunGenerators(compilation, cts.Token);
+
+            Assert.True(cts.Token.IsCancellationRequested);
+
+            var runResult = driver.GetRunResult();
+            Assert.IsType<GeneratorDriverRunResult.CancelledResult>(runResult);
+            Assert.Empty(runResult.Results);
+            Assert.False(invoked);
+        }
+
+        [Fact]
+        public void Generator_Driver_Uses_Partial_Results_From_Previous_Cancellation()
+        {
+            var source = @"
+class C { }
+";
+            var parseOptions = TestOptions.RegularPreview;
+            Compilation compilation = CreateCompilation(source, options: TestOptions.DebugDll, parseOptions: parseOptions);
+            compilation.VerifyDiagnostics();
+            Assert.Single(compilation.SyntaxTrees);
+
+            CancellationTokenSource cts = new CancellationTokenSource();
+            int callCount = 0;
+
+            var generator = new PipelineCallbackGenerator(ctx =>
+            {
+                ctx.RegisterSourceOutput(ctx.CompilationProvider, (context, ct) => { callCount++; context.AddSource("gen1", ""); });
+            });
+
+            var generator2 = new PipelineCallbackGenerator2(ctx =>
+            {
+                ctx.RegisterSourceOutput(ctx.CompilationProvider, (context, ct) => cts.Cancel());
+            });
+
+            GeneratorDriver driver = CSharpGeneratorDriver.Create(new ISourceGenerator[] { generator.AsSourceGenerator(), generator2.AsSourceGenerator() }, parseOptions: parseOptions, driverOptions: new GeneratorDriverOptions(IncrementalGeneratorOutputKind.None, true));
+            driver = driver.RunGenerators(compilation, cts.Token);
+
+            Assert.True(cts.Token.IsCancellationRequested);
+
+            var runResult = driver.GetRunResult();
+            Assert.IsType<GeneratorDriverRunResult.CancelledResult>(runResult);
+            Assert.Empty(runResult.Results);
+            Assert.Equal(1, callCount);
+
+            // re-run with the partial results and confirm that the cached results are not re-computed, but we still get run results
+            driver = driver.RunGenerators(compilation, cancellationToken: default);
+            var runResult2 = driver.GetRunResult();
+
+            Assert.Equal(2, runResult2.Results.Length);
+            Assert.Equal(1, callCount);
+        }
     }
 }

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

@@ -401,15 +401,16 @@ private DriverStateTable.Builder GetBuilder(DriverStateTable previous)
         {
             var options = CSharpParseOptions.Default.WithLanguageVersion(LanguageVersion.CSharp10);
             var c = CSharpCompilation.Create("empty");
-            var state = new GeneratorDriverState(options,
+            var state = new GeneratorDriverState(new GeneratorDriverOptions(IncrementalGeneratorOutputKind.None),
+                    options,
                     CompilerAnalyzerConfigOptionsProvider.Empty,
                     ImmutableArray<ISourceGenerator>.Empty,
                     ImmutableArray<IIncrementalGenerator>.Empty,
                     ImmutableArray<AdditionalText>.Empty,
                     ImmutableArray<GeneratorState>.Empty,
                     previous,
-                    disabledOutputs: IncrementalGeneratorOutputKind.None,
-                    runtime: TimeSpan.Zero);
+                    elapsedTime: TimeSpan.Zero,
+                    cancelled: false);
 
             return new DriverStateTable.Builder(c, state, ImmutableArray<ISyntaxInputNode>.Empty);
         }

src/Compilers/Core/Portable/PublicAPI.Unshipped.txt

@@ -13,9 +13,10 @@ Microsoft.CodeAnalysis.GeneratorDriver.WithUpdatedParseOptions(Microsoft.CodeAna
 Microsoft.CodeAnalysis.GeneratorDriverOptions
 Microsoft.CodeAnalysis.GeneratorDriverOptions.GeneratorDriverOptions() -> void
 Microsoft.CodeAnalysis.GeneratorDriverOptions.GeneratorDriverOptions(Microsoft.CodeAnalysis.IncrementalGeneratorOutputKind disabledOutputs) -> void
-Microsoft.CodeAnalysis.GeneratorDriverRunResult.Cancelled.get -> bool
+Microsoft.CodeAnalysis.GeneratorDriverOptions.GeneratorDriverOptions(Microsoft.CodeAnalysis.IncrementalGeneratorOutputKind disabledOutputs, bool enableGracefulCancellation) -> void
+Microsoft.CodeAnalysis.GeneratorDriverRunResult.CancelledResult
+Microsoft.CodeAnalysis.GeneratorDriverRunResult.CancelledResult.LastGeneratorRunning.get -> Microsoft.CodeAnalysis.GeneratorRunResult?
 Microsoft.CodeAnalysis.GeneratorDriverRunResult.ElapsedTime.get -> System.TimeSpan
-Microsoft.CodeAnalysis.GeneratorDriverRunResult.WasCancelled.get -> bool
 Microsoft.CodeAnalysis.GeneratorExtensions
 Microsoft.CodeAnalysis.GeneratorRunResult.ElapsedTime.get -> System.TimeSpan
 Microsoft.CodeAnalysis.IFieldSymbol.FixedSize.get -> int
@@ -126,6 +127,7 @@ Microsoft.CodeAnalysis.SyntaxValueProvider.CreateSyntaxProvider<T>(System.Func<M
 Microsoft.CodeAnalysis.SyntaxValueProvider.SyntaxValueProvider() -> void
 override Microsoft.CodeAnalysis.Text.TextChangeRange.ToString() -> string!
 readonly Microsoft.CodeAnalysis.GeneratorDriverOptions.DisabledOutputs -> Microsoft.CodeAnalysis.IncrementalGeneratorOutputKind
+readonly Microsoft.CodeAnalysis.GeneratorDriverOptions.EnableGracefulCancellation -> bool
 static Microsoft.CodeAnalysis.CaseInsensitiveComparison.Compare(System.ReadOnlySpan<char> left, System.ReadOnlySpan<char> right) -> int
 static Microsoft.CodeAnalysis.CaseInsensitiveComparison.Equals(System.ReadOnlySpan<char> left, System.ReadOnlySpan<char> right) -> bool
 override Microsoft.CodeAnalysis.Diagnostics.AnalyzerFileReference.GetGenerators(string! language) -> System.Collections.Immutable.ImmutableArray<Microsoft.CodeAnalysis.ISourceGenerator!>

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

@@ -36,7 +36,7 @@ internal GeneratorDriver(GeneratorDriverState state)
         internal GeneratorDriver(ParseOptions parseOptions, ImmutableArray<ISourceGenerator> generators, AnalyzerConfigOptionsProvider optionsProvider, ImmutableArray<AdditionalText> additionalTexts, GeneratorDriverOptions driverOptions)
         {
             (var filteredGenerators, var incrementalGenerators) = GetIncrementalGenerators(generators, SourceExtension);
-            _state = new GeneratorDriverState(parseOptions, optionsProvider, filteredGenerators, incrementalGenerators, additionalTexts, ImmutableArray.Create(new GeneratorState[filteredGenerators.Length]), DriverStateTable.Empty, driverOptions.DisabledOutputs, elapsedTime: TimeSpan.Zero, cancelled: false);
+            _state = new GeneratorDriverState(driverOptions, parseOptions, optionsProvider, filteredGenerators, incrementalGenerators, additionalTexts, ImmutableArray.Create(new GeneratorState[filteredGenerators.Length]), DriverStateTable.Empty, elapsedTime: TimeSpan.Zero, cancelled: false);
         }
 
         public GeneratorDriver RunGenerators(Compilation compilation, CancellationToken cancellationToken = default)
@@ -129,6 +129,13 @@ public GeneratorDriver WithUpdatedAnalyzerConfigOptions(AnalyzerConfigOptionsPro
 
         public GeneratorDriverRunResult GetRunResult()
         {
+            if (_state.Cancelled)
+            {
+                Debug.Assert(_state.Options.EnableGracefulCancellation);
+
+                return new GeneratorDriverRunResult.CancelledResult(_state.ElapsedTime, getCancelledRunResult(_state));
+            }
+
             var results = _state.Generators.ZipAsArray(
                             _state.GeneratorStates,
                             (generator, generatorState)
@@ -137,7 +144,7 @@ public GeneratorDriverRunResult GetRunResult()
                                                           exception: generatorState.Exception,
                                                           generatedSources: getGeneratorSources(generatorState),
                                                           elapsedTime: generatorState.ElapsedTime));
-            return new GeneratorDriverRunResult(results, _state.ElapsedTime, _state.Cancelled);
+            return new GeneratorDriverRunResult(results, _state.ElapsedTime);
 
             static ImmutableArray<GeneratedSourceResult> getGeneratorSources(GeneratorState generatorState)
             {
@@ -152,6 +159,19 @@ static ImmutableArray<GeneratedSourceResult> getGeneratorSources(GeneratorState
                 }
                 return sources.ToImmutableAndFree();
             }
+
+            static GeneratorRunResult? getCancelledRunResult(GeneratorDriverState driverState)
+            {
+                for (int i = 0; i < driverState.GeneratorStates.Length; i++)
+                {
+                    if (driverState.GeneratorStates[i].Cancelled)
+                    {
+                        return new GeneratorRunResult(driverState.Generators[i], ImmutableArray<GeneratedSourceResult>.Empty, ImmutableArray<Diagnostic>.Empty, exception: null, elapsedTime: driverState.GeneratorStates[i].ElapsedTime);
+                    }
+                }
+
+                return null;
+            }
         }
 
         internal GeneratorDriverState RunGeneratorsCore(Compilation compilation, DiagnosticBag? diagnosticsBag, CancellationToken cancellationToken = default)
@@ -260,11 +280,11 @@ internal GeneratorDriverState RunGeneratorsCore(Compilation compilation, Diagnos
                 {
                     stateBuilder[i] = SetGeneratorException(MessageProvider, stateBuilder[i], state.Generators[i], ufe.InnerException, diagnosticsBag, generatorTimer.Elapsed);
                 }
-                catch (OperationCanceledException)
+                catch (OperationCanceledException) when (state.Options.EnableGracefulCancellation)
                 {
                     // when cancelled, we record the time it spent generating, but don't include anything that was partially generated.
-                    // this allows us to know if a await runnning generator is frequently causing the cacnellation
-                    stateBuilder[i] = new GeneratorState(generatorState.Info, generatorState.PostInitTrees, generatorState.InputNodes, generatorState.OutputNodes, ImmutableArray<GeneratedSyntaxTree>.Empty, ImmutableArray<Diagnostic>.Empty, generatorTimer.Elapsed);
+                    // this allows us to know if a runnning generator is frequently causing the cancellation
+                    stateBuilder[i] = new GeneratorState(generatorState.Info, generatorState.PostInitTrees, generatorState.InputNodes, generatorState.OutputNodes, ImmutableArray<GeneratedSyntaxTree>.Empty, ImmutableArray<Diagnostic>.Empty, generatorTimer.Elapsed, cancelled: true);
                     cancelled = true;
                     break;
                 }
@@ -281,9 +301,10 @@ private IncrementalExecutionContext UpdateOutputs(ImmutableArray<IIncrementalGen
             foreach (var outputNode in outputNodes)
             {
                 // if we're looking for this output kind, and it has not been explicitly disabled
-                if (outputKind.HasFlag(outputNode.Kind) && !_state.DisabledOutputs.HasFlag(outputNode.Kind))
+                if (outputKind.HasFlag(outputNode.Kind) && !_state.Options.DisabledOutputs.HasFlag(outputNode.Kind))
                 {
                     outputNode.AppendOutputs(context, cancellationToken);
+                    cancellationToken.ThrowIfCancellationRequested();
                 }
             }
             return context;

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

@@ -15,9 +15,18 @@ public readonly struct GeneratorDriverOptions
     {
         public readonly IncrementalGeneratorOutputKind DisabledOutputs;
 
+        public readonly bool EnableGracefulCancellation;
+
         public GeneratorDriverOptions(IncrementalGeneratorOutputKind disabledOutputs)
+            : this(disabledOutputs, enableGracefulCancellation: false)
+        {
+        }
+
+        public GeneratorDriverOptions(IncrementalGeneratorOutputKind disabledOutputs, bool enableGracefulCancellation)
         {
             DisabledOutputs = disabledOutputs;
+            EnableGracefulCancellation = enableGracefulCancellation;
         }
+
     }
 }

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

@@ -11,14 +11,14 @@ namespace Microsoft.CodeAnalysis
 {
     internal readonly struct GeneratorDriverState
     {
-        internal GeneratorDriverState(ParseOptions parseOptions,
+        internal GeneratorDriverState(GeneratorDriverOptions options,
+                                      ParseOptions parseOptions,
                                       AnalyzerConfigOptionsProvider optionsProvider,
                                       ImmutableArray<ISourceGenerator> sourceGenerators,
                                       ImmutableArray<IIncrementalGenerator> incrementalGenerators,
                                       ImmutableArray<AdditionalText> additionalTexts,
                                       ImmutableArray<GeneratorState> generatorStates,
                                       DriverStateTable stateTable,
-                                      IncrementalGeneratorOutputKind disabledOutputs,
                                       TimeSpan elapsedTime,
                                       bool cancelled)
         {
@@ -27,9 +27,9 @@ internal GeneratorDriverState(ParseOptions parseOptions,
             GeneratorStates = generatorStates;
             AdditionalTexts = additionalTexts;
             ParseOptions = parseOptions;
+            Options = options;
             OptionsProvider = optionsProvider;
             StateTable = stateTable;
-            DisabledOutputs = disabledOutputs;
             ElapsedTime = elapsedTime;
             Cancelled = cancelled;
             Debug.Assert(Generators.Length == GeneratorStates.Length);
@@ -80,11 +80,6 @@ internal GeneratorDriverState(ParseOptions parseOptions,
 
         internal readonly DriverStateTable StateTable;
 
-        /// <summary>
-        /// A bit field containing the output kinds that should not be produced by this generator driver.
-        /// </summary>
-        internal readonly IncrementalGeneratorOutputKind DisabledOutputs;
-
         /// <summary>
         /// The time spent during the pass that created this state.
         /// </summary>
@@ -95,6 +90,11 @@ internal GeneratorDriverState(ParseOptions parseOptions,
         /// </summary>
         internal readonly bool Cancelled;
 
+        /// <summary>
+        /// The set of options passed when this driver was created.
+        /// </summary>
+        internal readonly GeneratorDriverOptions Options;
+
         internal GeneratorDriverState With(
             ImmutableArray<ISourceGenerator>? sourceGenerators = null,
             ImmutableArray<IIncrementalGenerator>? incrementalGenerators = null,
@@ -103,19 +103,18 @@ internal GeneratorDriverState With(
             DriverStateTable? stateTable = null,
             ParseOptions? parseOptions = null,
             AnalyzerConfigOptionsProvider? optionsProvider = null,
-            IncrementalGeneratorOutputKind? disabledOutputs = null,
             TimeSpan? elapsedTime = null,
             bool? cancelled = null)
         {
             return new GeneratorDriverState(
+                Options,
                 parseOptions ?? this.ParseOptions,
                 optionsProvider ?? this.OptionsProvider,
                 sourceGenerators ?? this.Generators,
                 incrementalGenerators ?? this.IncrementalGenerators,
                 additionalTexts ?? this.AdditionalTexts,
                 generatorStates ?? this.GeneratorStates,
                 stateTable ?? this.StateTable,
-                disabledOutputs ?? this.DisabledOutputs,
                 elapsedTime ?? this.ElapsedTime,
                 cancelled ?? this.Cancelled
                 );