Prevent race in LSP pull diagnostics handling that causes stale diagnostics that do not update.

src/Features/LanguageServer/Protocol/Handler/Diagnostics/AbstractPullDiagnosticHandler.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Diagnostics;
@@ -101,7 +102,9 @@ private void OnDiagnosticsUpdated(object? sender, DiagnosticsUpdatedArgs updateA
             if (updateArgs.DocumentId == null)
                 return;
 
-            lock (_documentIdToLastResultId)
+            // Ensure we do not clear the cached results while the handler is reading (and possibly then writing)
+            // to the cached results.
+            lock (_gate)
             {
                 // Whenever we hear about changes to a document, drop the data we've stored for it.  We'll recompute it as
                 // necessary on the next request.
@@ -146,7 +149,12 @@ private void OnDiagnosticsUpdated(object? sender, DiagnosticsUpdatedArgs updateA
                     continue;
                 }
 
-                if (DiagnosticsAreUnchanged(documentToPreviousDiagnosticParams, document))
+                if (HaveDiagnosticsChangedAndUpdateCache(documentToPreviousDiagnosticParams, document, out var newResultId))
+                {
+                    context.TraceInformation($"Diagnostics were changed for document: {document.FilePath}");
+                    progress.Report(await ComputeAndReportCurrentDiagnosticsAsync(context, document, newResultId, cancellationToken).ConfigureAwait(false));
+                }
+                else
                 {
                     context.TraceInformation($"Diagnostics were unchanged for document: {document.FilePath}");
 
@@ -156,11 +164,6 @@ private void OnDiagnosticsUpdated(object? sender, DiagnosticsUpdatedArgs updateA
                     var previousParams = documentToPreviousDiagnosticParams[document];
                     progress.Report(CreateReport(previousParams.TextDocument, diagnostics: null, previousParams.PreviousResultId));
                 }
-                else
-                {
-                    context.TraceInformation($"Diagnostics were changed for document: {document.FilePath}");
-                    await ComputeAndReportCurrentDiagnosticsAsync(context, progress, document, cancellationToken).ConfigureAwait(false);
-                }
             }
 
             // If we had a progress object, then we will have been reporting to that.  Otherwise, take what we've been
@@ -199,10 +202,10 @@ private static Dictionary<Document, DiagnosticParams> GetDocumentToPreviousDiagn
             return result;
         }
 
-        private async Task ComputeAndReportCurrentDiagnosticsAsync(
+        private async Task<TReport> ComputeAndReportCurrentDiagnosticsAsync(
             RequestContext context,
-            BufferedProgress<TReport> progress,
             Document document,
+            string resultId,
             CancellationToken cancellationToken)
         {
             // Being asked about this document for the first time.  Or being asked again and we have different
@@ -230,7 +233,7 @@ private async Task ComputeAndReportCurrentDiagnosticsAsync(
                     result.Add(ConvertDiagnostic(document, text, diagnostic));
             }
 
-            progress.Report(RecordDiagnosticReport(document, result.ToArray()));
+            return CreateReport(ProtocolConversions.DocumentToTextDocumentIdentifier(document), result.ToArray(), resultId);
         }
 
         private void HandleRemovedDocuments(RequestContext context, DiagnosticParams[] previousResults, BufferedProgress<TReport> progress)
@@ -257,30 +260,45 @@ private void HandleRemovedDocuments(RequestContext context, DiagnosticParams[] p
             }
         }
 
-        private bool DiagnosticsAreUnchanged(Dictionary<Document, DiagnosticParams> documentToPreviousDiagnosticParams, Document document)
+        /// <summary>
+        /// Returns true if diagnostics have changed since the last request and if so,
+        /// calculates a new resultId to use for subsequent computation and caches it.
+        /// </summary>
+        /// <param name="documentToPreviousDiagnosticParams">the resultIds the client sent us.</param>
+        /// <param name="document">the document we are currently calculating results for.</param>
+        /// <param name="newResultId">the resultId to report new diagnostics with if changed.</param>
+        private bool HaveDiagnosticsChangedAndUpdateCache(
+            Dictionary<Document, DiagnosticParams> documentToPreviousDiagnosticParams,
+            Document document,
+            [NotNullWhen(true)] out string? newResultId)
         {
+            // Read and write the cached resultId to _documentIdToLastResultId in a single 'transaction'
+            // to prevent in-between updates to _documentIdToLastResultId triggered by OnDiagnosticsUpdated.
             lock (_gate)
             {
                 var workspace = document.Project.Solution.Workspace;
-                return documentToPreviousDiagnosticParams.TryGetValue(document, out var previousParams) &&
+                if (documentToPreviousDiagnosticParams.TryGetValue(document, out var previousParams) &&
                        _documentIdToLastResultId.TryGetValue((workspace, document.Id), out var lastReportedResultId) &&
-                       lastReportedResultId == previousParams.PreviousResultId;
-            }
-        }
+                       lastReportedResultId == previousParams.PreviousResultId)
+                {
+                    // Our cached resultId for the document matches the resultId the client passed to us.
+                    // This means the diagnostics have not changed and we do not need to re-compute.
+                    newResultId = null;
+                    return false;
+                }
 
-        private TReport RecordDiagnosticReport(Document document, VSDiagnostic[] diagnostics)
-        {
-            lock (_gate)
-            {
                 // Keep track of the diagnostics we reported here so that we can short-circuit producing diagnostics for
                 // the same diagnostic set in the future.  Use a custom result-id per type (doc diagnostics or workspace
                 // diagnostics) so that clients of one don't errantly call into the other.  For example, a client
                 // getting document diagnostics should not ask for workspace diagnostics with the result-ids it got for
                 // doc-diagnostics.  The two systems are different and cannot share results, or do things like report
                 // what changed between each other.
-                var resultId = $"{GetType().Name}:{_nextDocumentResultId++}";
-                _documentIdToLastResultId[(document.Project.Solution.Workspace, document.Id)] = resultId;
-                return CreateReport(ProtocolConversions.DocumentToTextDocumentIdentifier(document), diagnostics, resultId);
+                //
+                // Note that we can safely update the map before computation as any cancellation or exception
+                // during computation means that the client will never recieve this resultId and so cannot ask us for it.
+                newResultId = $"{GetType().Name}:{_nextDocumentResultId++}";
+                _documentIdToLastResultId[(document.Project.Solution.Workspace, document.Id)] = newResultId;
+                return true;
             }
         }