Ensure rename cancellation is synchronous

src/EditorFeatures/Core/IInlineRenameSession.cs

@@ -44,7 +44,8 @@ internal InlineRenameSessionInfo(IInlineRenameSession session)
 internal interface IInlineRenameSession
 {
     /// <summary>
-    /// Cancels the rename session, and undoes any edits that had been performed by the session.
+    /// Cancels the rename session, and undoes any edits that had been performed by the session.  Must be called on the
+    /// UI thread.
     /// </summary>
     void Cancel();
 

src/EditorFeatures/Core/InlineRename/InlineRenameSession.cs

@@ -651,26 +651,35 @@ public void Cancel()
     {
         _threadingContext.ThrowIfNotOnUIThread();
 
-        // This wait is safe.  We are not passing the async callback to DismissUIAndRollbackEditsAndEndRenameSessionAsync.
-        // So everything in that method will happen synchronously.
-        DismissUIAndRollbackEditsAndEndRenameSessionAsync(
-            RenameLogMessage.UserActionOutcome.Canceled, previewChanges: false).Wait();
+        DismissUIAndRollbackEditsAndEndRenameSession_MustBeCalledOnUIThread(
+            RenameLogMessage.UserActionOutcome.Canceled, previewChanges: false);
     }
 
-    private async Task DismissUIAndRollbackEditsAndEndRenameSessionAsync(
+    private void DismissUIAndRollbackEditsAndEndRenameSession_MustBeCalledOnUIThread(
+        RenameLogMessage.UserActionOutcome outcome,
+        bool previewChanges)
+    {
+        _threadingContext.ThrowIfNotOnUIThread();
+
+        DismissUIAndRollbackEditsAndEndRenameSession_MustBeCalledOnUIThread(
+            outcome, previewChanges, applyChangesOpt: null);
+    }
+
+    /// <summary>
+    /// Dismisses the UI, rolls back any edits, and ends the rename session.
+    /// </summary>
+    private void DismissUIAndRollbackEditsAndEndRenameSession_MustBeCalledOnUIThread(
         RenameLogMessage.UserActionOutcome outcome,
         bool previewChanges,
-        Func<Task> finalCommitAction = null)
+        Action applyChangesOpt)
     {
         // Note: this entire sequence of steps is not cancellable.  We must perform it all to get back to a correct
         // state for all the editors the user is interacting with.
         var cancellationToken = CancellationToken.None;
-        await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+        _threadingContext.ThrowIfNotOnUIThread();
 
         if (_dismissed)
-        {
             return;
-        }
 
         _dismissed = true;
 
@@ -685,11 +694,7 @@ private async Task DismissUIAndRollbackEditsAndEndRenameSessionAsync(
         _keepAliveSession.Dispose();
 
         // Perform the actual commit step if we've been asked to.
-        if (finalCommitAction != null)
-        {
-            // ConfigureAwait(true) so we come back to the UI thread to finish work.
-            await finalCommitAction().ConfigureAwait(true);
-        }
+        applyChangesOpt?.Invoke();
 
         // Log the result so we know how well rename is going in practice.
         LogRenameSession(outcome, previewChanges);
@@ -827,8 +832,9 @@ private async Task<bool> CommitWorkerAsync(bool previewChanges, bool canUseBackg
         }
         catch (OperationCanceledException)
         {
-            await DismissUIAndRollbackEditsAndEndRenameSessionAsync(
-                RenameLogMessage.UserActionOutcome.Canceled | RenameLogMessage.UserActionOutcome.Committed, previewChanges).ConfigureAwait(false);
+            // We've used CA(true) consistently in this method.  So we should always be on the UI thread.
+            DismissUIAndRollbackEditsAndEndRenameSession_MustBeCalledOnUIThread(
+                RenameLogMessage.UserActionOutcome.Canceled | RenameLogMessage.UserActionOutcome.Committed, previewChanges);
             return false;
         }
 
@@ -867,56 +873,89 @@ private async Task CommitCoreAsync(IUIThreadOperationContext operationContext, b
             }
 
             // The user hasn't canceled by now, so we're done waiting for them. Off to rename!
-            using var _ = operationContext.AddScope(allowCancellation: false, EditorFeaturesResources.Updating_files);
+            using var _1 = operationContext.AddScope(allowCancellation: true, EditorFeaturesResources.Updating_files);
+
+            await TaskScheduler.Default;
+
+            // While on the background, attempt to figure out the actual changes to make to the workspace's current solution.
+            var documentChanges = await CalculateFinalDocumentChangesAsync(_baseSolution, newSolution, cancellationToken).ConfigureAwait(false);
+
+            // Now jump back to the UI thread to actually apply those changes.
+            await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+
+            // We're about to make irrevocable changes to the workspace and UI.  We're no longer cancellable at this point.
+            using var _2 = operationContext.AddScope(allowCancellation: false, EditorFeaturesResources.Updating_files);
+            cancellationToken = CancellationToken.None;
 
-            await DismissUIAndRollbackEditsAndEndRenameSessionAsync(
+            // Dismiss the rename UI and rollback any linked edits made.
+            DismissUIAndRollbackEditsAndEndRenameSession_MustBeCalledOnUIThread(
                 RenameLogMessage.UserActionOutcome.Committed, previewChanges,
-                async () =>
+                applyChangesOpt: () =>
                 {
-                    var error = await TryApplyRenameAsync(newSolution, cancellationToken).ConfigureAwait(false);
-
+                    // Now try to apply that change we computed to the workspace.
+                    var error = TryApplyRename(documentChanges);
                     if (error is not null)
                     {
-                        await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
                         var notificationService = Workspace.Services.GetService<INotificationService>();
                         notificationService.SendNotification(
                             error.Value.message, EditorFeaturesResources.Rename_Symbol, error.Value.severity);
                     }
-                }).ConfigureAwait(false);
+                });
         }
-    }
 
-    /// <summary>
-    /// Returns non-null error message if renaming fails.
-    /// </summary>
-    private async Task<(NotificationSeverity severity, string message)?> TryApplyRenameAsync(
-        Solution newSolution, CancellationToken cancellationToken)
-    {
-        var changes = _baseSolution.GetChanges(newSolution);
-        var changedDocumentIDs = changes.GetProjectChanges().SelectMany(c => c.GetChangedDocuments()).ToList();
+        async Task<ImmutableArray<(DocumentId documentId, string newName, SyntaxNode newRoot, SourceText newText)>> CalculateFinalDocumentChangesAsync(
+            Solution newSolution, CancellationToken cancellationToken)
+        {
+            var changes = _baseSolution.GetChanges(newSolution);
+            var changedDocumentIDs = changes.GetProjectChanges().SelectManyAsArray(c => c.GetChangedDocuments());
+
+            using var _ = PooledObjects.ArrayBuilder<(DocumentId documentId, string newName, SyntaxNode newRoot, SourceText newText)>.GetInstance(out var result);
+
+            foreach (var documentId in changes.GetProjectChanges().SelectMany(c => c.GetChangedDocuments()))
+            {
+                // If the document supports syntax tree, then create the new solution from the updated syntax root.
+                // This should ensure that annotations are preserved, and prevents the solution from having to reparse
+                // documents when we already have the trees for them.  If we don't support syntax, then just use the
+                // text of the document.
+                var newDocument = newSolution.GetDocument(documentId);
+
+                result.Add(newDocument.SupportsSyntaxTree
+                    ? (documentId, newDocument.Name, await newDocument.GetRequiredSyntaxRootAsync(cancellationToken).ConfigureAwait(false), newText: null)
+                    : (documentId, newDocument.Name, newRoot: null, await newDocument.GetTextAsync(cancellationToken).ConfigureAwait(false)));
+            }
 
-        // Go to the background thread for initial calculation of the final solution
-        await TaskScheduler.Default;
-        var finalSolution = CalculateFinalSolutionSynchronously(newSolution, newSolution.Workspace, changedDocumentIDs, cancellationToken);
+            return result.ToImmutableAndClear();
+        }
+    }
 
-        await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+    // Returns non-null error message if renaming fails.
+    private (NotificationSeverity severity, string message)? TryApplyRename(
+        ImmutableArray<(DocumentId documentId, string newName, SyntaxNode newRoot, SourceText newText)> documentChanges)
+    {
+        _threadingContext.ThrowIfNotOnUIThread();
 
         using var undoTransaction = Workspace.OpenGlobalUndoTransaction(EditorFeaturesResources.Inline_Rename);
 
-        if (!RenameInfo.TryOnBeforeGlobalSymbolRenamed(Workspace, changedDocumentIDs, this.ReplacementText))
+        if (!RenameInfo.TryOnBeforeGlobalSymbolRenamed(Workspace, documentChanges.SelectAsArray(t => t.documentId), this.ReplacementText))
             return (NotificationSeverity.Error, EditorFeaturesResources.Rename_operation_was_cancelled_or_is_not_valid);
 
-        if (!Workspace.TryApplyChanges(finalSolution))
+        // Grab the workspace's current solution, and make the document changes to it we computed.
+        var finalSolution = Workspace.CurrentSolution;
+        foreach (var (documentId, newName, newRoot, newText) in documentChanges)
         {
-            // If the workspace changed in TryOnBeforeGlobalSymbolRenamed retry, this prevents rename from failing for cases
-            // where text changes to other files or workspace state change doesn't impact the text changes being applied. 
-            Logger.Log(FunctionId.Rename_TryApplyRename_WorkspaceChanged, message: null, LogLevel.Information);
-            finalSolution = CalculateFinalSolutionSynchronously(newSolution, Workspace, changedDocumentIDs, cancellationToken);
+            finalSolution = newRoot != null
+                ? finalSolution.WithDocumentSyntaxRoot(documentId, newRoot)
+                : finalSolution.WithDocumentText(documentId, newText);
 
-            if (!Workspace.TryApplyChanges(finalSolution))
-                return (NotificationSeverity.Error, EditorFeaturesResources.Rename_operation_could_not_complete_due_to_external_change_to_workspace);
+            finalSolution = finalSolution.WithDocumentName(documentId, newName);
         }
 
+        // Now actually go and apply the changes to the workspace.  We expect this to succeed as we're on the UI
+        // thread, and nothing else should have been able to make a change to workspace since we we grabbed its
+        // current solution and forked it.
+        if (!Workspace.TryApplyChanges(finalSolution))
+            return (NotificationSeverity.Error, EditorFeaturesResources.Rename_operation_could_not_complete_due_to_external_change_to_workspace);
+
         try
         {
             // Since rename can apply file changes as well, and those file
@@ -927,8 +966,7 @@ await DismissUIAndRollbackEditsAndEndRenameSessionAsync(
 
             var finalChangedIds = finalChanges
                 .GetProjectChanges()
-                .SelectMany(c => c.GetChangedDocuments().Concat(c.GetAddedDocuments()))
-                .ToList();
+                .SelectManyAsArray(c => c.GetChangedDocuments().Concat(c.GetAddedDocuments()));
 
             if (!RenameInfo.TryOnAfterGlobalSymbolRenamed(Workspace, finalChangedIds, this.ReplacementText))
                 return (NotificationSeverity.Information, EditorFeaturesResources.Rename_operation_was_not_properly_completed_Some_file_might_not_have_been_updated);
@@ -941,36 +979,6 @@ await DismissUIAndRollbackEditsAndEndRenameSessionAsync(
             // always able to undo anything any other external listener did.
             undoTransaction.Commit();
         }
-
-        static Solution CalculateFinalSolutionSynchronously(Solution newSolution, Workspace workspace, List<DocumentId> changedDocumentIDs, CancellationToken cancellationToken)
-        {
-            var finalSolution = workspace.CurrentSolution;
-            foreach (var id in changedDocumentIDs)
-            {
-                // If the document supports syntax tree, then create the new solution from the
-                // updated syntax root.  This should ensure that annotations are preserved, and
-                // prevents the solution from having to reparse documents when we already have
-                // the trees for them.  If we don't support syntax, then just use the text of
-                // the document.
-                var newDocument = newSolution.GetDocument(id);
-
-                if (newDocument.SupportsSyntaxTree)
-                {
-                    var root = newDocument.GetRequiredSyntaxRootSynchronously(cancellationToken);
-                    finalSolution = finalSolution.WithDocumentSyntaxRoot(id, root);
-                }
-                else
-                {
-                    var newText = newDocument.GetTextSynchronously(cancellationToken);
-                    finalSolution = finalSolution.WithDocumentText(id, newText);
-                }
-
-                // Make sure to include any document rename as well
-                finalSolution = finalSolution.WithDocumentName(id, newDocument.Name);
-            }
-
-            return finalSolution;
-        }
     }
 
     internal bool TryGetContainingEditableSpan(SnapshotPoint point, out SnapshotSpan editableSpan)