Further reduce thread switches during package load. (#77602)

Further reduce thread switches during package load.

The Rolsyn package load sequence is quite confusing, and currently does quite a few thread switches. Trying re-order statements is difficult, and fraught with danger if one of the calls is entered or returns on an unexpected thread. For example, one component created during load, the ColorSchemeApplier, needs to be initialized during this sequence, but does up to three thread switches during it's initialization. Every time work is added to the package load sequence, we run a pretty high risk of adding new thread switches during this critical period.

Instead, this change allows packages to add work to main thread or background thread work queues. Items in each of these queues are batched together and can use a single thread switch for the entire batch. These queues can be added to during processing, allowing packages to "discover" new work that needs to be processed as they proceed.

Author: ToddGrun

src/VisualStudio/CSharp/Impl/CSharpPackage.cs

@@ -58,12 +58,17 @@ internal sealed class CSharpPackage : AbstractPackage<CSharpPackage, CSharpLangu
         private ObjectBrowserLibraryManager? _libraryManager;
         private uint _libraryManagerCookie;
 
-        protected override async Task InitializeAsync(CancellationToken cancellationToken, IProgress<ServiceProgressData> progress)
+        protected override void RegisterInitializationWork(PackageRegistrationTasks packageRegistrationTasks)
+        {
+            base.RegisterInitializationWork(packageRegistrationTasks);
+
+            packageRegistrationTasks.AddTask(isMainThreadTask: false, task: PackageInitializationBackgroundThreadAsync);
+        }
+
+        private Task PackageInitializationBackgroundThreadAsync(IProgress<ServiceProgressData> progress, PackageRegistrationTasks packageRegistrationTasks, CancellationToken cancellationToken)
         {
             try
             {
-                await base.InitializeAsync(cancellationToken, progress).ConfigureAwait(false);
-
                 this.RegisterService<ICSharpTempPECompilerService>(async ct =>
                 {
                     var workspace = this.ComponentModel.GetService<VisualStudioWorkspace>();
@@ -74,6 +79,8 @@ protected override async Task InitializeAsync(CancellationToken cancellationToke
             catch (Exception e) when (FatalError.ReportAndPropagateUnlessCanceled(e, ErrorSeverity.General))
             {
             }
+
+            return Task.CompletedTask;
         }
 
         protected override void RegisterObjectBrowserLibraryManager()

src/VisualStudio/Core/Def/ColorSchemes/ColorSchemeApplier.cs

@@ -15,10 +15,10 @@
 using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.Threading;
 using Microsoft.VisualStudio;
+using Microsoft.VisualStudio.LanguageServices.Implementation.LanguageService;
 using Microsoft.VisualStudio.Settings;
 using Microsoft.VisualStudio.Shell;
 using Microsoft.VisualStudio.Shell.Interop;
-using Microsoft.VisualStudio.Threading;
 using Task = System.Threading.Tasks.Task;
 
 namespace Microsoft.CodeAnalysis.ColorSchemes;
@@ -60,7 +60,7 @@ public ColorSchemeApplier(
             threadingContext.DisposalToken);
     }
 
-    public async Task InitializeAsync(CancellationToken cancellationToken)
+    public void RegisterInitializationWork(PackageRegistrationTasks packageRegistrationTasks)
     {
         lock (_gate)
         {
@@ -70,15 +70,16 @@ public async Task InitializeAsync(CancellationToken cancellationToken)
             _isInitialized = true;
         }
 
-        // We need to update the theme whenever the Editor Color Scheme setting changes.
-        await TaskScheduler.Default;
+        packageRegistrationTasks.AddTask(isMainThreadTask: false, task: PackageInitializationBackgroundThreadAsync);
+    }
+
+    private async Task PackageInitializationBackgroundThreadAsync(IProgress<ServiceProgressData> progress, PackageRegistrationTasks packageRegistrationTasks, CancellationToken cancellationToken)
+    {
         var settingsManager = await _asyncServiceProvider.GetServiceAsync<SVsSettingsPersistenceManager, ISettingsManager>(_threadingContext.JoinableTaskFactory).ConfigureAwait(false);
 
-        await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+        // We need to update the theme whenever the Editor Color Scheme setting changes.
         settingsManager.GetSubset(ColorSchemeOptionsStorage.ColorSchemeSettingKey).SettingChangedAsync += ColorSchemeChangedAsync;
 
-        await TaskScheduler.Default;
-
         // Try to migrate the `useEnhancedColorsSetting` to the new `ColorSchemeName` setting.
         _settings.MigrateToColorSchemeSetting();
 

src/VisualStudio/Core/Def/LanguageService/AbstractPackage.cs

@@ -4,11 +4,11 @@
 
 using System;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.VisualStudio.ComponentModelHost;
 using Microsoft.VisualStudio.Shell;
 using Microsoft.VisualStudio.Threading;
-using Task = System.Threading.Tasks.Task;
 
 namespace Microsoft.VisualStudio.LanguageServices.Implementation.LanguageService;
 
@@ -25,10 +25,38 @@ internal IComponentModel ComponentModel
         }
     }
 
-    protected override async Task InitializeAsync(CancellationToken cancellationToken, IProgress<ServiceProgressData> progress)
+    protected virtual void RegisterInitializationWork(PackageRegistrationTasks packageRegistrationTasks)
     {
-        _componentModel_doNotAccessDirectly = (IComponentModel?)await GetServiceAsync(typeof(SComponentModel)).ConfigureAwait(false);
-        Assumes.Present(_componentModel_doNotAccessDirectly);
+        // This treatment of registering work on the bg/main threads is a bit unique as we want the component model initialized at the beginning
+        // of whichever context is invoked first. The current architecture doesn't execute any of the registered tasks concurrently,
+        // so that isn't a concern for running calculating or setting _componentModel_doNotAccessDirectly multiple times.
+        packageRegistrationTasks.AddTask(isMainThreadTask: false, task: EnsureComponentModelAsync);
+        packageRegistrationTasks.AddTask(isMainThreadTask: true, task: EnsureComponentModelAsync);
+
+        async Task EnsureComponentModelAsync(IProgress<ServiceProgressData> progress, PackageRegistrationTasks packageRegistrationTasks, CancellationToken token)
+        {
+            if (_componentModel_doNotAccessDirectly == null)
+            {
+                _componentModel_doNotAccessDirectly = (IComponentModel?)await GetServiceAsync(typeof(SComponentModel)).ConfigureAwait(false);
+                Assumes.Present(_componentModel_doNotAccessDirectly);
+            }
+        }
+    }
+
+    /// This method is called upon package creation and is the mechanism by which roslyn packages calculate and
+    /// process all package initialization work. Do not override this sealed method, instead override RegisterInitializationWork
+    /// to indicate the work your package needs upon initialization.
+    protected sealed override async Task InitializeAsync(CancellationToken cancellationToken, IProgress<ServiceProgressData> progress)
+    {
+        var packageRegistrationTasks = new PackageRegistrationTasks(JoinableTaskFactory);
+
+        // Request all initially known work, classified into whether it should be processed on the main or
+        // background thread. These lists can be modified by the work itself to add more work for subsequent processing.
+        // Requesting this information is useful as it lets us batch up work on these threads, significantly
+        // reducing thread switches during package load.
+        RegisterInitializationWork(packageRegistrationTasks);
+
+        await packageRegistrationTasks.ProcessTasksAsync(progress, cancellationToken).ConfigureAwait(false);
     }
 
     protected override async Task OnAfterPackageLoadedAsync(CancellationToken cancellationToken)

src/VisualStudio/Core/Def/LanguageService/AbstractPackage`2.cs

@@ -34,16 +34,21 @@ protected AbstractPackage()
     {
     }
 
-    protected override async Task InitializeAsync(CancellationToken cancellationToken, IProgress<ServiceProgressData> progress)
+    protected override void RegisterInitializationWork(PackageRegistrationTasks packageRegistrationTasks)
     {
-        await base.InitializeAsync(cancellationToken, progress).ConfigureAwait(false);
+        base.RegisterInitializationWork(packageRegistrationTasks);
 
-        await JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+        packageRegistrationTasks.AddTask(isMainThreadTask: true, task: PackageInitializationMainThreadAsync);
+    }
+
+    private async Task PackageInitializationMainThreadAsync(IProgress<ServiceProgressData> progress, PackageRegistrationTasks packageRegistrationTasks, CancellationToken cancellationToken)
+    {
+        // This code uses various main thread only services, so it must run completely on the main thread
+        // (thus the CA(true) usage throughout)
+        Contract.ThrowIfFalse(JoinableTaskFactory.Context.IsOnMainThread);
 
         var shell = (IVsShell7?)await GetServiceAsync(typeof(SVsShell)).ConfigureAwait(true);
-        var solution = (IVsSolution?)await GetServiceAsync(typeof(SVsSolution)).ConfigureAwait(true);
         Assumes.Present(shell);
-        Assumes.Present(solution);
 
         _shell = (IVsShell?)shell;
         Assumes.Present(_shell);
@@ -66,9 +71,11 @@ protected override async Task InitializeAsync(CancellationToken cancellationToke
             return _languageService.ComAggregate!;
         });
 
+        var miscellaneousFilesWorkspace = this.ComponentModel.GetService<MiscellaneousFilesWorkspace>();
+
+        // awaiting an IVsTask guarantees to return on the captured context
         await shell.LoadPackageAsync(Guids.RoslynPackageId);
 
-        var miscellaneousFilesWorkspace = this.ComponentModel.GetService<MiscellaneousFilesWorkspace>();
         RegisterMiscellaneousFilesWorkspaceInformation(miscellaneousFilesWorkspace);
 
         if (!_shell.IsInCommandLineMode())
@@ -77,6 +84,13 @@ protected override async Task InitializeAsync(CancellationToken cancellationToke
             // this is a no op
             RegisterObjectBrowserLibraryManager();
         }
+    }
+
+    protected override async Task OnAfterPackageLoadedAsync(CancellationToken cancellationToken)
+    {
+        await base.OnAfterPackageLoadedAsync(cancellationToken).ConfigureAwait(false);
+
+        await JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
 
         LoadComponentsInUIContextOnceSolutionFullyLoadedAsync(cancellationToken).Forget();
     }

src/VisualStudio/Core/Def/LanguageService/PackageRegistrationTasks.cs

@@ -0,0 +1,74 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// 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.Collections.Concurrent;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.VisualStudio.Shell;
+using Microsoft.VisualStudio.Threading;
+
+namespace Microsoft.VisualStudio.LanguageServices.Implementation.LanguageService;
+
+using WorkTask = Func<IProgress<ServiceProgressData>, PackageRegistrationTasks, CancellationToken, Task>;
+
+/// <summary>
+/// Provides a mechanism for registering work to be done during package initialization. Work is registered
+/// as either main thread or background thread appropriate. This allows processing of these work items
+/// in a batched manner, reducing the number of thread switches required during the performance sensitive
+/// package loading timeframe.
+/// 
+/// Note that currently the processing of these tasks isn't done concurrently. A future optimization may
+/// allow parallel background thread task execution, or even concurrent main and background thread work.
+/// </summary>
+internal sealed class PackageRegistrationTasks(JoinableTaskFactory jtf)
+{
+    private readonly ConcurrentQueue<WorkTask> _backgroundThreadWorkTasks = [];
+    private readonly ConcurrentQueue<WorkTask> _mainThreadWorkTasks = [];
+    private readonly JoinableTaskFactory _jtf = jtf;
+
+    public void AddTask(bool isMainThreadTask, WorkTask task)
+    {
+        var workTasks = GetWorkTasks(isMainThreadTask);
+        workTasks.Enqueue(task);
+    }
+
+    public async Task ProcessTasksAsync(IProgress<ServiceProgressData> progress, CancellationToken cancellationToken)
+    {
+        // prime the pump by doing the first group of bg thread work if the initiating thread is not the main thread
+        if (!_jtf.Context.IsOnMainThread)
+            await PerformWorkAsync(isMainThreadTask: false, progress, cancellationToken).ConfigureAwait(false);
+
+        // Continue processing work until everything is completed, switching between main and bg threads as needed.
+        while (!_mainThreadWorkTasks.IsEmpty || !_backgroundThreadWorkTasks.IsEmpty)
+        {
+            await PerformWorkAsync(isMainThreadTask: true, progress, cancellationToken).ConfigureAwait(false);
+            await PerformWorkAsync(isMainThreadTask: false, progress, cancellationToken).ConfigureAwait(false);
+        }
+    }
+
+    private ConcurrentQueue<WorkTask> GetWorkTasks(bool isMainThreadTask)
+        => isMainThreadTask ? _mainThreadWorkTasks : _backgroundThreadWorkTasks;
+
+    private async Task PerformWorkAsync(bool isMainThreadTask, IProgress<ServiceProgressData> progress, CancellationToken cancellationToken)
+    {
+        var workTasks = GetWorkTasks(isMainThreadTask);
+        if (workTasks.IsEmpty)
+            return;
+
+        // Ensure we're invoking the task on the right thread
+        if (isMainThreadTask)
+            await _jtf.SwitchToMainThreadAsync(cancellationToken);
+        else if (_jtf.Context.IsOnMainThread)
+            await TaskScheduler.Default;
+
+        while (workTasks.TryDequeue(out var work))
+        {
+            // CA(true) is important here, as we want to ensure that each iteration is done in the same
+            // captured context. Thus, even poorly behaving tasks (ie, those that do their own thread switching)
+            // don't effect the next loop iteration.
+            await work(progress, this, cancellationToken).ConfigureAwait(true);
+        }
+    }
+}

src/VisualStudio/Core/Def/Options/VisualStudioSettingsOptionPersister.cs

@@ -35,7 +35,7 @@ internal sealed class VisualStudioSettingsOptionPersister
         = ImmutableDictionary<string, (OptionKey2, string)>.Empty;
 
     /// <remarks>
-    /// We make sure this code is from the UI by asking for all <see cref="IOptionPersister"/> in <see cref="RoslynPackage.InitializeAsync"/>
+    /// We make sure this code is from the UI by asking for all <see cref="IOptionPersister"/> in <see cref="RoslynPackage.OnAfterPackageLoadedAsync"/>
     /// </remarks>
     public VisualStudioSettingsOptionPersister(Action<OptionKey2, object?> refreshOption, ImmutableDictionary<string, Lazy<IVisualStudioStorageReadFallback, OptionNameMetadata>> readFallbacks, ISettingsManager settingsManager)
     {

src/VisualStudio/Core/Def/ProjectSystem/MiscellaneousFilesWorkspace.cs

@@ -74,7 +74,6 @@ public MiscellaneousFilesWorkspace(
 
     public async Task InitializeAsync()
     {
-        await TaskScheduler.Default;
         _textManager = await _textManagerService.GetValueAsync().ConfigureAwait(false);
     }