[release/dev18.0] Try to reuse initialized generator drivers when we can

src/Workspaces/Core/Portable/Workspace/Solution/Solution.cs

@@ -62,7 +62,8 @@ internal Solution(
         ImmutableDictionary<string, StructuredAnalyzerConfigOptions> fallbackAnalyzerOptions)
         : this(new SolutionCompilationState(
             new SolutionState(workspace.Kind, workspace.Services.SolutionServices, solutionAttributes, options, analyzerReferences, fallbackAnalyzerOptions),
-            workspace.PartialSemanticsEnabled))
+            workspace.PartialSemanticsEnabled,
+            workspace.GeneratorDriverCreationCache))
     {
     }
 

src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.GeneratorDriverInitializationCache.cs

@@ -0,0 +1,107 @@
+// 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.Generic;
+using System.Collections.Immutable;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Host;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis;
+
+internal sealed partial class SolutionCompilationState
+{
+    internal sealed class GeneratorDriverInitializationCache
+    {
+        /// <summary>
+        /// A set of GeneratorDriver instances that have been created for the keyed project in the solution. Any time we create a GeneratorDriver the first time for
+        /// a project, we'll put it into this map. If other requests come in to get a GeneratorDriver for the same project (but from different Solution snapshots),
+        /// well reuse this GeneratorDriver rather than creating a new one. This allows some first time initialization of a generator (like walking metadata references)
+        /// to be shared rather than doing that initialization multiple times. In the case we are reusing a GeneratorDriver, we'll still always update the GeneratorDriver with
+        /// the current state of the project, so the results are still correct.
+        /// 
+        /// Since these entries are going to be holding onto non-trivial amounts of state, we get rid of the cached entries once there's a belief that we won't be
+        /// creating further GeneratorDrivers for a given project. See uses of <see cref="EmptyCacheForProjectsThatHaveGeneratorDriversInSolution"/>
+        /// for details.
+        /// 
+        /// Any additions/removals to this map must be done via ImmutableInterlocked methods.
+        /// </summary>
+        private ImmutableDictionary<ProjectId, AsyncLazy<GeneratorDriver>> _driverCache = ImmutableDictionary<ProjectId, AsyncLazy<GeneratorDriver>>.Empty;
+
+        public async Task<GeneratorDriver> CreateAndRunGeneratorDriverAsync(
+            ProjectState projectState,
+            Compilation compilation,
+            Func<GeneratorFilterContext, bool> generatorFilter,
+            CancellationToken cancellationToken)
+        {
+            // The AsyncLazy we create here implicitly creates a GeneratorDriver that will run generators for the compilation passed to this method.
+            // If the one that is added to _driverCache is the one we created, then it's ready to go. If the AsyncLazy is one created by some
+            // other call, then we'll still need to run generators for the compilation passed.
+            var createdAsyncLazy = AsyncLazy.Create(CreateGeneratorDriverAndRunGenerators);
+            var asyncLazy = ImmutableInterlocked.GetOrAdd(ref _driverCache, projectState.Id, static (_, created) => created, createdAsyncLazy);
+
+            if (asyncLazy == createdAsyncLazy)
+            {
+                // We want to ensure that the driver is always created and initialized at least once, so we'll ensure that runs even if we cancel the request here.
+                // Otherwise the concern is we might keep starting and cancelling the work which is just wasteful to keep doing it over and over again. We do this
+                // in a Task.Run() so if the underlying computation were to run on our thread, we're not blocking our caller from observing cancellation
+                // if they request it.
+                _ = Task.Run(() => asyncLazy.GetValueAsync(CancellationToken.None));
+
+                return await asyncLazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
+            }
+            else
+            {
+                var driver = await asyncLazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
+
+                driver = UpdateGeneratorDriverToMatchState(driver, projectState);
+
+                return driver.RunGenerators(compilation, generatorFilter, cancellationToken);
+            }
+
+            GeneratorDriver CreateGeneratorDriverAndRunGenerators(CancellationToken cancellationToken)
+            {
+                var generatedFilesBaseDirectory = projectState.CompilationOutputInfo.GetEffectiveGeneratedFilesOutputDirectory();
+                var additionalTexts = projectState.AdditionalDocumentStates.SelectAsArray(static documentState => documentState.AdditionalText);
+                var compilationFactory = projectState.LanguageServices.GetRequiredService<ICompilationFactoryService>();
+
+                var generatorDriver = compilationFactory.CreateGeneratorDriver(
+                    projectState.ParseOptions!,
+                    GetSourceGenerators(projectState),
+                    projectState.ProjectAnalyzerOptions.AnalyzerConfigOptionsProvider,
+                    additionalTexts,
+                    generatedFilesBaseDirectory);
+
+                return generatorDriver.RunGenerators(compilation, generatorFilter, cancellationToken);
+            }
+        }
+
+        public void EmptyCacheForProjectsThatHaveGeneratorDriversInSolution(SolutionCompilationState state)
+        {
+            // If we don't have any cached drivers, then just return before we loop through all the projects
+            // in the solution. This is to ensure that once we hit a steady-state case of a Workspace's CurrentSolution
+            // having generators for all projects, we won't need to keep anything further in our cache since the cache
+            // will never be used -- any running of generators in the future will use the GeneratorDrivers already held by
+            // the Solutions.
+            //
+            // This doesn't need to be synchronized against other mutations to _driverCache. If we see it as empty when
+            // in reality something was just being added, we'll just do the cleanup the next time this method is called.
+            if (_driverCache.IsEmpty)
+                return;
+
+            foreach (var (projectId, tracker) in state._projectIdToTrackerMap)
+            {
+                if (tracker.GeneratorDriver is not null)
+                    EmptyCacheForProject(projectId);
+            }
+        }
+
+        public void EmptyCacheForProject(ProjectId projectId)
+        {
+            ImmutableInterlocked.TryRemove(ref _driverCache, projectId, out _);
+        }
+    }
+}

src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.RegularCompilationTracker_Generators.cs

@@ -3,8 +3,8 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
-using System.Collections.Immutable;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Linq;
 using System.Text;
@@ -13,10 +13,8 @@
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Collections;
 using Microsoft.CodeAnalysis.Diagnostics;
-using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Remote;
-using Microsoft.CodeAnalysis.Shared.Collections;
 using Microsoft.CodeAnalysis.SourceGeneration;
 using Microsoft.CodeAnalysis.SourceGeneratorTelemetry;
 using Microsoft.CodeAnalysis.Text;
@@ -269,16 +267,6 @@ await newGeneratedDocuments.States.Values.SelectAsArrayAsync(
             if (!await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false))
                 return (compilationWithoutGeneratedFiles, TextDocumentStates<SourceGeneratedDocumentState>.Empty, generatorDriver);
 
-            // Hold onto the prior results so we can compare when filtering
-            var priorRunResult = generatorDriver?.GetRunResult();
-
-            // If we don't already have an existing generator driver, create one from scratch
-            generatorDriver ??= CreateGeneratorDriver(this.ProjectState);
-
-            CheckGeneratorDriver(generatorDriver, this.ProjectState);
-
-            Contract.ThrowIfNull(generatorDriver);
-
             // HACK HACK HACK HACK to address https://github.com/dotnet/roslyn/issues/59818. There, we were running into issues where
             // a generator being present and consuming syntax was causing all red nodes to be processed. This was problematic when
             // Razor design time files are also fed in, since those files tend to be quite large. The Razor design time files
@@ -298,9 +286,19 @@ await newGeneratedDocuments.States.Values.SelectAsArrayAsync(
             var compilationToRunGeneratorsOn = compilationWithoutGeneratedFiles.RemoveSyntaxTrees(treesToRemove);
             // END HACK HACK HACK HACK.
 
-            generatorDriver = generatorDriver.RunGenerators(compilationToRunGeneratorsOn, ShouldGeneratorRun, cancellationToken);
+            // Hold onto the prior results so we can compare when filtering
+            var priorRunResult = generatorDriver?.GetRunResult();
 
-            Contract.ThrowIfNull(generatorDriver);
+            if (generatorDriver == null)
+            {
+                generatorDriver = await compilationState.GeneratorDriverCache.CreateAndRunGeneratorDriverAsync(this.ProjectState, compilationToRunGeneratorsOn, ShouldGeneratorRun, cancellationToken).ConfigureAwait(false);
+            }
+            else
+            {
+                generatorDriver = generatorDriver.RunGenerators(compilationToRunGeneratorsOn, ShouldGeneratorRun, cancellationToken);
+            }
+
+            CheckGeneratorDriver(generatorDriver, this.ProjectState);
 
             var runResult = generatorDriver.GetRunResult();
 
@@ -426,20 +424,6 @@ await newGeneratedDocuments.States.Values.SelectAsArrayAsync(
                 return null;
             }
 
-            static GeneratorDriver CreateGeneratorDriver(ProjectState projectState)
-            {
-                var generatedFilesBaseDirectory = projectState.CompilationOutputInfo.GetEffectiveGeneratedFilesOutputDirectory();
-                var additionalTexts = projectState.AdditionalDocumentStates.SelectAsArray(static documentState => documentState.AdditionalText);
-                var compilationFactory = projectState.LanguageServices.GetRequiredService<ICompilationFactoryService>();
-
-                return compilationFactory.CreateGeneratorDriver(
-                    projectState.ParseOptions!,
-                    GetSourceGenerators(projectState),
-                    projectState.ProjectAnalyzerOptions.AnalyzerConfigOptionsProvider,
-                    additionalTexts,
-                    generatedFilesBaseDirectory);
-            }
-
             [Conditional("DEBUG")]
             static void CheckGeneratorDriver(GeneratorDriver generatorDriver, ProjectState projectState)
             {

src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.TranslationAction_Actions.cs

@@ -362,11 +362,7 @@ public override GeneratorDriver TransformGeneratorDriver(GeneratorDriver _)
             {
                 // The GeneratorDriver that we have here is from a prior version of the Project, it may be missing state changes due
                 // to changes to the project. We'll update everything here.
-                var generatorDriver = oldGeneratorDriver
-                    .ReplaceAdditionalTexts(this.NewProjectState.AdditionalDocumentStates.SelectAsArray(static documentState => documentState.AdditionalText))
-                    .WithUpdatedParseOptions(this.NewProjectState.ParseOptions!)
-                    .WithUpdatedAnalyzerConfigOptions(this.NewProjectState.ProjectAnalyzerOptions.AnalyzerConfigOptionsProvider)
-                    .ReplaceGenerators(GetSourceGenerators(this.NewProjectState));
+                var generatorDriver = UpdateGeneratorDriverToMatchState(oldGeneratorDriver, NewProjectState);
 
                 return generatorDriver;
             }

src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.cs

@@ -43,6 +43,8 @@ internal sealed partial class SolutionCompilationState
     public bool PartialSemanticsEnabled { get; }
     public TextDocumentStates<SourceGeneratedDocumentState> FrozenSourceGeneratedDocumentStates { get; }
 
+    public GeneratorDriverInitializationCache GeneratorDriverCache { get; }
+
     // Values for all these are created on demand.
     private ImmutableSegmentedDictionary<ProjectId, ICompilationTracker> _projectIdToTrackerMap;
 
@@ -62,13 +64,15 @@ private SolutionCompilationState(
         ImmutableSegmentedDictionary<ProjectId, ICompilationTracker> projectIdToTrackerMap,
         SourceGeneratorExecutionVersionMap sourceGeneratorExecutionVersionMap,
         TextDocumentStates<SourceGeneratedDocumentState> frozenSourceGeneratedDocumentStates,
+        GeneratorDriverInitializationCache generatorDriverCreationCache,
         AsyncLazy<SolutionCompilationState>? cachedFrozenSnapshot = null)
     {
         SolutionState = solution;
         PartialSemanticsEnabled = partialSemanticsEnabled;
         _projectIdToTrackerMap = projectIdToTrackerMap;
         SourceGeneratorExecutionVersionMap = sourceGeneratorExecutionVersionMap;
         FrozenSourceGeneratedDocumentStates = frozenSourceGeneratedDocumentStates;
+        GeneratorDriverCache = generatorDriverCreationCache;
 
         // when solution state is changed, we recalculate its checksum
         _lazyChecksums = AsyncLazy.Create(static async (self, cancellationToken) =>
@@ -87,12 +91,14 @@ private SolutionCompilationState(
 
     public SolutionCompilationState(
         SolutionState solution,
-        bool partialSemanticsEnabled)
+        bool partialSemanticsEnabled,
+        GeneratorDriverInitializationCache generatorDriverCreationCache)
         : this(
               solution,
               partialSemanticsEnabled,
               projectIdToTrackerMap: ImmutableSegmentedDictionary<ProjectId, ICompilationTracker>.Empty,
               sourceGeneratorExecutionVersionMap: SourceGeneratorExecutionVersionMap.Empty,
+              generatorDriverCreationCache: generatorDriverCreationCache,
               frozenSourceGeneratedDocumentStates: TextDocumentStates<SourceGeneratedDocumentState>.Empty)
     {
     }
@@ -137,6 +143,7 @@ private SolutionCompilationState Branch(
             projectIdToTrackerMap.Value,
             sourceGeneratorExecutionVersionMap,
             frozenSourceGeneratedDocumentStates,
+            GeneratorDriverCache,
             cachedFrozenSnapshot);
     }
 
@@ -1524,6 +1531,11 @@ public SolutionCompilationState UpdateSpecificSourceGeneratorExecutionVersions(
                 if (newTracker != existingTracker)
                     newIdToTrackerMapBuilder[projectId] = newTracker;
             }
+
+            // Clear out the cache of any previously initialized GeneratorDriver. Otherwise we might reuse a
+            // driver which will not count as a new "run" in some of our unit tests. We have tests that very explicitly count
+            // and assert the number of invocations of a generator.
+            GeneratorDriverCache.EmptyCacheForProject(projectId);
         }
 
         if (!changed)

src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState_SourceGenerators.cs

@@ -60,6 +60,15 @@ private static ImmutableArray<ISourceGenerator> GetSourceGenerators(ProjectState
         return map is null ? [] : map.SourceGenerators;
     }
 
+    private static GeneratorDriver UpdateGeneratorDriverToMatchState(GeneratorDriver driver, ProjectState projectState)
+    {
+        return driver
+            .ReplaceAdditionalTexts(projectState.AdditionalDocumentStates.SelectAsArray(static documentState => documentState.AdditionalText))
+            .WithUpdatedParseOptions(projectState.ParseOptions!)
+            .WithUpdatedAnalyzerConfigOptions(projectState.ProjectAnalyzerOptions.AnalyzerConfigOptionsProvider)
+            .ReplaceGenerators(GetSourceGenerators(projectState));
+    }
+
     /// <summary>
     /// This method should only be called in a .net core host like our out of process server.
     /// </summary>

src/Workspaces/Core/Portable/Workspace/Workspace.cs

@@ -50,6 +50,11 @@ public abstract partial class Workspace : IDisposable
     // this lock guards all the mutable fields (do not share lock with derived classes)
     private readonly NonReentrantLock _stateLock = new(useThisInstanceForSynchronization: true);
 
+    /// <summary>
+    /// Cache for initializing generator drivers across different Solution instances from this Workspace.
+    /// </summary>
+    internal SolutionCompilationState.GeneratorDriverInitializationCache GeneratorDriverCreationCache { get; } = new();
+
     /// <summary>
     /// Current solution.  Must be locked with <see cref="_serializationLock"/> when writing to it.
     /// </summary>
@@ -275,6 +280,13 @@ internal bool SetCurrentSolution(
             {
                 data.onAfterUpdate?.Invoke(oldSolution, newSolution);
 
+                // The GeneratorDriverCreationCache holds onto a primordial GeneratorDriver for a project when we first creat one. That way, if another fork
+                // of the Solution also needs to run generators, it's able to reuse that primordial driver rather than recreating one from scratch. We want to
+                // clean up that cache at some point so we're not holding onto unneeded GeneratorDrivers. We'll clean out some cached entries here for projects
+                // that have a GeneratorDriver held in CurrentSolution. The idea being that once a project has a GeneratorDriver in the CurrentSolution, all future
+                // requests for generated documents will just use the updated generator that project already has, so there will never be another need to create one.
+                data.@this.GeneratorDriverCreationCache.EmptyCacheForProjectsThatHaveGeneratorDriversInSolution(newSolution.CompilationState);
+
                 // Queue the event but don't execute its handlers on this thread.
                 // Doing so under the serialization lock guarantees the same ordering of the events
                 // as the order of the changes made to the solution.
@@ -2370,7 +2382,7 @@ private static void CheckProjectIsInSolution(Solution solution, ProjectId projec
     }
 
     /// <summary>
-    /// Throws an exception is the project is part of the current solution.
+    /// Throws an exception if the project is part of the current solution.
     /// </summary>
     protected void CheckProjectIsNotInCurrentSolution(ProjectId projectId)
         => CheckProjectIsNotInSolution(this.CurrentSolution, projectId);