Simplify model computation threading (#76461)

Author: CyrusNajmabadi

src/EditorFeatures/Core.Wpf/SignatureHelp/Controller.Session.cs

@@ -15,7 +15,7 @@ internal partial class Controller
         internal partial class Session : Session<Controller, Model, ISignatureHelpPresenterSession>
         {
             public Session(Controller controller, ISignatureHelpPresenterSession presenterSession)
-                : base(controller, new ModelComputation<Model>(controller.ThreadingContext, controller, TaskScheduler.Default), presenterSession)
+                : base(controller, new ModelComputation<Model>(controller.ThreadingContext, controller), presenterSession)
             {
                 this.PresenterSession.ItemSelected += OnPresenterSessionItemSelected;
             }

src/EditorFeatures/Core.Wpf/SignatureHelp/Controller.Session_ComputeModel.cs

@@ -13,10 +13,10 @@
 using Microsoft.CodeAnalysis.ErrorReporting;
 using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.CodeAnalysis.LanguageService;
-using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.SignatureHelp;
 using Microsoft.VisualStudio.Text;
+using Microsoft.VisualStudio.Threading;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.IntelliSense.SignatureHelp
@@ -37,9 +37,8 @@ public void ComputeModel(
                 // If we've already computed a model, then just use that.  Otherwise, actually
                 // compute a new model and send that along.
                 Computation.ChainTaskAndNotifyControllerWhenFinished(
-                    (model, cancellationToken) => ComputeModelInBackgroundAsync(
-                        model, providers, caretPosition, disconnectedBufferGraph,
-                        triggerInfo, cancellationToken));
+                    (currentModel, cancellationToken) => ComputeModelInBackgroundAsync(
+                        currentModel, providers, caretPosition, disconnectedBufferGraph, triggerInfo, cancellationToken));
             }
 
             private async Task<Model> ComputeModelInBackgroundAsync(

src/EditorFeatures/Core.Wpf/SignatureHelp/Controller.Session_SetModelSelectedItem.cs

@@ -2,11 +2,11 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-#nullable disable
-
 using System;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.SignatureHelp;
+using Microsoft.VisualStudio.Threading;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.IntelliSense.SignatureHelp
@@ -20,20 +20,17 @@ private void SetModelExplicitlySelectedItem(Func<Model, SignatureHelpItem> selec
                 this.Computation.ThreadingContext.ThrowIfNotOnUIThread();
 
                 Computation.ChainTaskAndNotifyControllerWhenFinished(
-                    model => SetModelExplicitlySelectedItemInBackground(model, selector),
+                    (model, cancellationToken) => Task.FromResult(SetModelExplicitlySelectedItemInBackground(model, selector)),
                     updateController: false);
             }
 
-            private Model SetModelExplicitlySelectedItemInBackground(
-                Model model,
+            private Model? SetModelExplicitlySelectedItemInBackground(
+                Model? model,
                 Func<Model, SignatureHelpItem> selector)
             {
                 this.Computation.ThreadingContext.ThrowIfNotOnBackgroundThread();
-
                 if (model == null)
-                {
                     return null;
-                }
 
                 var selectedItem = selector(model);
                 Contract.ThrowIfFalse(model.Items.Contains(selectedItem));

src/EditorFeatures/Core/IntelliSense/ModelComputation.cs

@@ -11,6 +11,7 @@
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
+using Microsoft.VisualStudio.Threading;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.IntelliSense;
@@ -32,16 +33,6 @@ internal class ModelComputation<TModel> where TModel : class
     #region Fields that can only be accessed from the foreground thread
 
     private readonly IController<TModel> _controller;
-    private readonly TaskScheduler __taskScheduler;
-
-    private TaskScheduler _taskScheduler
-    {
-        get
-        {
-            ThreadingContext.ThrowIfNotOnUIThread();
-            return __taskScheduler;
-        }
-    }
 
     private readonly CancellationTokenSource _stopTokenSource;
 
@@ -57,12 +48,10 @@ private TaskScheduler _taskScheduler
 
     public ModelComputation(
         IThreadingContext threadingContext,
-        IController<TModel> controller,
-        TaskScheduler computationTaskScheduler)
+        IController<TModel> controller)
     {
         ThreadingContext = threadingContext;
         _controller = controller;
-        __taskScheduler = computationTaskScheduler;
 
         _stopTokenSource = new CancellationTokenSource();
         _stopCancellationToken = _stopTokenSource.Token;
@@ -119,13 +108,6 @@ public virtual void Stop()
         _notifyControllerTask = _lastTask = SpecializedTasks.Null<TModel>();
     }
 
-    public void ChainTaskAndNotifyControllerWhenFinished(
-            Func<TModel, TModel> transformModel,
-            bool updateController = true)
-    {
-        ChainTaskAndNotifyControllerWhenFinished((m, c) => Task.FromResult(transformModel(m)), updateController);
-    }
-
     public void ChainTaskAndNotifyControllerWhenFinished(
         Func<TModel, CancellationToken, Task<TModel>> transformModelAsync,
         bool updateController = true)
@@ -140,9 +122,8 @@ public void ChainTaskAndNotifyControllerWhenFinished(
         var asyncToken = _controller.BeginAsyncOperation();
 
         // Create the task that will actually run the transformation step.
-        var nextTask = _lastTask.SafeContinueWithFromAsync(
-            t => transformModelAsync(t.Result, _stopCancellationToken),
-            _stopCancellationToken, TaskContinuationOptions.OnlyOnRanToCompletion, _taskScheduler);
+        var nextTask = TransformModelAsync(_lastTask);// transformModelAsync(_lastTask, _stopCancellationToken);
+        nextTask.ReportNonFatalErrorAsync();
 
         // The next task is now the last task in the chain.
         _lastTask = nextTask;
@@ -174,6 +155,14 @@ public void ChainTaskAndNotifyControllerWhenFinished(
         // When we've notified the controller of our result, we consider the async operation
         // to be completed.
         _notifyControllerTask.CompletesAsyncOperation(asyncToken);
+
+        async Task<TModel> TransformModelAsync(Task<TModel> lastTask)
+        {
+            // Ensure we're on the BG before doing any model transformation work.
+            await TaskScheduler.Default;
+            var model = await lastTask.ConfigureAwait(false);
+            return await transformModelAsync(model, _stopCancellationToken).ConfigureAwait(false);
+        }
     }
 
     private void OnModelUpdated(TModel result, bool updateController)

src/EditorFeatures/Test2/IntelliSense/ModelTests.vb

@@ -6,21 +6,22 @@ Imports System.Threading
 Imports Microsoft.CodeAnalysis.Editor.Implementation.IntelliSense
 Imports Microsoft.CodeAnalysis.Editor.Shared.Utilities
 Imports Microsoft.CodeAnalysis.Shared.TestHooks
+Imports Microsoft.VisualStudio.Threading
 Imports Moq
 
 Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
 
     <UseExportProvider>
     <Trait(Traits.Feature, Traits.Features.DebuggingIntelliSense)>
-    Public Class ModelTests
+    Public NotInheritable Class ModelTests
         Public Class Model
         End Class
 
         Private Class TestModelComputation
             Inherits ModelComputation(Of Model)
 
             Public Sub New(threadingContext As IThreadingContext, controller As IController(Of Model))
-                MyBase.New(threadingContext, controller, TaskScheduler.Default)
+                MyBase.New(threadingContext, controller)
             End Sub
 
             Friend Shared Function Create(threadingContext As IThreadingContext, Optional controller As IController(Of Model) = Nothing) As TestModelComputation
@@ -44,7 +45,7 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
             controller.Setup(Function(c) c.BeginAsyncOperation("", Nothing, It.IsAny(Of String), It.IsAny(Of Integer))).Returns(EmptyAsyncToken.Instance)
             Dim modelComputation = TestModelComputation.Create(threadingContext, controller:=controller.Object)
 
-            modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m) m)
+            modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m, c) Task.FromResult(m))
 
             controller.Verify(Sub(c) c.BeginAsyncOperation(
                                   It.IsAny(Of String),
@@ -62,7 +63,7 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
             controller.Setup(Sub(c) c.OnModelUpdated(model, True))
             Dim modelComputation = TestModelComputation.Create(threadingContext, controller:=controller.Object)
 
-            modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m) model)
+            modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m, c) Task.FromResult(model))
             modelComputation.Wait()
 
             controller.Verify(Sub(c) c.OnModelUpdated(model, True))
@@ -79,12 +80,12 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
             Dim gate = New Object
 
             Monitor.Enter(gate)
-            modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m)
+            modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m, c)
                                                                           SyncLock gate
-                                                                              Return Nothing
+                                                                              Return Task.FromResult(Of Model)(Nothing)
                                                                           End SyncLock
                                                                       End Function)
-            modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m) model)
+            modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m, c) Task.FromResult(model))
             Monitor.Exit(gate)
             modelComputation.Wait()
 
@@ -109,12 +110,13 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
 
             token.Setup(Sub(t) t.Dispose()).Callback(Sub() checkpoint3.Release())
 
-            modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m, c)
-                                                                          checkpoint1.Release()
-                                                                          checkpoint2.Task.Wait()
-                                                                          c.ThrowIfCancellationRequested()
-                                                                          Return Task.FromResult(model)
-                                                                      End Function)
+            modelComputation.ChainTaskAndNotifyControllerWhenFinished(
+                Function(model1, c)
+                    checkpoint1.Release()
+                    checkpoint2.Task.Wait()
+                    c.ThrowIfCancellationRequested()
+                    Return Task.FromResult(model1)
+                End Function)
             Await checkpoint1.Task
             modelComputation.Stop()
             checkpoint2.Release()

src/EditorFeatures/Test2/IntelliSense/SessionTests.vb

@@ -24,7 +24,7 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
             controller.Setup(Sub(c) c.StopModelComputation())
             Dim session = New Session(Of IController(Of Model), Model, IIntelliSensePresenterSession)(
                 controller.Object,
-                New ModelComputation(Of Model)(threadingContext, controller.Object, TaskScheduler.Default),
+                New ModelComputation(Of Model)(threadingContext, controller.Object),
                 presenter.Object)
 
             presenter.Raise(Sub(p) AddHandler p.Dismissed, Nothing, New EventArgs())
@@ -41,7 +41,7 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
             Dim controller = New Mock(Of IController(Of Model))(MockBehavior.Strict)
             Dim session = New Session(Of IController(Of Model), Model, IIntelliSensePresenterSession)(
                 controller.Object,
-                New ModelComputation(Of Model)(threadingContext, controller.Object, TaskScheduler.Default),
+                New ModelComputation(Of Model)(threadingContext, controller.Object),
                 presenter.Object)
 
             session.Stop()