Remove the cancellation token from the calculation of InProgressState…

….LazyStaleCompilationWithGeneratedDocuments (dotnet#75001)

* Remove strong the cancellation token used in the calculation of InProgressState.LazyStaleCompilationWithGeneratedDocuments

This will hopefully address multiple recent feedback tickets around Go To Definition occasionally failing. The umbrella pri 0 feedback ticket is https://developercommunity.visualstudio.com/t/Go-to-definition-doesnt-work/10729503

I was able to reproduce this by following the repro outlined by Garry in the feedback ticket, using the "Code Search" window and doing go to definition multiple times.

The issue here is that RegularCompilationTracker.FinalizeCompilationWorkerAsync tries to get the value from the Lazy inProgressState.LazyStaleCompilationWithGeneratedDocuments. This lazy holds onto an action that uses a CancellationToken that has already been cancelled, whereas the CancellationToken used by FinalizeCompilationWorkerAsync has not been cancelled. When the action for the lazy is performed, the corresponding OperationCanceledException ends up escaping out of FinalizeCompilationWorker due to the FatalError.ReportAndPropagateUnlessCanceled exception clause failing due to the cancellation token passed into FinalizeCompilationWorkerAsync not being in a cancelled state.

The offending Lazy is constructed by RegularCompilationTracker.WithDoNotCreateCreationPolicy's call to the InProgressState ctor. Instead of passing a closure capturing a cancellation token, I've changed the code to instead use CancellableLazy, which takes in a cancellation token at the time the value is requested.

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

@@ -72,7 +72,7 @@ private sealed class InProgressState : CompilationTrackerState
             /// correct snapshot in that the generators have not been rerun, but may be reusable if the generators
             /// are later found to give the same output.
             /// </summary>
-            public readonly Lazy<Compilation?> LazyStaleCompilationWithGeneratedDocuments;
+            public readonly CancellableLazy<Compilation?> LazyStaleCompilationWithGeneratedDocuments;
 
             /// <summary>
             /// The list of changes that have happened since we last computed a compilation. The oldState corresponds to
@@ -86,7 +86,7 @@ public InProgressState(
                 CreationPolicy creationPolicy,
                 Lazy<Compilation> compilationWithoutGeneratedDocuments,
                 CompilationTrackerGeneratorInfo generatorInfo,
-                Lazy<Compilation?> staleCompilationWithGeneratedDocuments,
+                CancellableLazy<Compilation?> staleCompilationWithGeneratedDocuments,
                 ImmutableList<TranslationAction> pendingTranslationActions)
                 : base(creationPolicy, generatorInfo)
             {
@@ -123,8 +123,8 @@ public InProgressState(
             {
             }
 
-            private static Lazy<Compilation?> CreateLazyCompilation(Compilation? staleCompilationWithGeneratedDocuments)
-                => new(() => staleCompilationWithGeneratedDocuments);
+            private static CancellableLazy<Compilation?> CreateLazyCompilation(Compilation? staleCompilationWithGeneratedDocuments)
+                => new(staleCompilationWithGeneratedDocuments);
         }
 
         /// <summary>

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

@@ -49,7 +49,7 @@ bool ContainsAssemblyOrModuleOrDynamic(
         /// <summary>
         /// Updates the creation policy for this tracker.  Setting it to <see cref="CreationPolicy.DoNotCreate"/>.
         /// </summary>
-        ICompilationTracker WithDoNotCreateCreationPolicy(CancellationToken cancellationToken);
+        ICompilationTracker WithDoNotCreateCreationPolicy();
 
         Task<VersionStamp> GetDependentVersionAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken);
         Task<VersionStamp> GetDependentSemanticVersionAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken);

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

@@ -32,7 +32,7 @@ private sealed partial class RegularCompilationTracker : ICompilationTracker
             private static readonly Func<ProjectState, string> s_logBuildCompilationAsync =
                 state => string.Join(",", state.AssemblyName, state.DocumentStates.Count);
 
-            private static readonly Lazy<Compilation?> s_lazyNullCompilation = new Lazy<Compilation?>(() => null);
+            private static readonly CancellableLazy<Compilation?> s_lazyNullCompilation = new CancellableLazy<Compilation?>((Compilation?)null);
 
             public ProjectState ProjectState { get; }
 
@@ -145,7 +145,7 @@ public ICompilationTracker Fork(
                     var (compilationWithoutGeneratedDocuments, staleCompilationWithGeneratedDocuments) = state switch
                     {
                         InProgressState inProgressState => (inProgressState.LazyCompilationWithoutGeneratedDocuments, inProgressState.LazyStaleCompilationWithGeneratedDocuments),
-                        FinalCompilationTrackerState finalState => (new Lazy<Compilation>(() => finalState.CompilationWithoutGeneratedDocuments), new Lazy<Compilation?>(() => finalState.FinalCompilationWithGeneratedDocuments)),
+                        FinalCompilationTrackerState finalState => (new Lazy<Compilation>(() => finalState.CompilationWithoutGeneratedDocuments), new CancellableLazy<Compilation?>(finalState.FinalCompilationWithGeneratedDocuments)),
                         _ => throw ExceptionUtilities.UnexpectedValue(state.GetType()),
                     };
 
@@ -404,7 +404,7 @@ async Task<InProgressState> CollapseInProgressStateAsync(InProgressState initial
                         var translationAction = inProgressState.PendingTranslationActions[0];
 
                         var compilationWithoutGeneratedDocuments = inProgressState.CompilationWithoutGeneratedDocuments;
-                        var staleCompilationWithGeneratedDocuments = inProgressState.LazyStaleCompilationWithGeneratedDocuments.Value;
+                        var staleCompilationWithGeneratedDocuments = inProgressState.LazyStaleCompilationWithGeneratedDocuments.GetValue(cancellationToken);
 
                         // If staleCompilationWithGeneratedDocuments is the same as compilationWithoutGeneratedDocuments,
                         // then it means a prior run of generators didn't produce any files. In that case, we'll just make
@@ -476,7 +476,7 @@ async Task<FinalCompilationTrackerState> FinalizeCompilationWorkerAsync(InProgre
                     var creationPolicy = inProgressState.CreationPolicy;
                     var generatorInfo = inProgressState.GeneratorInfo;
                     var compilationWithoutGeneratedDocuments = inProgressState.CompilationWithoutGeneratedDocuments;
-                    var staleCompilationWithGeneratedDocuments = inProgressState.LazyStaleCompilationWithGeneratedDocuments.Value;
+                    var staleCompilationWithGeneratedDocuments = inProgressState.LazyStaleCompilationWithGeneratedDocuments.GetValue(cancellationToken);
 
                     // Project is complete only if the following are all true:
                     //  1. HasAllInformation flag is set for the project
@@ -735,7 +735,7 @@ public ICompilationTracker WithCreateCreationPolicy(bool forceRegeneration)
                     skeletonReferenceCacheToClone: _skeletonReferenceCache);
             }
 
-            public ICompilationTracker WithDoNotCreateCreationPolicy(CancellationToken cancellationToken)
+            public ICompilationTracker WithDoNotCreateCreationPolicy()
             {
                 var state = this.ReadState();
 
@@ -792,8 +792,7 @@ public ICompilationTracker WithDoNotCreateCreationPolicy(CancellationToken cance
                     var alreadyParsedTrees = alreadyParsedTreesBuilder.ToImmutableAndClear();
                     var lazyCompilationWithoutGeneratedDocuments = new Lazy<Compilation>(() => this.CreateEmptyCompilation().AddSyntaxTrees(alreadyParsedTrees));
 
-                    // Safe cast to appease NRT system.
-                    var lazyCompilationWithGeneratedDocuments = (Lazy<Compilation?>)lazyCompilationWithoutGeneratedDocuments!;
+                    var lazyCompilationWithGeneratedDocuments = new CancellableLazy<Compilation?>(cancellationToken => lazyCompilationWithoutGeneratedDocuments.Value);
 
                     return new RegularCompilationTracker(
                         frozenProjectState,
@@ -818,8 +817,12 @@ public ICompilationTracker WithDoNotCreateCreationPolicy(CancellationToken cance
                     // us to a frozen state with that information.
                     var generatorInfo = inProgressState.GeneratorInfo;
                     var compilationWithoutGeneratedDocuments = inProgressState.LazyCompilationWithoutGeneratedDocuments;
-                    var compilationWithGeneratedDocuments = new Lazy<Compilation?>(() => compilationWithoutGeneratedDocuments.Value.AddSyntaxTrees(
-                        generatorInfo.Documents.States.Values.Select(state => state.GetSyntaxTree(cancellationToken))));
+
+                    var compilationWithGeneratedDocuments = new CancellableLazy<Compilation?>(cancellationToken =>
+                    {
+                        var syntaxTrees = generatorInfo.Documents.States.Values.Select(state => state.GetSyntaxTree(cancellationToken));
+                        return compilationWithoutGeneratedDocuments.Value.AddSyntaxTrees(syntaxTrees);
+                    });
 
                     return new RegularCompilationTracker(
                         frozenProjectState,
@@ -931,9 +934,9 @@ private void ValidateState(CompilationTrackerState? state)
 
                     ValidateCompilationTreesMatchesProjectState(inProgressState.CompilationWithoutGeneratedDocuments, projectState, generatorInfo: null);
 
-                    if (inProgressState.LazyStaleCompilationWithGeneratedDocuments.Value != null)
+                    if (inProgressState.LazyStaleCompilationWithGeneratedDocuments.GetValue(CancellationToken.None) is Compilation staleCompilationWithGeneratedDocuments)
                     {
-                        ValidateCompilationTreesMatchesProjectState(inProgressState.LazyStaleCompilationWithGeneratedDocuments.Value, projectState, inProgressState.GeneratorInfo);
+                        ValidateCompilationTreesMatchesProjectState(staleCompilationWithGeneratedDocuments, projectState, inProgressState.GeneratorInfo);
                     }
                 }
                 else

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

@@ -89,9 +89,9 @@ public ICompilationTracker WithCreateCreationPolicy(bool forceRegeneration)
                 : new WithFrozenSourceGeneratedDocumentsCompilationTracker(underlyingTracker, _replacementDocumentStates);
         }
 
-        public ICompilationTracker WithDoNotCreateCreationPolicy(CancellationToken cancellationToken)
+        public ICompilationTracker WithDoNotCreateCreationPolicy()
         {
-            var underlyingTracker = this.UnderlyingTracker.WithDoNotCreateCreationPolicy(cancellationToken);
+            var underlyingTracker = this.UnderlyingTracker.WithDoNotCreateCreationPolicy();
             return underlyingTracker == this.UnderlyingTracker
                 ? this
                 : new WithFrozenSourceGeneratedDocumentsCompilationTracker(underlyingTracker, _replacementDocumentStates);

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

@@ -1481,7 +1481,7 @@ private SolutionCompilationState ComputeFrozenSnapshot(CancellationToken cancell
 
             // Since we're freezing, set both generators and skeletons to not be created.  We don't want to take any
             // perf hit on either of those at all for our clients.
-            var newTracker = oldTracker.WithDoNotCreateCreationPolicy(cancellationToken);
+            var newTracker = oldTracker.WithDoNotCreateCreationPolicy();
             if (oldTracker == newTracker)
                 continue;