Merge pull request #187 from manne/feature/cancellation

Make the generator cancellable

Author: AArnott

src/CodeGeneration.Roslyn.Engine/CompilationGenerator.cs

@@ -123,13 +123,13 @@ public async Task GenerateAsync(IProgress<Diagnostic> progress = null, Cancellat
                                     inputSyntaxTree,
                                     this.ProjectDirectory,
                                     this.LoadPlugin,
-                                    progress);
-
-                                var outputText = generatedSyntaxTree.GetText(cancellationToken);
+                                    progress,
+                                    cancellationToken);
+                                var outputText = await generatedSyntaxTree.GetTextAsync(cancellationToken);
                                 using (var outputFileStream = File.OpenWrite(outputFilePath))
                                 using (var outputWriter = new StreamWriter(outputFileStream))
                                 {
-                                    outputText.Write(outputWriter);
+                                    outputText.Write(outputWriter, cancellationToken);
 
                                     // Truncate any data that may be beyond this point if the file existed previously.
                                     outputWriter.Flush();
@@ -150,9 +150,9 @@ public async Task GenerateAsync(IProgress<Diagnostic> progress = null, Cancellat
                             catch (IOException ex) when (ex.HResult == ProcessCannotAccessFileHR && retriesLeft > 0)
                             {
                                 retriesLeft--;
-                                Task.Delay(200).Wait();
+                                await Task.Delay(200, cancellationToken);
                             }
-                            catch (Exception ex)
+                            catch (Exception ex) when (!(ex is OperationCanceledException))
                             {
                                 ReportError(progress, "CGR001", inputSyntaxTree, ex);
                                 fileFailures.Add(ex);

src/CodeGeneration.Roslyn.Engine/DocumentTransform.cs

@@ -42,13 +42,15 @@ public static class DocumentTransform
         /// <param name="projectDirectory">The path of the <c>.csproj</c> project file.</param>
         /// <param name="assemblyLoader">A function that can load an assembly with the given name.</param>
         /// <param name="progress">Reports warnings and errors in code generation.</param>
+        /// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
         /// <returns>A task whose result is the generated document.</returns>
         public static async Task<SyntaxTree> TransformAsync(
             CSharpCompilation compilation,
             SyntaxTree inputDocument,
             string projectDirectory,
             Func<AssemblyName, Assembly> assemblyLoader,
-            IProgress<Diagnostic> progress)
+            IProgress<Diagnostic> progress,
+            CancellationToken cancellationToken)
         {
             Requires.NotNull(compilation, nameof(compilation));
             Requires.NotNull(inputDocument, nameof(inputDocument));
@@ -70,17 +72,19 @@ public static async Task<SyntaxTree> TransformAsync(
             var emittedAttributeLists = ImmutableArray<AttributeListSyntax>.Empty;
             var emittedMembers = ImmutableArray<MemberDeclarationSyntax>.Empty;
 
-            var root = await inputDocument.GetRootAsync();
+            var root = await inputDocument.GetRootAsync(cancellationToken);
             var memberNodes = root
                 .DescendantNodesAndSelf(n => n is CompilationUnitSyntax || n is NamespaceDeclarationSyntax || n is TypeDeclarationSyntax)
                 .OfType<CSharpSyntaxNode>();
 
             foreach (var memberNode in memberNodes)
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 var attributeData = GetAttributeData(compilation, inputSemanticModel, memberNode);
                 var generators = FindCodeGenerators(attributeData, assemblyLoader);
                 foreach (var generator in generators)
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     var context = new TransformationContext(
                         memberNode,
                         inputSemanticModel,
@@ -91,7 +95,7 @@ public static async Task<SyntaxTree> TransformAsync(
 
                     var richGenerator = generator as IRichCodeGenerator ?? new EnrichingCodeGeneratorProxy(generator);
 
-                    var emitted = await richGenerator.GenerateRichAsync(context, progress, CancellationToken.None);
+                    var emitted = await richGenerator.GenerateRichAsync(context, progress, cancellationToken);
 
                     emittedExterns = emittedExterns.AddRange(emitted.Externs);
                     emittedUsings = emittedUsings.AddRange(emitted.Usings);
@@ -244,7 +248,7 @@ public Task<SyntaxList<MemberDeclarationSyntax>> GenerateAsync(
 
             public async Task<RichGenerationResult> GenerateRichAsync(TransformationContext context, IProgress<Diagnostic> progress, CancellationToken cancellationToken)
             {
-                var generatedMembers = await CodeGenerator.GenerateAsync(context, progress, CancellationToken.None);
+                var generatedMembers = await CodeGenerator.GenerateAsync(context, progress, cancellationToken);
 
                 // Figure out ancestry for the generated type, including nesting types and namespaces.
                 var wrappedMembers = context.ProcessingNode.Ancestors().Aggregate(generatedMembers, WrapInAncestor);

src/CodeGeneration.Roslyn.Tests/Helpers/CompilationTestsBase.cs

@@ -6,6 +6,7 @@
 using System.IO;
 using System.Linq;
 using System.Reflection;
+using System.Threading;
 using System.Threading.Tasks;
 using CodeGeneration.Roslyn;
 using CodeGeneration.Roslyn.Engine;
@@ -78,7 +79,7 @@ protected static async Task<SyntaxTree> GenerateAsync(string source)
         var diagnostics = compilation.GetDiagnostics();
         Assert.Empty(diagnostics.Where(x => x.Severity >= DiagnosticSeverity.Warning));
         var progress = new Progress<Diagnostic>();
-        var result = await DocumentTransform.TransformAsync(compilation, tree, null, Assembly.Load, progress);
+        var result = await DocumentTransform.TransformAsync(compilation, tree, null, Assembly.Load, progress, CancellationToken.None);
         return result;
     }
 

src/CodeGeneration.Roslyn.Tool/Program.cs

@@ -1,6 +1,8 @@
 // Copyright (c) Andrew Arnott. All rights reserved.
 // Licensed under the MS-PL license. See LICENSE.txt file in the project root for full license information.
 
+using System.Threading;
+
 namespace CodeGeneration.Roslyn.Generate
 {
     using System;
@@ -14,6 +16,24 @@ namespace CodeGeneration.Roslyn.Generate
     internal static class Program
     {
         private static async Task<int> Main(string[] args)
+        {
+            using (var cancellationTokenSource = new CancellationTokenSource())
+            {
+                var cancellationToken = cancellationTokenSource.Token;
+                Console.CancelKeyPress += ConsoleOnCancelKeyPress;
+
+                void ConsoleOnCancelKeyPress(object sender, ConsoleCancelEventArgs e)
+                {
+                    Console.CancelKeyPress -= ConsoleOnCancelKeyPress;
+                    cancellationTokenSource.Cancel();
+                    e.Cancel = true;
+                }
+
+                return await Core(args, cancellationToken);
+            }
+        }
+
+        private static async Task<int> Core(string[] args, CancellationToken cancellationToken)
         {
             IReadOnlyList<string> compile = Array.Empty<string>();
             IReadOnlyList<string> refs = Array.Empty<string>();
@@ -66,7 +86,7 @@ private static async Task<int> Main(string[] args)
 
             try
             {
-                await generator.GenerateAsync(progress);
+                await generator.GenerateAsync(progress, cancellationToken);
             }
             catch (Exception e)
             {
@@ -77,7 +97,7 @@ private static async Task<int> Main(string[] args)
 
             if (generatedCompileItemFile != null)
             {
-                File.WriteAllLines(generatedCompileItemFile, generator.GeneratedFiles);
+                await File.WriteAllLinesAsync(generatedCompileItemFile, generator.GeneratedFiles, cancellationToken);
             }
 
             foreach (var file in generator.GeneratedFiles)