Clean up blocking code in snippets (#76482)

Author: CyrusNajmabadi

src/VisualStudio/Core/Def/Snippets/SnippetExpansionClient.cs

@@ -705,8 +705,13 @@ private static XDocument CreateMethodCallSnippet(string methodName, bool include
     private void OnModelUpdated(object sender, ModelUpdatedEventsArgs e)
     {
         _threadingContext.ThrowIfNotOnUIThread();
+        _threadingContext.JoinableTaskFactory.Run(() => OnModelUpdatedAsync(e.NewModel, CancellationToken.None));
+    }
 
-        if (e.NewModel is null)
+    private async Task OnModelUpdatedAsync(
+        Model? newModel, CancellationToken cancellationToken)
+    {
+        if (newModel is null)
         {
             // Signature Help was dismissed, but it's possible for a user to bring it back with Ctrl+Shift+Space.
             // Leave the snippet session (if any) in its current state to allow it to process either a subsequent
@@ -721,7 +726,7 @@ private void OnModelUpdated(object sender, ModelUpdatedEventsArgs e)
             return;
         }
 
-        if (!e.NewModel.UserSelected && _state._method is not null)
+        if (!newModel.UserSelected && _state._method is not null)
         {
             // This was an implicit signature change which was not triggered by user pressing up/down, and we are
             // already showing an initialized argument completion snippet session, so avoid switching sessions.
@@ -739,13 +744,17 @@ private void OnModelUpdated(object sender, ModelUpdatedEventsArgs e)
         // TODO: The following blocks the UI thread without cancellation, but it only occurs when an argument value
         // completion session is active, which is behind an experimental feature flag.
         // https://github.com/dotnet/roslyn/issues/50634
-        var compilation = _threadingContext.JoinableTaskFactory.Run(() => document.Project.GetRequiredCompilationAsync(CancellationToken.None));
-        var newSymbolKey = (e.NewModel.SelectedItem as AbstractSignatureHelpProvider.SymbolKeySignatureHelpItem)?.SymbolKey ?? default;
-        var newSymbol = newSymbolKey.Resolve(compilation, cancellationToken: CancellationToken.None).GetAnySymbol();
+        var compilation = await document.Project.GetRequiredCompilationAsync(cancellationToken).ConfigureAwait(true);
+        var newSymbolKey = (newModel.SelectedItem as AbstractSignatureHelpProvider.SymbolKeySignatureHelpItem)?.SymbolKey ?? default;
+        var newSymbol = newSymbolKey.Resolve(compilation, cancellationToken: cancellationToken).GetAnySymbol();
         if (newSymbol is not IMethodSymbol method)
             return;
 
-        MoveToSpecificMethod(method, CancellationToken.None);
+        await MoveToSpecificMethodAsync(
+            document, method, cancellationToken).ConfigureAwait(true);
+
+        // Don't let the compilation drop as MoveToSpecificMethodAsync wants to get the semantic model for the document.
+        GC.KeepAlive(compilation);
     }
 
     private static async Task<ImmutableArray<ISymbol>> GetReferencedSymbolsToLeftOfCaretAsync(
@@ -771,10 +780,12 @@ private static async Task<ImmutableArray<ISymbol>> GetReferencedSymbolsToLeftOfC
     /// </summary>
     /// <param name="method">The currently-selected method in Signature Help.</param>
     /// <param name="cancellationToken">A cancellation token the operation may observe.</param>
-    public void MoveToSpecificMethod(IMethodSymbol method, CancellationToken cancellationToken)
+    public async Task MoveToSpecificMethodAsync(
+        Document document,
+        IMethodSymbol method,
+        CancellationToken cancellationToken)
     {
-        _threadingContext.ThrowIfNotOnUIThread();
-
+        await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
         if (ExpansionSession is null)
         {
             return;
@@ -804,14 +815,6 @@ public void MoveToSpecificMethod(IMethodSymbol method, CancellationToken cancell
             return;
         }
 
-        var document = SubjectBuffer.CurrentSnapshot.GetOpenDocumentInCurrentContextWithChanges();
-        if (document is null)
-        {
-            // Couldn't identify the current document
-            ExpansionSession.EndCurrentExpansion(fLeaveCaret: 1);
-            return;
-        }
-
         var textViewModel = TextView.TextViewModel;
         if (textViewModel == null)
         {
@@ -893,7 +896,7 @@ public void MoveToSpecificMethod(IMethodSymbol method, CancellationToken cancell
         }
 
         // Now compute the new arguments for the new call
-        var semanticModel = document.GetRequiredSemanticModelAsync(cancellationToken).AsTask().WaitAndGetResult(cancellationToken);
+        var semanticModel = await document.GetRequiredSemanticModelAsync(cancellationToken).ConfigureAwait(true);
         var position = SubjectBuffer.CurrentSnapshot.GetPosition(adjustedTextSpan.iStartLine, adjustedTextSpan.iStartIndex);
 
         foreach (var parameter in method.Parameters)
@@ -903,7 +906,7 @@ public void MoveToSpecificMethod(IMethodSymbol method, CancellationToken cancell
             foreach (var provider in GetArgumentProviders(document.Project.Solution.Workspace))
             {
                 var context = new ArgumentContext(provider, semanticModel, position, parameter, value, cancellationToken);
-                _threadingContext.JoinableTaskFactory.Run(() => provider.ProvideArgumentAsync(context));
+                await provider.ProvideArgumentAsync(context).ConfigureAwait(true);
 
                 if (context.DefaultValue is not null)
                 {