Reduce main thread switches in VisualStudioProjectFactory.CreateAndAddToWorkspaceAsync

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

@@ -14,6 +14,7 @@
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Workspaces.AnalyzerRedirecting;
 using Microsoft.CodeAnalysis.Workspaces.ProjectSystem;
+using Microsoft.Internal.VisualStudio.Shell.Interop;
 using Microsoft.VisualStudio.LanguageServices.ExternalAccess.VSTypeScript.Api;
 using Microsoft.VisualStudio.LanguageServices.Implementation.Diagnostics;
 using Microsoft.VisualStudio.Shell;
@@ -35,7 +36,9 @@ internal sealed class VisualStudioProjectFactory : IVsTypeScriptVisualStudioProj
     private readonly ImmutableArray<Lazy<IDynamicFileInfoProvider, FileExtensionsMetadata>> _dynamicFileInfoProviders;
     private readonly IVisualStudioDiagnosticAnalyzerProviderFactory _vsixAnalyzerProviderFactory;
     private readonly ImmutableArray<IAnalyzerAssemblyRedirector> _analyzerAssemblyRedirectors;
-    private readonly IVsService<SVsSolution, IVsSolution2> _solution2;
+    private readonly IVsService<SVsBackgroundSolution, IVsBackgroundSolution> _solution;
+
+    private readonly JoinableTask<VisualStudioDiagnosticAnalyzerProvider> _initializationTask;
 
     [ImportingConstructor]
     [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
@@ -45,14 +48,33 @@ public VisualStudioProjectFactory(
         [ImportMany] IEnumerable<Lazy<IDynamicFileInfoProvider, FileExtensionsMetadata>> fileInfoProviders,
         IVisualStudioDiagnosticAnalyzerProviderFactory vsixAnalyzerProviderFactory,
         [ImportMany] IEnumerable<IAnalyzerAssemblyRedirector> analyzerAssemblyRedirectors,
-        IVsService<SVsSolution, IVsSolution2> solution2)
+        IVsService<SVsBackgroundSolution, IVsBackgroundSolution> solution)
     {
         _threadingContext = threadingContext;
         _visualStudioWorkspaceImpl = visualStudioWorkspaceImpl;
         _dynamicFileInfoProviders = fileInfoProviders.AsImmutableOrEmpty();
         _vsixAnalyzerProviderFactory = vsixAnalyzerProviderFactory;
         _analyzerAssemblyRedirectors = analyzerAssemblyRedirectors.AsImmutableOrEmpty();
-        _solution2 = solution2;
+        _solution = solution;
+
+        _initializationTask = _threadingContext.JoinableTaskFactory.RunAsync(
+            async () =>
+            {
+                var cancellationToken = _threadingContext.DisposalToken;
+
+                // HACK: Fetch this service to ensure it's still created on the UI thread; once this is
+                // moved off we'll need to fix up it's constructor to be free-threaded.
+
+                // yield if on the main thread, as the VisualStudioMetadataReferenceManager construction can be fairly expensive
+                // and we don't want the case where VisualStudioProjectFactory is constructed on the main thread to block on that.
+                await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(alwaysYield: true, cancellationToken);
+                _visualStudioWorkspaceImpl.Services.GetRequiredService<VisualStudioMetadataReferenceManager>();
+
+                _visualStudioWorkspaceImpl.SubscribeExternalErrorDiagnosticUpdateSourceToSolutionBuildEvents();
+                _visualStudioWorkspaceImpl.SubscribeToSourceGeneratorImpactingEvents();
+
+                return await _vsixAnalyzerProviderFactory.GetOrCreateProviderAsync(cancellationToken).ConfigureAwait(true);
+            });
     }
 
     public Task<ProjectSystemProject> CreateAndAddToWorkspaceAsync(string projectSystemName, string language, CancellationToken cancellationToken)
@@ -61,24 +83,7 @@ public Task<ProjectSystemProject> CreateAndAddToWorkspaceAsync(string projectSys
     public async Task<ProjectSystemProject> CreateAndAddToWorkspaceAsync(
         string projectSystemName, string language, VisualStudioProjectCreationInfo creationInfo, CancellationToken cancellationToken)
     {
-        // HACK: Fetch this service to ensure it's still created on the UI thread; once this is
-        // moved off we'll need to fix up it's constructor to be free-threaded.
-
-        await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
-        _visualStudioWorkspaceImpl.Services.GetRequiredService<VisualStudioMetadataReferenceManager>();
-
-        _visualStudioWorkspaceImpl.SubscribeExternalErrorDiagnosticUpdateSourceToSolutionBuildEvents();
-        _visualStudioWorkspaceImpl.SubscribeToSourceGeneratorImpactingEvents();
-
-#pragma warning disable CA2007 // Consider calling ConfigureAwait on the awaited task
-        // Since we're on the UI thread here anyways, use that as an opportunity to grab the
-        // IVsSolution object and solution file path.
-        var solution = await _solution2.GetValueOrNullAsync(cancellationToken);
-        var solutionFilePath = solution != null && ErrorHandler.Succeeded(solution.GetSolutionInfo(out _, out var filePath, out _))
-            ? filePath
-            : null;
-
-        var vsixAnalyzerProvider = await _vsixAnalyzerProviderFactory.GetOrCreateProviderAsync(cancellationToken).ConfigureAwait(false);
+        var vsixAnalyzerProvider = await _initializationTask.JoinAsync(cancellationToken).ConfigureAwait(false);
 
         // The rest of this method can be ran off the UI thread. We'll only switch though if the UI thread isn't already blocked -- the legacy project
         // system creates project synchronously, and during solution load we've seen traces where the thread pool is sufficiently saturated that this
@@ -89,27 +94,24 @@ public async Task<ProjectSystemProject> CreateAndAddToWorkspaceAsync(
             await TaskScheduler.Default;
         }
 
+        var solution = await _solution.GetValueOrNullAsync(cancellationToken).ConfigureAwait(true);
+
         // From this point on, we start mutating the solution.  So make us non cancellable.
-#pragma warning disable IDE0059 // Unnecessary assignment of a value
         cancellationToken = CancellationToken.None;
-#pragma warning restore IDE0059 // Unnecessary assignment of a value
 
-        _visualStudioWorkspaceImpl.ProjectSystemProjectFactory.SolutionPath = solutionFilePath;
+        _visualStudioWorkspaceImpl.ProjectSystemProjectFactory.SolutionPath = solution?.SolutionFileName;
         _visualStudioWorkspaceImpl.ProjectSystemProjectFactory.SolutionTelemetryId = GetSolutionSessionId();
 
         var hostInfo = new ProjectSystemHostInfo(_dynamicFileInfoProviders, vsixAnalyzerProvider, _analyzerAssemblyRedirectors);
-        var project = await _visualStudioWorkspaceImpl.ProjectSystemProjectFactory.CreateAndAddToWorkspaceAsync(projectSystemName, language, creationInfo, hostInfo);
+        var project = await _visualStudioWorkspaceImpl.ProjectSystemProjectFactory.CreateAndAddToWorkspaceAsync(projectSystemName, language, creationInfo, hostInfo).ConfigureAwait(true);
 
         _visualStudioWorkspaceImpl.AddProjectToInternalMaps(project, creationInfo.Hierarchy, creationInfo.ProjectGuid, projectSystemName);
 
         // Ensure that other VS contexts get accurate information that the UIContext for this language is now active.
-        // This is not cancellable as we have already mutated the solution.
-        await _visualStudioWorkspaceImpl.RefreshProjectExistsUIContextForLanguageAsync(language, CancellationToken.None);
+        await _visualStudioWorkspaceImpl.RefreshProjectExistsUIContextForLanguageAsync(language, cancellationToken).ConfigureAwait(true);
 
         return project;
 
-#pragma warning restore CA2007 // Consider calling ConfigureAwait on the awaited task
-
         static Guid GetSolutionSessionId()
         {
             var dataModelTelemetrySession = TelemetryService.DefaultSession;

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

@@ -92,9 +92,15 @@ internal abstract partial class VisualStudioWorkspaceImpl : VisualStudioWorkspac
     private readonly Dictionary<string, List<ProjectSystemProject>> _projectSystemNameToProjectsMap = [];
 
     /// <summary>
-    /// Only safe to use on the UI thread.
+    /// Mapping from language name to an existing UIContext's active state.
+    /// Only access when holding <see cref="_gate"/>
     /// </summary>
-    private readonly Dictionary<string, UIContext?> _languageToProjectExistsUIContext = [];
+    private readonly Dictionary<string, bool> _languageToProjectExistsUIContextState = [];
+
+    /// <summary>
+    /// Joinable task collection to await to ensure language ui contexts are updated.
+    /// </summary>
+    private readonly JoinableTaskCollection _updateUIContextJoinableTasks;
 
     private OpenFileTracker? _openFileTracker;
     internal IFileChangeWatcher FileChangeWatcher { get; }
@@ -131,6 +137,8 @@ public VisualStudioWorkspaceImpl(ExportProvider exportProvider, IAsyncServicePro
             exportProvider.GetExportedValue<ExternalErrorDiagnosticUpdateSource>(),
             isThreadSafe: true);
 
+        _updateUIContextJoinableTasks = new JoinableTaskCollection(_threadingContext.JoinableTaskContext);
+
         _workspaceListener = Services.GetRequiredService<IWorkspaceAsynchronousOperationListenerProvider>().GetListener();
     }
 
@@ -1562,12 +1570,35 @@ internal void RemoveProjectFromMaps(CodeAnalysis.Project project)
     }
 
     internal async Task RefreshProjectExistsUIContextForLanguageAsync(string language, CancellationToken cancellationToken)
+    {
+        using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
+        {
+            var isContextActive = _languageToProjectExistsUIContextState.GetOrAdd(language, false);
+
+            // Determine if there is a project with a matching language. Uses _projectSystemNameToProjectsMap as 
+            // that data structure is updated before calling into this method, whereas CurrentSolution may not be.
+            var projectExistsWithLanguage = _projectSystemNameToProjectsMap.Any(projects => projects.Value.Any(project => project.Language == language));
+            if (projectExistsWithLanguage != isContextActive)
+            {
+                _languageToProjectExistsUIContextState[language] = projectExistsWithLanguage;
+
+                // Create a task to update the UI context, and add it to the task collection that all callers to 
+                // this method will wait on before returning.
+                var joinableTask = _threadingContext.JoinableTaskFactory.RunAsync(() => UpdateUIContextAsync(language, cancellationToken));
+
+                _updateUIContextJoinableTasks.Add(joinableTask);
+            }
+        }
+
+        // Ensure any pending ui context updates have occurred before returning
+        await _updateUIContextJoinableTasks.JoinTillEmptyAsync(cancellationToken).ConfigureAwait(false);
+    }
+
+    private async Task UpdateUIContextAsync(string language, CancellationToken cancellationToken)
     {
         await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(alwaysYield: true, cancellationToken);
 
-        var uiContext = _languageToProjectExistsUIContext.GetOrAdd(
-            language,
-            language => Services.GetLanguageServices(language).GetService<IProjectExistsUIContextProviderLanguageService>()?.GetUIContext());
+        var uiContext = Services.GetLanguageServices(language).GetService<IProjectExistsUIContextProviderLanguageService>()?.GetUIContext();
 
         // UIContexts can be "zombied" if UIContexts aren't supported because we're in a command line build or in
         // other scenarios.
@@ -1578,7 +1609,9 @@ internal async Task RefreshProjectExistsUIContextForLanguageAsync(string languag
         // thread, so that acts as a natural ordering mechanism here.  If, say, a BG piece of work was mutating this
         // solution (either adding or removing a project) then that work will also have enqueued the next refresh
         // operation on the UI thread.  So we'll always eventually reach a fixed point where the task for that
-        // language will check the latest CurrentSolution we have and will set the IsActive bit accordingly.
+        // language will check the latest CurrentSolution we have and will set the IsActive bit accordingly. We
+        // don't use the isContextActive value here specifically for this case as it may not reflect the desired
+        // value after the main thread switch.
         uiContext.IsActive = this.CurrentSolution.Projects.Any(p => p.Language == language);
     }
 }