Move usings on paste off the UI thread (#61092)

Utilize background work indicator for add usings on paste feature. Users should no longer see a blocking dialog with this feature.

src/EditorFeatures/CSharp/AddImports/CSharpAddImportsPasteCommandHandler.cs

@@ -8,6 +8,7 @@
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.VisualStudio.Commanding;
 using Microsoft.VisualStudio.Utilities;
 
@@ -29,8 +30,8 @@ internal class CSharpAddImportsPasteCommandHandler : AbstractAddImportsPasteComm
     {
         [ImportingConstructor]
         [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
-        public CSharpAddImportsPasteCommandHandler(IThreadingContext threadingContext, IGlobalOptionService globalOptions)
-            : base(threadingContext, globalOptions)
+        public CSharpAddImportsPasteCommandHandler(IThreadingContext threadingContext, IGlobalOptionService globalOptions, IAsynchronousOperationListenerProvider listnerProvider)
+            : base(threadingContext, globalOptions, listnerProvider)
         {
         }
 

src/EditorFeatures/Core/AddImports/AbstractAddImportsPasteCommandHandler.cs

@@ -3,19 +3,25 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.AddMissingImports;
 using Microsoft.CodeAnalysis.CodeCleanup;
 using Microsoft.CodeAnalysis.Completion;
+using Microsoft.CodeAnalysis.Editor.BackgroundWorkIndicator;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Editor.Shared.Options;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.Shared.Extensions;
+using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.Commanding;
 using Microsoft.VisualStudio.Text;
+using Microsoft.VisualStudio.Text.Editor;
 using Microsoft.VisualStudio.Text.Editor.Commanding.Commands;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.AddImport
 {
@@ -33,11 +39,16 @@ internal abstract class AbstractAddImportsPasteCommandHandler : IChainedCommandH
 
         private readonly IThreadingContext _threadingContext;
         private readonly IGlobalOptionService _globalOptions;
+        private readonly IAsynchronousOperationListener _listener;
 
-        public AbstractAddImportsPasteCommandHandler(IThreadingContext threadingContext, IGlobalOptionService globalOptions)
+        public AbstractAddImportsPasteCommandHandler(
+            IThreadingContext threadingContext,
+            IGlobalOptionService globalOptions,
+            IAsynchronousOperationListenerProvider listenerProvider)
         {
             _threadingContext = threadingContext;
             _globalOptions = globalOptions;
+            _listener = listenerProvider.GetListener(FeatureAttribute.AddImportsOnPaste);
         }
 
         public CommandState GetCommandState(PasteCommandArgs args, Func<CommandState> nextCommandHandler)
@@ -101,7 +112,6 @@ private void ExecuteCommandWorker(
 
             // Applying the post-paste snapshot to the tracking span gives us the span of pasted text.
             var snapshotSpan = trackingSpan.GetSpan(args.SubjectBuffer.CurrentSnapshot);
-            var textSpan = snapshotSpan.Span.ToTextSpan();
 
             var sourceTextContainer = args.SubjectBuffer.AsTextContainer();
             if (!Workspace.TryGetWorkspace(sourceTextContainer, out var workspace))
@@ -115,32 +125,53 @@ private void ExecuteCommandWorker(
                 return;
             }
 
-            using var _ = executionContext.OperationContext.AddScope(allowCancellation: true, DialogText);
-            var cancellationToken = executionContext.OperationContext.UserCancellationToken;
+            // We're showing our own UI, ensure the editor doesn't show anything itself.
+            executionContext.OperationContext.TakeOwnership();
+
+            var token = _listener.BeginAsyncOperation(nameof(ExecuteAsync));
+
+            ExecuteAsync(document, snapshotSpan, args.TextView)
+                .ReportNonFatalErrorAsync()
+                .CompletesAsyncOperation(token);
+        }
+
+        private async Task ExecuteAsync(Document document, SnapshotSpan snapshotSpan, ITextView textView)
+        {
+            _threadingContext.ThrowIfNotOnUIThread();
+
+            var indicatorFactory = document.Project.Solution.Workspace.Services.GetRequiredService<IBackgroundWorkIndicatorFactory>();
+            using var backgroundWorkContext = indicatorFactory.Create(
+                textView,
+                snapshotSpan,
+                DialogText,
+                cancelOnEdit: true,
+                cancelOnFocusLost: true);
+
+            var cancellationToken = backgroundWorkContext.UserCancellationToken;
 
             // We're going to log the same thing on success or failure since this blocks the UI thread. This measurement is 
             // intended to tell us how long we're blocking the user from typing with this action. 
             using var blockLogger = Logger.LogBlock(FunctionId.CommandHandler_Paste_ImportsOnPaste, KeyValueLogMessage.Create(LogType.UserAction), cancellationToken);
 
             var addMissingImportsService = document.GetRequiredLanguageService<IAddMissingImportsFeatureService>();
-#pragma warning disable VSTHRD102 // Implement internal logic asynchronously
-            var updatedDocument = _threadingContext.JoinableTaskFactory.Run(async () =>
-            {
-                var cleanupOptions = await document.GetCodeCleanupOptionsAsync(_globalOptions, cancellationToken).ConfigureAwait(false);
 
-                var options = new AddMissingImportsOptions(
-                    CleanupOptions: cleanupOptions,
-                    HideAdvancedMembers: _globalOptions.GetOption(CompletionOptionsStorage.HideAdvancedMembers, document.Project.Language));
+            var cleanupOptions = await document.GetCodeCleanupOptionsAsync(_globalOptions, cancellationToken).ConfigureAwait(false);
+
+            var options = new AddMissingImportsOptions(
+                CleanupOptions: cleanupOptions,
+                HideAdvancedMembers: _globalOptions.GetOption(CompletionOptionsStorage.HideAdvancedMembers, document.Project.Language));
+
+            var textSpan = snapshotSpan.Span.ToTextSpan();
+            var updatedDocument = await addMissingImportsService.AddMissingImportsAsync(document, textSpan, options, cancellationToken).ConfigureAwait(false);
 
-                return await addMissingImportsService.AddMissingImportsAsync(document, textSpan, options, cancellationToken).ConfigureAwait(false);
-            });
-#pragma warning restore VSTHRD102 // Implement internal logic asynchronously
             if (updatedDocument is null)
             {
                 return;
             }
 
-            workspace.TryApplyChanges(updatedDocument.Project.Solution);
+            // Required to switch back to the UI thread to call TryApplyChanges
+            await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+            document.Project.Solution.Workspace.TryApplyChanges(updatedDocument.Project.Solution);
         }
     }
 }

src/EditorFeatures/VisualBasic/AddImports/VisualBasicAddImportsOnPasteCommandHandler.vb

@@ -6,6 +6,7 @@ Imports System.ComponentModel.Composition
 Imports Microsoft.CodeAnalysis.AddImport
 Imports Microsoft.CodeAnalysis.Host.Mef
 Imports Microsoft.CodeAnalysis.Options
+Imports Microsoft.CodeAnalysis.Shared.TestHooks
 Imports Microsoft.VisualStudio.Commanding
 Imports Microsoft.VisualStudio.Utilities
 
@@ -22,8 +23,9 @@ Namespace Microsoft.CodeAnalysis.Editor.VisualBasic.AddImports
         <ImportingConstructor>
         <Obsolete(MefConstruction.ImportingConstructorMessage, True)>
         Public Sub New(threadingContext As [Shared].Utilities.IThreadingContext,
-                       globalOptions As IGlobalOptionService)
-            MyBase.New(threadingContext, globalOptions)
+                       globalOptions As IGlobalOptionService,
+                       listenerProvider As IAsynchronousOperationListenerProvider)
+            MyBase.New(threadingContext, globalOptions, listenerProvider)
         End Sub
 
         Public Overrides ReadOnly Property DisplayName As String = VBEditorResources.Add_Missing_Imports_on_Paste

src/VisualStudio/IntegrationTest/New.IntegrationTests/CSharp/CSharpAddMissingUsingsOnPaste.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.Threading;
 using System.Threading.Tasks;
 using System.Windows;
@@ -160,9 +161,14 @@ static void Main(string[] args)
 
         private async Task PasteAsync(string text, CancellationToken cancellationToken)
         {
+            var provider = await TestServices.Shell.GetComponentModelServiceAsync<IAsynchronousOperationListenerProvider>(HangMitigatingCancellationToken);
+            var waiter = (IAsynchronousOperationWaiter)provider.GetListener(FeatureAttribute.AddImportsOnPaste);
+
             await TestServices.Workspace.WaitForAllAsyncOperationsAsync(new[] { FeatureAttribute.Workspace, FeatureAttribute.SolutionCrawler }, cancellationToken);
             Clipboard.SetText(text);
             await TestServices.Shell.ExecuteCommandAsync(VSConstants.VSStd97CmdID.Paste, cancellationToken);
+
+            await waiter.ExpeditedWaitAsync();
         }
     }
 }

src/Workspaces/Core/Portable/Shared/TestHooks/FeatureAttribute.cs

@@ -57,5 +57,6 @@ internal static class FeatureAttribute
         public const string LanguageServer = nameof(LanguageServer);
         public const string ValueTracking = nameof(ValueTracking);
         public const string Workspace = nameof(Workspace);
+        public const string AddImportsOnPaste = nameof(AddImportsOnPaste);
     }
 }