From 7d1ea1b1c78c1f547ddc2b53710911a6e3f541a1 Mon Sep 17 00:00:00 2001
From: Sam Harwell <Sam.Harwell@microsoft.com>
Date: Thu, 17 Jun 2021 13:29:29 -0700
Subject: [PATCH] Initialize MEF services before switching to the main thread

Fixes #54159
---
 .../Core/Impl/CodeModel/CodeModelState.cs     |  2 ++
 .../Impl/CodeModel/ProjectCodeModelFactory.cs | 28 +++++++++++++++++--
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/VisualStudio/Core/Impl/CodeModel/CodeModelState.cs b/src/VisualStudio/Core/Impl/CodeModel/CodeModelState.cs
index 653f5f78d8f3b..f1771f36fcc1a 100644
--- a/src/VisualStudio/Core/Impl/CodeModel/CodeModelState.cs
+++ b/src/VisualStudio/Core/Impl/CodeModel/CodeModelState.cs
@@ -38,6 +38,8 @@ public CodeModelState(
             Debug.Assert(languageServices != null);
             Debug.Assert(workspace != null);
 
+            // ⚠ This code runs on the main thread. Language services accessed here should be preloaded in
+            // ProjectCodemodelFactory to avoid blocking MEF operations.
             this.ThreadingContext = threadingContext;
             this.ServiceProvider = serviceProvider;
             this.CodeModelService = languageServices.GetService<ICodeModelService>();
diff --git a/src/VisualStudio/Core/Impl/CodeModel/ProjectCodeModelFactory.cs b/src/VisualStudio/Core/Impl/CodeModel/ProjectCodeModelFactory.cs
index 3d679fbd43688..1b37aae74fc7f 100644
--- a/src/VisualStudio/Core/Impl/CodeModel/ProjectCodeModelFactory.cs
+++ b/src/VisualStudio/Core/Impl/CodeModel/ProjectCodeModelFactory.cs
@@ -7,15 +7,18 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.ComponentModel.Composition;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis;
-using Microsoft.CodeAnalysis.Editor;
+using Microsoft.CodeAnalysis.CodeGeneration;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
-using Microsoft.CodeAnalysis.PooledObjects;
+using Microsoft.CodeAnalysis.LanguageServices;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.SolutionCrawler;
 using Microsoft.VisualStudio.Shell;
+using Microsoft.VisualStudio.Threading;
 using Roslyn.Utilities;
 using Task = System.Threading.Tasks.Task;
 
@@ -76,6 +79,27 @@ private async Task ProcessNextDocumentBatchAsync(
             const int MaxTimeSlice = 15;
             var delayBetweenProcessing = TimeSpan.FromMilliseconds(50);
 
+            Debug.Assert(!_threadingContext.JoinableTaskContext.IsOnMainThread, "The following context switch is not expected to cause runtime overhead.");
+            await TaskScheduler.Default;
+
+            // Ensure MEF services used by the code model are initially obtained on a background thread.
+            // This code avoids allocations where possible.
+            // https://github.com/dotnet/roslyn/issues/54159
+            string? previousLanguage = null;
+            foreach (var (_, projectState) in _visualStudioWorkspace.CurrentSolution.State.ProjectStates)
+            {
+                if (projectState.Language == previousLanguage)
+                {
+                    // Avoid duplicate calls if the language did not change
+                    continue;
+                }
+
+                previousLanguage = projectState.Language;
+                projectState.LanguageServices.GetService<ICodeModelService>();
+                projectState.LanguageServices.GetService<ISyntaxFactsService>();
+                projectState.LanguageServices.GetService<ICodeGenerationService>();
+            }
+
             await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
 
             var stopwatch = SharedStopwatch.StartNew();