Merge pull request #54558 from dibarbet/lsp_diagnostics_race

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;
@@ -100,7 +101,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.
@@ -145,7 +148,12 @@ private void OnDiagnosticsUpdated(object? sender, DiagnosticsUpdatedArgs updateA
  continue;
  }
 
- if (DiagnosticsAreUnchanged(documentToPreviousDiagnosticParams, document))
+ if (HaveDiagnosticsChanged(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}");
 
@@ -155,11 +163,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
@@ -198,10 +201,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
@@ -229,7 +232,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)
@@ -256,30 +259,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 HaveDiagnosticsChanged(
+ 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;
  }
  }