Merge pull request #18863 from CyrusNajmabadi/lightBulbAsyncGetSelection

Acquire the user's selection up-front if we're on the UI thread for lightbulbs.

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

@@ -42,6 +42,8 @@ private class SuggestedActionsSource : ForegroundThreadAffinitizedObject, ISugge
             private Workspace _workspace;
             private int _lastSolutionVersionReported;
 
+            public event EventHandler<EventArgs> SuggestedActionsChanged;
+
             public SuggestedActionsSource(SuggestedActionsSourceProvider owner, ITextView textView, ITextBuffer textBuffer)
             {
                 _owner = owner;
@@ -64,7 +66,37 @@ public SuggestedActionsSource(SuggestedActionsSourceProvider owner, ITextView te
                 _registration.WorkspaceChanged += OnWorkspaceChanged;
             }
 
-            public event EventHandler<EventArgs> SuggestedActionsChanged;
+            public void Dispose()
+            {
+                if (_owner != null)
+                {
+                    var updateSource = (IDiagnosticUpdateSource)_owner._diagnosticService;
+                    updateSource.DiagnosticsUpdated -= OnDiagnosticsUpdated;
+                }
+
+                if (_workspace != null)
+                {
+                    _workspace.DocumentActiveContextChanged -= OnActiveContextChanged;
+                }
+
+                if (_registration != null)
+                {
+                    _registration.WorkspaceChanged -= OnWorkspaceChanged;
+                }
+
+                if (_textView != null)
+                {
+                    _textView.Closed -= OnTextViewClosed;
+                }
+
+                _owner = null;
+                _workspace = null;
+                _registration = null;
+                _textView = null;
+                _subjectBuffer = null;
+            }
+
+            private bool IsDisposed => _subjectBuffer == null;
 
             public bool TryGetTelemetryId(out Guid telemetryId)
             {
@@ -111,6 +143,11 @@ public bool TryGetTelemetryId(out Guid telemetryId)
             {
                 AssertIsForeground();
 
+                if (IsDisposed)
+                {
+                    return null;
+                }
+
                 using (Logger.LogBlock(FunctionId.SuggestedActions_GetSuggestedActions, cancellationToken))
                 {
                     var document = range.Snapshot.GetOpenDocumentInCurrentContextWithChanges();
@@ -529,7 +566,7 @@ private static SuggestedActionSetPriority GetSuggestedActionSetPriority(CodeActi
                     requestedActionCategories.Contains(PredefinedSuggestedActionCategoryNames.Refactoring))
                 {
                     // Get the selection while on the UI thread.
-                    var selection = TryGetCodeRefactoringSelection(_subjectBuffer, _textView, range);
+                    var selection = TryGetCodeRefactoringSelection(range);
                     if (!selection.HasValue)
                     {
                         // this is here to fail test and see why it is failed.
@@ -593,20 +630,60 @@ private static SuggestedActionSetPriority GetSuggestedActionSetPriority(CodeActi
                 SnapshotSpan range,
                 CancellationToken cancellationToken)
             {
-                // Explicitly hold onto below fields in locals and use these locals throughout this code path to avoid crashes
-                // if these fields happen to be cleared by Dispose() below. This is required since this code path involves
-                // code that can run asynchronously from background thread.
-                var view = _textView;
-                var buffer = _subjectBuffer;
                 var provider = _owner;
 
-                if (view == null || buffer == null || provider == null)
+                if (IsDisposed)
                 {
+                    // We've already been disposed.  No point in continuing.
                     return false;
                 }
 
-                using (var asyncToken = provider.OperationListener.BeginAsyncOperation("HasSuggestedActionsAsync"))
-                {
+                using (var asyncToken = _owner.OperationListener.BeginAsyncOperation("HasSuggestedActionsAsync"))
+                {
+                    // Acquire the user's selection up front, before we do any async work, so that we can
+                    // directly grab it when we're on the UI thread.  That way we don't have any reentrancy
+                    // blocking concerns if VS wants to block on this call (for example, if the user 
+                    // explicitly invokes the 'show smart tag' command).
+                    //
+                    // This work must happen on the UI thread as it needs to access the _textView's mutable
+                    // state.
+                    //
+                    // Note: we may be called in one of two VS scenarios:
+                    //      1) User has moved caret to a new line.  In this case VS will call into us in the
+                    //         bg to see if we have any suggested actions for this line.  In order to figure
+                    //         this out, we need to see what selectoin the user has (for refactorings), which
+                    //         necessitates going back to the fg.
+                    //
+                    //      2) User moves to a line and immediately hits ctrl-dot.  In this case, on the UI
+                    //         thread VS will kick us off and then immediately block to get the results so
+                    //         that they can expand the lightbulb.  In this case we cannot do BG work first,
+                    //         then call back into the UI thread to try to get the user selection.  This will
+                    //         deadlock as the UI thread is blocked on us.  
+                    //
+                    // There are two solution to '2'.  Either introduce reentrancy (which we really don't
+                    // like to do), or just ensure that we acquire and get the users selection up front.
+                    // This means that when we're called from the UI therad, we never try to go back to the
+                    // UI thread.
+                    TextSpan? selection = null;
+                    if (IsForeground())
+                    {
+                        selection = TryGetCodeRefactoringSelection(range);
+                    }
+                    else
+                    {
+                        await InvokeBelowInputPriority(() =>
+                        {
+                            // Make sure we were not disposed between kicking off this work and getting
+                            // to this point.
+                            if (IsDisposed)
+                            {
+                                return;
+                            }
+
+                            selection = TryGetCodeRefactoringSelection(range);
+                        }).ConfigureAwait(false);
+                    }
+
                     var document = range.Snapshot.GetOpenDocumentInCurrentContextWithChanges();
                     if (document == null)
                     {
@@ -617,7 +694,7 @@ private static SuggestedActionSetPriority GetSuggestedActionSetPriority(CodeActi
 
                     return
                         await HasFixesAsync(provider, document, range, cancellationToken).ConfigureAwait(false) ||
-                        await HasRefactoringsAsync(provider, document, buffer, view, range, cancellationToken).ConfigureAwait(false);
+                        await HasRefactoringsAsync(provider, document, selection, cancellationToken).ConfigureAwait(false);
                 }
             }
 
@@ -658,11 +735,16 @@ private static SuggestedActionSetPriority GetSuggestedActionSetPriority(CodeActi
             private async Task<bool> HasRefactoringsAsync(
                 SuggestedActionsSourceProvider provider,
                 Document document,
-                ITextBuffer buffer,
-                ITextView view,
-                SnapshotSpan range,
+                TextSpan? selection,
                 CancellationToken cancellationToken)
             {
+                if (!selection.HasValue)
+                {
+                    // this is here to fail test and see why it is failed.
+                    Trace.WriteLine("given range is not current");
+                    return false;
+                }
+
                 var workspace = document.Project.Solution.Workspace;
                 var supportsFeatureService = workspace.Services.GetService<IDocumentSupportsFeatureService>();
 
@@ -670,28 +752,6 @@ private static SuggestedActionSetPriority GetSuggestedActionSetPriority(CodeActi
                     provider._codeRefactoringService != null &&
                     supportsFeatureService.SupportsRefactorings(document))
                 {
-                    TextSpan? selection = null;
-                    if (IsForeground())
-                    {
-                        // This operation needs to happen on UI thread because it needs to access textView.Selection.
-                        selection = TryGetCodeRefactoringSelection(buffer, view, range);
-                    }
-                    else
-                    {
-                        await InvokeBelowInputPriority(() =>
-                        {
-                                // This operation needs to happen on UI thread because it needs to access textView.Selection.
-                                selection = TryGetCodeRefactoringSelection(buffer, view, range);
-                        }).ConfigureAwait(false);
-                    }
-
-                    if (!selection.HasValue)
-                    {
-                        // this is here to fail test and see why it is failed.
-                        Trace.WriteLine("given range is not current");
-                        return false;
-                    }
-
                     return await Task.Run(
                         () => provider._codeRefactoringService.HasRefactoringsAsync(
                             document, selection.Value, cancellationToken),
@@ -701,11 +761,14 @@ private static SuggestedActionSetPriority GetSuggestedActionSetPriority(CodeActi
                 return false;
             }
 
-            private static TextSpan? TryGetCodeRefactoringSelection(ITextBuffer buffer, ITextView view, SnapshotSpan range)
+            private TextSpan? TryGetCodeRefactoringSelection(SnapshotSpan range)
             {
-                var selectedSpans = view.Selection.SelectedSpans
-                    .SelectMany(ss => view.BufferGraph.MapDownToBuffer(ss, SpanTrackingMode.EdgeExclusive, buffer))
-                    .Where(ss => !view.IsReadOnlyOnSurfaceBuffer(ss))
+                this.AssertIsForeground();
+                Debug.Assert(!this.IsDisposed);
+
+                var selectedSpans = _textView.Selection.SelectedSpans
+                    .SelectMany(ss => _textView.BufferGraph.MapDownToBuffer(ss, SpanTrackingMode.EdgeExclusive, _subjectBuffer))
+                    .Where(ss => !_textView.IsReadOnlyOnSurfaceBuffer(ss))
                     .ToList();
 
                 // We only support refactorings when there is a single selection in the document.
@@ -796,39 +859,6 @@ private void OnSuggestedActionsChanged(Workspace currentWorkspace, DocumentId cu
 
                 Volatile.Write(ref _lastSolutionVersionReported, solutionVersion);
             }
-
-            public void Dispose()
-            {
-                if (_owner != null)
-                {
-                    var updateSource = (IDiagnosticUpdateSource)_owner._diagnosticService;
-                    updateSource.DiagnosticsUpdated -= OnDiagnosticsUpdated;
-                    _owner = null;
-                }
-
-                if (_workspace != null)
-                {
-                    _workspace.DocumentActiveContextChanged -= OnActiveContextChanged;
-                    _workspace = null;
-                }
-
-                if (_registration != null)
-                {
-                    _registration.WorkspaceChanged -= OnWorkspaceChanged;
-                    _registration = null;
-                }
-
-                if (_textView != null)
-                {
-                    _textView.Closed -= OnTextViewClosed;
-                    _textView = null;
-                }
-
-                if (_subjectBuffer != null)
-                {
-                    _subjectBuffer = null;
-                }
-            }
         }
     }
 }
\ No newline at end of file