Switch OnAfterPackageLoadedAsync to use a TaskQueue

src/VisualStudio/CSharp/Impl/CSharpPackage.cs

@@ -58,14 +58,14 @@ internal sealed class CSharpPackage : AbstractPackage<CSharpPackage, CSharpLangu
         private ObjectBrowserLibraryManager? _libraryManager;
         private uint _libraryManagerCookie;
 
-        protected override void RegisterInitializationWork(PackageRegistrationTasks packageRegistrationTasks)
+        protected override void RegisterInitializeAsyncWork(PackageLoadTasks packageInitializationTasks)
         {
-            base.RegisterInitializationWork(packageRegistrationTasks);
+            base.RegisterInitializeAsyncWork(packageInitializationTasks);
 
-            packageRegistrationTasks.AddTask(isMainThreadTask: false, task: PackageInitializationBackgroundThreadAsync);
+            packageInitializationTasks.AddTask(isMainThreadTask: false, task: PackageInitializationBackgroundThreadAsync);
         }
 
-        private Task PackageInitializationBackgroundThreadAsync(IProgress<ServiceProgressData> progress, PackageRegistrationTasks packageRegistrationTasks, CancellationToken cancellationToken)
+        private Task PackageInitializationBackgroundThreadAsync(PackageLoadTasks packageInitializationTasks, CancellationToken cancellationToken)
         {
             try
             {

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

@@ -57,7 +57,7 @@ public ColorSchemeApplier(
             threadingContext.DisposalToken);
     }
 
-    public void RegisterInitializationWork(PackageRegistrationTasks packageRegistrationTasks)
+    public void RegisterInitializationWork(PackageLoadTasks packageInitializationTasks)
     {
         lock (_gate)
         {
@@ -67,10 +67,10 @@ public void RegisterInitializationWork(PackageRegistrationTasks packageRegistrat
             _isInitialized = true;
         }
 
-        packageRegistrationTasks.AddTask(isMainThreadTask: false, task: PackageInitializationBackgroundThreadAsync);
+        packageInitializationTasks.AddTask(isMainThreadTask: false, task: PackageInitializationBackgroundThreadAsync);
     }
 
-    private async Task PackageInitializationBackgroundThreadAsync(IProgress<ServiceProgressData> progress, PackageRegistrationTasks packageRegistrationTasks, CancellationToken cancellationToken)
+    private async Task PackageInitializationBackgroundThreadAsync(PackageLoadTasks packageInitializationTasks, CancellationToken cancellationToken)
     {
         var settingsManager = await _asyncServiceProvider.GetServiceAsync<SVsSettingsPersistenceManager, ISettingsManager>(_threadingContext.JoinableTaskFactory).ConfigureAwait(false);
 

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

@@ -8,7 +8,6 @@
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.VisualStudio.ComponentModelHost;
 using Microsoft.VisualStudio.Shell;
-using Microsoft.VisualStudio.Threading;
 
 namespace Microsoft.VisualStudio.LanguageServices.Implementation.LanguageService;
 
@@ -25,15 +24,40 @@ internal IComponentModel ComponentModel
         }
     }
 
-    protected virtual void RegisterInitializationWork(PackageRegistrationTasks packageRegistrationTasks)
+    /// 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 RegisterOnAfterPackageLoadedAsyncWork
+    /// to indicate the work your package needs upon initialization.
+    protected sealed override Task InitializeAsync(CancellationToken cancellationToken, IProgress<ServiceProgressData> progress)
+        => RegisterAndProcessTasksAsync(RegisterInitializeAsyncWork, cancellationToken);
+
+    /// This method is called after package load and is the mechanism by which roslyn packages calculate and
+    /// process all post load package work. Do not override this sealed method, instead override RegisterOnAfterPackageLoadedAsyncWork
+    /// to indicate the work your package needs after load.
+    protected sealed override Task OnAfterPackageLoadedAsync(CancellationToken cancellationToken)
+        => RegisterAndProcessTasksAsync(RegisterOnAfterPackageLoadedAsyncWork, cancellationToken);
+
+    private Task RegisterAndProcessTasksAsync(Action<PackageLoadTasks> registerTasks, CancellationToken cancellationToken)
+    {
+        var packageTasks = new PackageLoadTasks(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.
+        registerTasks(packageTasks);
+
+        return packageTasks.ProcessTasksAsync(cancellationToken);
+    }
+
+    protected virtual void RegisterInitializeAsyncWork(PackageLoadTasks packageInitializationTasks)
     {
         // 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);
+        // so that isn't a concern for calculating or setting _componentModel_doNotAccessDirectly multiple times.
+        packageInitializationTasks.AddTask(isMainThreadTask: false, task: EnsureComponentModelAsync);
+        packageInitializationTasks.AddTask(isMainThreadTask: true, task: EnsureComponentModelAsync);
 
-        async Task EnsureComponentModelAsync(IProgress<ServiceProgressData> progress, PackageRegistrationTasks packageRegistrationTasks, CancellationToken token)
+        async Task EnsureComponentModelAsync(PackageLoadTasks packageInitializationTasks, CancellationToken token)
         {
             if (_componentModel_doNotAccessDirectly == null)
             {
@@ -43,31 +67,28 @@ async Task EnsureComponentModelAsync(IProgress<ServiceProgressData> progress, Pa
         }
     }
 
-    /// 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)
+    protected virtual void RegisterOnAfterPackageLoadedAsyncWork(PackageLoadTasks afterPackageLoadedTasks)
     {
-        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);
+        afterPackageLoadedTasks.AddTask(
+            isMainThreadTask: false,
+            task: (afterPackageLoadedTasks, cancellationToken) =>
+            {
+                // TODO: remove, workaround for https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1985204
+                var globalOptions = ComponentModel.GetService<IGlobalOptionService>();
+                if (globalOptions.GetOption(SemanticSearchFeatureFlag.Enabled))
+                {
+                    afterPackageLoadedTasks.AddTask(
+                        isMainThreadTask: true,
+                        task: (packageLoadedTasks, cancellationToken) =>
+                        {
+                            UIContext.FromUIContextGuid(new Guid(SemanticSearchFeatureFlag.UIContextId)).IsActive = true;
 
-        await packageRegistrationTasks.ProcessTasksAsync(progress, cancellationToken).ConfigureAwait(false);
-    }
+                            return Task.CompletedTask;
+                        });
+                }
 
-    protected override async Task OnAfterPackageLoadedAsync(CancellationToken cancellationToken)
-    {
-        // TODO: remove, workaround for https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1985204
-        var globalOptions = ComponentModel.GetService<IGlobalOptionService>();
-        if (globalOptions.GetOption(SemanticSearchFeatureFlag.Enabled))
-        {
-            await JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
-            UIContext.FromUIContextGuid(new Guid(SemanticSearchFeatureFlag.UIContextId)).IsActive = true;
-        }
+                return Task.CompletedTask;
+            });
     }
 
     protected async Task LoadComponentsInUIContextOnceSolutionFullyLoadedAsync(CancellationToken cancellationToken)

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

@@ -34,14 +34,14 @@ protected AbstractPackage()
     {
     }
 
-    protected override void RegisterInitializationWork(PackageRegistrationTasks packageRegistrationTasks)
+    protected override void RegisterInitializeAsyncWork(PackageLoadTasks packageInitializationTasks)
     {
-        base.RegisterInitializationWork(packageRegistrationTasks);
+        base.RegisterInitializeAsyncWork(packageInitializationTasks);
 
-        packageRegistrationTasks.AddTask(isMainThreadTask: true, task: PackageInitializationMainThreadAsync);
+        packageInitializationTasks.AddTask(isMainThreadTask: true, task: PackageInitializationMainThreadAsync);
     }
 
-    private async Task PackageInitializationMainThreadAsync(IProgress<ServiceProgressData> progress, PackageRegistrationTasks packageRegistrationTasks, CancellationToken cancellationToken)
+    private async Task PackageInitializationMainThreadAsync(PackageLoadTasks packageInitializationTasks, 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)
@@ -63,9 +63,9 @@ private async Task PackageInitializationMainThreadAsync(IProgress<ServiceProgres
         // awaiting an IVsTask guarantees to return on the captured context
         await shell.LoadPackageAsync(Guids.RoslynPackageId);
 
-        packageRegistrationTasks.AddTask(
+        packageInitializationTasks.AddTask(
              isMainThreadTask: false,
-             task: (IProgress<ServiceProgressData> progress, PackageRegistrationTasks packageRegistrationTasks, CancellationToken cancellationToken) =>
+             task: (PackageLoadTasks packageInitializationTasks, CancellationToken cancellationToken) =>
              {
                  RegisterLanguageService(typeof(TLanguageService), async cancellationToken =>
                  {
@@ -86,20 +86,25 @@ private async Task PackageInitializationMainThreadAsync(IProgress<ServiceProgres
              });
     }
 
-    protected override async Task OnAfterPackageLoadedAsync(CancellationToken cancellationToken)
+    protected override void RegisterOnAfterPackageLoadedAsyncWork(PackageLoadTasks afterPackageLoadedTasks)
     {
-        await base.OnAfterPackageLoadedAsync(cancellationToken).ConfigureAwait(false);
+        base.RegisterOnAfterPackageLoadedAsyncWork(afterPackageLoadedTasks);
 
-        await JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+        afterPackageLoadedTasks.AddTask(
+            isMainThreadTask: true,
+            task: (packageLoadedTasks, cancellationToken) =>
+            {
+                if (_shell != null && !_shell.IsInCommandLineMode())
+                {
+                    // not every derived package support object browser and for those languages
+                    // this is a no op
+                    RegisterObjectBrowserLibraryManager();
+                }
 
-        if (_shell != null && !_shell.IsInCommandLineMode())
-        {
-            // not every derived package support object browser and for those languages
-            // this is a no op
-            RegisterObjectBrowserLibraryManager();
-        }
+                LoadComponentsInUIContextOnceSolutionFullyLoadedAsync(cancellationToken).Forget();
 
-        LoadComponentsInUIContextOnceSolutionFullyLoadedAsync(cancellationToken).Forget();
+                return Task.CompletedTask;
+            });
     }
 
     protected override async Task LoadComponentsAsync(CancellationToken cancellationToken)

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

@@ -11,7 +11,7 @@
 
 namespace Microsoft.VisualStudio.LanguageServices.Implementation.LanguageService;
 
-using WorkTask = Func<IProgress<ServiceProgressData>, PackageRegistrationTasks, CancellationToken, Task>;
+using WorkTask = Func<PackageLoadTasks, CancellationToken, Task>;
 
 /// <summary>
 /// Provides a mechanism for registering work to be done during package initialization. Work is registered
@@ -22,7 +22,7 @@ namespace Microsoft.VisualStudio.LanguageServices.Implementation.LanguageService
 /// 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)
+internal sealed class PackageLoadTasks(JoinableTaskFactory jtf)
 {
     private readonly ConcurrentQueue<WorkTask> _backgroundThreadWorkTasks = [];
     private readonly ConcurrentQueue<WorkTask> _mainThreadWorkTasks = [];
@@ -34,24 +34,24 @@ public void AddTask(bool isMainThreadTask, WorkTask task)
         workTasks.Enqueue(task);
     }
 
-    public async Task ProcessTasksAsync(IProgress<ServiceProgressData> progress, CancellationToken cancellationToken)
+    public async Task ProcessTasksAsync(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);
+            await PerformWorkAsync(isMainThreadTask: false, 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);
+            await PerformWorkAsync(isMainThreadTask: true, cancellationToken).ConfigureAwait(false);
+            await PerformWorkAsync(isMainThreadTask: false, cancellationToken).ConfigureAwait(false);
         }
     }
 
     private ConcurrentQueue<WorkTask> GetWorkTasks(bool isMainThreadTask)
         => isMainThreadTask ? _mainThreadWorkTasks : _backgroundThreadWorkTasks;
 
-    private async Task PerformWorkAsync(bool isMainThreadTask, IProgress<ServiceProgressData> progress, CancellationToken cancellationToken)
+    private async Task PerformWorkAsync(bool isMainThreadTask, CancellationToken cancellationToken)
     {
         var workTasks = GetWorkTasks(isMainThreadTask);
         if (workTasks.IsEmpty)
@@ -68,7 +68,7 @@ private async Task PerformWorkAsync(bool isMainThreadTask, IProgress<ServiceProg
             // 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);
+            await work(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.OnAfterPackageLoadedAsync"/>
+    /// We make sure this code is from the UI by asking for all <see cref="IOptionPersister"/> in <see cref="RoslynPackage.RegisterOnAfterPackageLoadedAsyncWork"/>
     /// </remarks>
     public VisualStudioSettingsOptionPersister(Action<OptionKey2, object?> refreshOption, ImmutableDictionary<string, Lazy<IVisualStudioStorageReadFallback, OptionNameMetadata>> readFallbacks, ISettingsManager settingsManager)
     {