Add mitigations for SG documents not being present

dotnet · Sep 15, 2023 · f6d6c20 · f6d6c20
1 parent b19ba09
commit f6d6c20
Show file tree

Hide file tree

Showing 14 changed files with 171 additions and 25 deletions.
diff --git a/src/Features/Core/Portable/DocumentHighlighting/IRemoteDocumentHighlightsService.cs b/src/Features/Core/Portable/DocumentHighlighting/IRemoteDocumentHighlightsService.cs
@@ -31,8 +31,17 @@ public SerializableDocumentHighlights(DocumentId documentId, ImmutableArray<High
             HighlightSpans = highlightSpans;
         }
 
-        public async ValueTask<DocumentHighlights> RehydrateAsync(Solution solution)
-            => new(await solution.GetRequiredDocumentAsync(DocumentId, includeSourceGenerated: true).ConfigureAwait(false), HighlightSpans);
+        public async ValueTask<DocumentHighlights?> RehydrateAsync(Solution solution, CancellationToken cancellationToken)
+        {
+            // https://github.com/dotnet/roslyn/issues/69964
+            //
+            // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+            var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(DocumentId, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+            if (document is null)
+                return null;
+
+            return new(document, HighlightSpans);
+        }
 
         public static SerializableDocumentHighlights Dehydrate(DocumentHighlights highlights)
             => new(highlights.Document.Id, highlights.HighlightSpans);

diff --git a/src/Features/Core/Portable/NavigateTo/RoslynNavigateToItem.cs b/src/Features/Core/Portable/NavigateTo/RoslynNavigateToItem.cs
@@ -91,8 +91,14 @@ public RoslynNavigateToItem(
             }
             else
             {
-                var document = await solution.GetRequiredDocumentAsync(
-                    DocumentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+                // https://github.com/dotnet/roslyn/issues/69964
+                //
+                // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+                var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(
+                    DocumentId, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+                if (document == null)
+                    return null;
+
                 return new NavigateToSearchResult(this, document, activeDocument);
             }
         }

diff --git a/src/Features/Core/Portable/Navigation/INavigableItem.cs b/src/Features/Core/Portable/Navigation/INavigableItem.cs
@@ -65,18 +65,30 @@ public static NavigableDocument FromDocument(Document document)
             /// this navigable item. The document is required to exist within the solution, e.g. a case where the
             /// navigable item was constructed during a Find Symbols operation on the same solution instance.
             /// </summary>
-            internal ValueTask<Document> GetRequiredDocumentAsync(Solution solution, CancellationToken cancellationToken)
-                => solution.GetRequiredDocumentAsync(Id, includeSourceGenerated: IsSourceGeneratedDocument, cancellationToken);
+            internal async ValueTask<Document?> GetRequiredDocumentAsync(Solution solution, CancellationToken cancellationToken)
+            {
+                if (!IsSourceGeneratedDocument)
+                    return solution.GetRequiredDocument(Id);
+
+                // https://github.com/dotnet/roslyn/issues/69964
+                //
+                // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+                return await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(
+                    Id, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+            }
 
             /// <summary>
             /// Get the <see cref="SourceText"/> of the <see cref="CodeAnalysis.Document"/> within
             /// <paramref name="solution"/> which is referenced by this navigable item. The document is required to
             /// exist within the solution, e.g. a case where the navigable item was constructed during a Find Symbols
             /// operation on the same solution instance.
             /// </summary>
-            internal async ValueTask<SourceText> GetTextAsync(Solution solution, CancellationToken cancellationToken)
+            internal async ValueTask<SourceText?> GetTextAsync(Solution solution, CancellationToken cancellationToken)
             {
                 var document = await GetRequiredDocumentAsync(solution, cancellationToken).ConfigureAwait(false);
+                if (document is null)
+                    return null;
+
                 return await document.GetValueTextAsync(cancellationToken).ConfigureAwait(false);
             }
 

diff --git a/src/Workspaces/Core/Portable/FindSymbols/FindReferences/DependentTypeFinder.cs b/src/Workspaces/Core/Portable/FindSymbols/FindReferences/DependentTypeFinder.cs
@@ -249,7 +249,15 @@ async Task AddMatchingTypesAsync(
                         cancellationToken.ThrowIfCancellationRequested();
 
                         Debug.Assert(infos.Count > 0);
-                        var document = await solution.GetRequiredDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+
+                        // https://github.com/dotnet/roslyn/issues/69964
+                        //
+                        // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+                        var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(
+                            documentId, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+                        if (document is null)
+                            continue;
+
                         var semanticModel = await document.GetRequiredSemanticModelAsync(cancellationToken).ConfigureAwait(false);
                         cachedModels.Add(semanticModel);
 
@@ -276,7 +284,14 @@ async Task AddSourceTypesThatDeriveFromNameAsync(SymbolSet result, string name)
                     {
                         cancellationToken.ThrowIfCancellationRequested();
 
-                        var document = await solution.GetRequiredDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+                        // https://github.com/dotnet/roslyn/issues/69964
+                        //
+                        // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+                        var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(
+                            documentId, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+                        if (document is null)
+                            continue;
+
                         var semanticModel = await document.GetRequiredSemanticModelAsync(cancellationToken).ConfigureAwait(false);
                         cachedModels.Add(semanticModel);
 

diff --git a/src/Workspaces/Core/Portable/FindSymbols/SymbolFinder.FindReferencesServerCallback.cs b/src/Workspaces/Core/Portable/FindSymbols/SymbolFinder.FindReferencesServerCallback.cs
@@ -42,13 +42,27 @@ public ValueTask OnCompletedAsync(CancellationToken cancellationToken)
 
             public async ValueTask OnFindInDocumentStartedAsync(DocumentId documentId, CancellationToken cancellationToken)
             {
-                var document = await solution.GetRequiredDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+                // https://github.com/dotnet/roslyn/issues/69964
+                //
+                // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+                var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(
+                    documentId, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+                if (document is null)
+                    return;
+
                 await progress.OnFindInDocumentStartedAsync(document, cancellationToken).ConfigureAwait(false);
             }
 
             public async ValueTask OnFindInDocumentCompletedAsync(DocumentId documentId, CancellationToken cancellationToken)
             {
-                var document = await solution.GetRequiredDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+                // https://github.com/dotnet/roslyn/issues/69964
+                //
+                // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+                var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(
+                    documentId, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+                if (document is null)
+                    return;
+
                 await progress.OnFindInDocumentCompletedAsync(document, cancellationToken).ConfigureAwait(false);
             }
 
@@ -102,10 +116,11 @@ public async ValueTask OnReferenceFoundAsync(
                     }
                 }
 
-                var referenceLocation = await reference.RehydrateAsync(
-                    solution, cancellationToken).ConfigureAwait(false);
+                var referenceLocation = await reference.RehydrateAsync(solution, cancellationToken).ConfigureAwait(false);
+                if (referenceLocation is null)
+                    return;
 
-                await progress.OnReferenceFoundAsync(symbolGroup, symbol, referenceLocation, cancellationToken).ConfigureAwait(false);
+                await progress.OnReferenceFoundAsync(symbolGroup, symbol, referenceLocation.Value, cancellationToken).ConfigureAwait(false);
             }
         }
     }

diff --git a/src/Workspaces/Core/Portable/Remote/RemoteArguments.cs b/src/Workspaces/Core/Portable/Remote/RemoteArguments.cs
@@ -169,10 +169,17 @@ public static SerializableReferenceLocation Dehydrate(
                 referenceLocation.CandidateReason);
         }
 
-        public async ValueTask<ReferenceLocation> RehydrateAsync(
+        public async ValueTask<ReferenceLocation?> RehydrateAsync(
             Solution solution, CancellationToken cancellationToken)
         {
-            var document = await solution.GetRequiredDocumentAsync(this.Document, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+            // https://github.com/dotnet/roslyn/issues/69964
+            //
+            // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+            var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(
+                this.Document, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+            if (document is null)
+                return null;
+
             var syntaxTree = await document.GetRequiredSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
             var aliasSymbol = await RehydrateAliasAsync(solution, cancellationToken).ConfigureAwait(false);
             var additionalProperties = this.AdditionalProperties;

diff --git a/src/Workspaces/Core/Portable/Rename/IRemoteRenamerService.cs b/src/Workspaces/Core/Portable/Rename/IRemoteRenamerService.cs
@@ -162,7 +162,7 @@ internal partial class SymbolicRenameLocations
                 serializableLocations.Options,
                 fallbackOptions,
                 locations,
-                implicitLocations,
+                implicitLocations.WhereNotNull(),
                 referencedSymbols);
         }
     }

diff --git a/src/Workspaces/Core/Portable/Rename/LightweightRenameLocations.cs b/src/Workspaces/Core/Portable/Rename/LightweightRenameLocations.cs
@@ -67,7 +67,7 @@ private LightweightRenameLocations(
                 Options,
                 FallbackOptions,
                 Locations,
-                implicitLocations,
+                implicitLocations.WhereNotNull(),
                 referencedSymbols);
         }
 

diff --git a/...rkspaces/Remote/ServiceHub/Services/DocumentHighlights/RemoteDocumentHighlightsService.cs b/...rkspaces/Remote/ServiceHub/Services/DocumentHighlights/RemoteDocumentHighlightsService.cs
@@ -35,7 +35,14 @@ public ValueTask<ImmutableArray<SerializableDocumentHighlights>> GetDocumentHigh
             // need to be revisited if we someday support FAR between these languages.
             return RunServiceAsync(solutionChecksum, async solution =>
             {
-                var document = await solution.GetRequiredDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+                // https://github.com/dotnet/roslyn/issues/69964
+                //
+                // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+                var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(
+                    documentId, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+                if (document is null)
+                    return ImmutableArray<SerializableDocumentHighlights>.Empty;
+
                 var documentsToSearch = await documentIdsToSearch.SelectAsArrayAsync(id => solution.GetDocumentAsync(id, includeSourceGenerated: true, cancellationToken)).ConfigureAwait(false);
                 var documentsToSearchSet = ImmutableHashSet.CreateRange(documentsToSearch.WhereNotNull());
 

diff --git a/...Workspaces/Remote/ServiceHub/Services/InheritanceMargin/RemoteInheritanceMarginService.cs b/...Workspaces/Remote/ServiceHub/Services/InheritanceMargin/RemoteInheritanceMarginService.cs
@@ -42,7 +42,14 @@ public ValueTask<ImmutableArray<InheritanceMarginItem>> GetInheritanceMarginItem
                 //
                 // Tracked by https://github.com/dotnet/roslyn/issues/67065.
                 frozenPartialSemantics = false;
-                var document = await solution.GetRequiredDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+
+                // https://github.com/dotnet/roslyn/issues/69964
+                //
+                // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+                var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(documentId, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+                if (document is null)
+                    return ImmutableArray<InheritanceMarginItem>.Empty;
+
                 var service = document.GetRequiredLanguageService<IInheritanceMarginService>();
                 return await service.GetInheritanceMemberItemsAsync(document, spanToSearch, includeGlobalImports, frozenPartialSemantics, cancellationToken).ConfigureAwait(false);
             }, cancellationToken);

diff --git a/src/Workspaces/Remote/ServiceHub/Services/NavigationBar/RemoteNavigationBarItemService.cs b/src/Workspaces/Remote/ServiceHub/Services/NavigationBar/RemoteNavigationBarItemService.cs
@@ -29,8 +29,12 @@ public ValueTask<ImmutableArray<SerializableNavigationBarItem>> GetItemsAsync(
         {
             return RunServiceAsync(solutionChecksum, async solution =>
             {
-                var document = await solution.GetDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
-                Contract.ThrowIfNull(document);
+                // https://github.com/dotnet/roslyn/issues/69964
+                //
+                // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+                var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(documentId, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+                if (document is null)
+                    return ImmutableArray<SerializableNavigationBarItem>.Empty;
 
                 if (forceFrozenPartialSemanticsForCrossProcessOperations)
                 {

diff --git a/src/Workspaces/Remote/ServiceHub/Services/TaskList/RemoteTaskListService.cs b/src/Workspaces/Remote/ServiceHub/Services/TaskList/RemoteTaskListService.cs
@@ -30,7 +30,14 @@ public ValueTask<ImmutableArray<TaskListItem>> GetTaskListItemsAsync(
         {
             return RunServiceAsync(solutionChecksum, async solution =>
             {
-                var document = await solution.GetRequiredDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+                // https://github.com/dotnet/roslyn/issues/69964
+                //
+                // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+                var document = await solution.GetRequiredDocumentIncludingSourceGeneratedAsync(
+                    documentId, throwForMissingSourceGenerated: false, cancellationToken).ConfigureAwait(false);
+                if (document is null)
+                    return ImmutableArray<TaskListItem>.Empty;
+
                 var service = document.GetRequiredLanguageService<ITaskListService>();
                 return await service.GetTaskListItemsAsync(document, descriptors, cancellationToken).ConfigureAwait(false);
             }, cancellationToken);

diff --git a/...kspaces/SharedUtilitiesAndExtensions/Compiler/Core/Extensions/ImmutableArrayExtensions.cs b/...kspaces/SharedUtilitiesAndExtensions/Compiler/Core/Extensions/ImmutableArrayExtensions.cs
@@ -35,5 +35,20 @@ public static ImmutableArray<T> TakeAsArray<T>(this ImmutableArray<T> array, int
 
             return result.ToImmutableAndClear();
         }
+
+        public static ImmutableArray<T> WhereNotNull<T>(this ImmutableArray<T?> array) where T : struct
+        {
+            var count = array.Count(static t => t != null);
+
+            using var _ = ArrayBuilder<T>.GetInstance(count, out var result);
+
+            foreach (var value in array)
+            {
+                if (value != null)
+                    result.Add(value.Value);
+            }
+
+            return result.ToImmutableAndClear();
+        }
     }
 }
diff --git a/src/Workspaces/SharedUtilitiesAndExtensions/Workspace/Core/Extensions/ISolutionExtensions.cs b/src/Workspaces/SharedUtilitiesAndExtensions/Workspace/Core/Extensions/ISolutionExtensions.cs
@@ -4,8 +4,10 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.ErrorReporting;
 
 namespace Microsoft.CodeAnalysis.Shared.Extensions
 {
@@ -48,8 +50,48 @@ public static Document GetRequiredDocument(this Solution solution, DocumentId do
             => solution.GetDocument(documentId) ?? throw CreateDocumentNotFoundException();
 
 #if !CODE_STYLE
-        public static async ValueTask<Document> GetRequiredDocumentAsync(this Solution solution, DocumentId documentId, bool includeSourceGenerated = false, CancellationToken cancellationToken = default)
-            => (await solution.GetDocumentAsync(documentId, includeSourceGenerated, cancellationToken).ConfigureAwait(false)) ?? throw CreateDocumentNotFoundException();
+        public static async ValueTask<Document?> GetRequiredDocumentIncludingSourceGeneratedAsync(
+            this Solution solution,
+            DocumentId documentId,
+            bool throwForMissingSourceGenerated = true,
+            CancellationToken cancellationToken = default)
+        {
+            var document = await solution.GetDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+
+            // https://github.com/dotnet/roslyn/issues/69964
+            //
+            // Remove this once we solve root cause issue of the hosts disagreeing on source generated documents.
+            if (document is null)
+            {
+                if (documentId.IsSourceGenerated && !throwForMissingSourceGenerated)
+                {
+                    // Create a crash report so we can better hunt this down.
+                    try
+                    {
+                        throw CreateDocumentNotFoundException();
+                    }
+                    catch (Exception ex) when (FatalError.ReportAndCatch(ex, ErrorSeverity.Critical))
+                    {
+                    }
+
+                    return null;
+                }
+
+                throw CreateDocumentNotFoundException();
+            }
+
+            return document;
+        }
+
+        public static async ValueTask<Document> GetRequiredDocumentAsync(
+            this Solution solution,
+            DocumentId documentId,
+            bool includeSourceGenerated = false,
+            CancellationToken cancellationToken = default)
+        {
+            return await solution.GetDocumentAsync(documentId, includeSourceGenerated, cancellationToken).ConfigureAwait(false) ??
+                throw CreateDocumentNotFoundException();
+        }
 
         public static async ValueTask<TextDocument> GetRequiredTextDocumentAsync(this Solution solution, DocumentId documentId, CancellationToken cancellationToken = default)
             => (await solution.GetTextDocumentAsync(documentId, cancellationToken).ConfigureAwait(false)) ?? throw CreateDocumentNotFoundException();