Handle cancellation tokens correctly in 'fix all occurences'

Use the correct cancellation token in 'fix all occurences' and also bail on cancellation in more places in the 'fix all occurences' code path.

src/EditorFeatures/Core/Implementation/Suggestions/FixAllGetFixesService.cs

@@ -45,6 +45,7 @@ public async Task<IEnumerable<CodeActionOperation>> GetFixAllOperationsAsync(Fix
                     allowCancel: true,
                     action: waitContext =>
                     {
+                        fixAllContext.CancellationToken.ThrowIfCancellationRequested();
                         using (var linkedCts =
                             CancellationTokenSource.CreateLinkedTokenSource(waitContext.CancellationToken, fixAllContext.CancellationToken))
                         {
@@ -85,14 +86,17 @@ private async Task<IEnumerable<CodeActionOperation>> GetFixAllOperationsAsync(Co
             var cancellationToken = fixAllContext.CancellationToken;
             var workspace = fixAllContext.Project.Solution.Workspace;
 
+            cancellationToken.ThrowIfCancellationRequested();
             var operations = await codeAction.GetOperationsAsync(cancellationToken).ConfigureAwait(false);
             if (operations == null)
             {
                 return null;
             }
 
+            cancellationToken.ThrowIfCancellationRequested();
             var newSolution = await codeAction.GetChangedSolutionInternalAsync(cancellationToken).ConfigureAwait(false);
 
+            cancellationToken.ThrowIfCancellationRequested();
             using (Logger.LogBlock(FunctionId.CodeFixes_FixAllOccurrencesPreviewChanges, cancellationToken))
             {
                 var previewService = workspace.Services.GetService<IPreviewDialogService>();
@@ -121,14 +125,16 @@ private async Task<IEnumerable<CodeActionOperation>> GetFixAllOperationsAsync(Co
             }
 
             // Get a code action, with apply changes operation replaced with the newSolution.
-            return GetNewFixAllOperations(operations, newSolution);
+            return GetNewFixAllOperations(operations, newSolution, cancellationToken);
         }
 
-        private IEnumerable<CodeActionOperation> GetNewFixAllOperations(IEnumerable<CodeActionOperation> operations, Solution newSolution)
+        private IEnumerable<CodeActionOperation> GetNewFixAllOperations(IEnumerable<CodeActionOperation> operations, Solution newSolution, CancellationToken cancellationToken)
         {
             bool foundApplyChanges = false;
             foreach (var operation in operations)
             {
+                cancellationToken.ThrowIfCancellationRequested();
+
                 if (!foundApplyChanges)
                 {
                     var applyChangesOperation = operation as ApplyChangesOperation;

src/EditorFeatures/Core/Implementation/Suggestions/PreviewChangesCodeAction.cs

@@ -31,6 +31,7 @@ public override string Title
 
         protected override async Task<IEnumerable<CodeActionOperation>> ComputeOperationsAsync(CancellationToken cancellationToken)
         {
+            cancellationToken.ThrowIfCancellationRequested();
             var previewDialogService = _workspace.Services.GetService<IPreviewDialogService>();
             if (previewDialogService == null)
             {
@@ -53,6 +54,7 @@ protected override async Task<IEnumerable<CodeActionOperation>> ComputeOperation
                 return null;
             }
 
+            cancellationToken.ThrowIfCancellationRequested();
             return await _originalCodeAction.GetOperationsAsync(cancellationToken).ConfigureAwait(false);
         }
     }

src/Features/Core/CodeFixes/FixAllOccurrences/FixAllCodeAction.cs

@@ -43,10 +43,12 @@ public override string Title
 
         protected override async Task<IEnumerable<CodeActionOperation>> ComputeOperationsAsync(CancellationToken cancellationToken)
         {
+            cancellationToken.ThrowIfCancellationRequested();
             FixAllLogger.LogContext(_fixAllContext, IsInternalCodeFixProvider(_fixAllContext.CodeFixProvider));
 
             var service = _fixAllContext.Project.Solution.Workspace.Services.GetService<IFixAllGetFixesService>();
-            return await service.GetFixAllOperationsAsync(_fixAllProvider, _fixAllContext).ConfigureAwait(false);
+            // Use the new cancellation token instead of the stale one present inside _fixAllContext.
+            return await service.GetFixAllOperationsAsync(_fixAllProvider, _fixAllContext.WithCancellationToken(cancellationToken)).ConfigureAwait(false);
         }
 
         private static bool IsInternalCodeFixProvider(CodeFixProvider fixer)

src/Workspaces/Core/Portable/CodeFixes/FixAllOccurrences/BatchFixAllProvider.cs

@@ -88,6 +88,7 @@ public async virtual Task AddDocumentFixesAsync(Document document, ImmutableArra
 
             for (var i = 0; i < diagnostics.Length; i++)
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 var diagnostic = diagnostics[i];
                 fixerTasks[i] = Task.Run(async () =>
                 {
@@ -112,6 +113,7 @@ public async virtual Task AddDocumentFixesAsync(Document document, ImmutableArra
 
                     foreach (var fix in fixes)
                     {
+                        cancellationToken.ThrowIfCancellationRequested();
                         if (fix != null && fix.EquivalenceKey == fixAllContext.CodeActionEquivalenceKey)
                         {
                             addFix(fix);
@@ -250,12 +252,14 @@ public virtual string GetFixAllTitle(FixAllContext fixAllContext)
                         var tasks = new Task[projectsToFix.Length];
                         for (int i = 0; i < projectsToFix.Length; i++)
                         {
+                            fixAllContext.CancellationToken.ThrowIfCancellationRequested();
                             var projectToFix = projectsToFix[i];
                             tasks[i] = Task.Run(async () =>
                             {
                                 var projectDiagnostics = await fixAllContext.GetAllDiagnosticsAsync(projectToFix).ConfigureAwait(false);
                                 foreach (var diagnostic in projectDiagnostics)
                                 {
+                                    fixAllContext.CancellationToken.ThrowIfCancellationRequested();
                                     diagnostics.Add(diagnostic);
                                 }
                             }, fixAllContext.CancellationToken);
@@ -285,6 +289,7 @@ public virtual string GetFixAllTitle(FixAllContext fixAllContext)
             var builder = ImmutableDictionary.CreateBuilder<Document, ImmutableArray<Diagnostic>>();
             foreach (var documentAndDiagnostics in diagnostics.GroupBy(d => GetReportedDocument(d, treeToDocumentMap)))
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 var document = documentAndDiagnostics.Key;
                 if (!isGeneratedCode(document))
                 {
@@ -301,8 +306,10 @@ public virtual string GetFixAllTitle(FixAllContext fixAllContext)
             var builder = ImmutableDictionary.CreateBuilder<SyntaxTree, Document>();
             foreach (var project in projects)
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 foreach (var document in project.Documents)
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     var tree = await document.GetSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
                     builder.Add(tree, document);
                 }
@@ -368,6 +375,7 @@ public virtual async Task<Solution> TryMergeFixesAsync(Solution oldSolution, IEn
 
             foreach (var codeAction in codeActions)
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 // TODO: Parallelize GetChangedSolutionInternalAsync for codeActions
                 var changedSolution = await codeAction.GetChangedSolutionInternalAsync(cancellationToken).ConfigureAwait(false);
 
@@ -382,6 +390,7 @@ public virtual async Task<Solution> TryMergeFixesAsync(Solution oldSolution, IEn
 
                 foreach (var documentId in documentIdsWithChanges)
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     var document = changedSolution.GetDocument(documentId);
 
                     Document existingDocument;
@@ -411,6 +420,7 @@ public virtual async Task<Solution> TryMergeFixesAsync(Solution oldSolution, IEn
             var currentSolution = oldSolution;
             foreach (var kvp in changedDocumentsMap)
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 var document = kvp.Value;
                 if (document != null)
                 {
@@ -422,21 +432,23 @@ public virtual async Task<Solution> TryMergeFixesAsync(Solution oldSolution, IEn
             if (documentsToMergeMap != null)
             {
                 var mergedDocuments = new ConcurrentDictionary<DocumentId, SourceText>();
-                var documentsToMergeArray = documentsToMergeMap.ToImmutableArray();                
+                var documentsToMergeArray = documentsToMergeMap.ToImmutableArray();
                 var mergeTasks = new Task[documentsToMergeArray.Length];
                 for (int i = 0; i < documentsToMergeArray.Length; i++)
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     var kvp = documentsToMergeArray[i];
                     var documentId = kvp.Key;
                     var documentsToMerge = kvp.Value;
                     var oldDocument = oldSolution.GetDocument(documentId);
 
-                    mergeTasks[i] = Task.Run(async() =>
+                    mergeTasks[i] = Task.Run(async () =>
                     {
                         var appliedChanges = (await documentsToMerge[0].GetTextChangesAsync(oldDocument, cancellationToken).ConfigureAwait(false)).ToList();
 
                         foreach (var document in documentsToMerge.Skip(1))
                         {
+                            cancellationToken.ThrowIfCancellationRequested();
                             appliedChanges = await TryAddDocumentMergeChangesAsync(
                                 oldDocument,
                                 document,
@@ -454,6 +466,7 @@ public virtual async Task<Solution> TryMergeFixesAsync(Solution oldSolution, IEn
 
                 foreach (var kvp in mergedDocuments)
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     currentSolution = currentSolution.WithDocumentText(kvp.Key, kvp.Value);
                 }
             }
@@ -481,8 +494,10 @@ public virtual async Task<Solution> TryMergeFixesAsync(Solution oldSolution, IEn
             int cumulativeChangeIndex = 0;
             foreach (var change in await newDocument.GetTextChangesAsync(oldDocument, cancellationToken).ConfigureAwait(false))
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 while (cumulativeChangeIndex < cumulativeChanges.Count && cumulativeChanges[cumulativeChangeIndex].Span.End < change.Span.Start)
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
                     // Existing change that does not overlap with the current change in consideration
                     successfullyMergedChanges.Add(cumulativeChanges[cumulativeChangeIndex]);
                     cumulativeChangeIndex++;
@@ -522,6 +537,7 @@ public virtual async Task<Solution> TryMergeFixesAsync(Solution oldSolution, IEn
 
             while (cumulativeChangeIndex < cumulativeChanges.Count)
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 // Existing change that does not overlap with the current change in consideration
                 successfullyMergedChanges.Add(cumulativeChanges[cumulativeChangeIndex]);
                 cumulativeChangeIndex++;