Fix race causing workspace diagnostics to be stale

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

@@ -106,7 +106,7 @@ protected abstract ValueTask<ImmutableArray<IDiagnosticSource>> GetOrderedDiagno
     /// Used by public workspace pull diagnostics to allow it to keep the connection open until
     /// changes occur to avoid the client spamming the server with requests.
     /// </summary>
-    protected virtual Task WaitForChangesAsync(RequestContext context, CancellationToken cancellationToken)
+    protected virtual Task WaitForChangesAsync(string? category, RequestContext context, CancellationToken cancellationToken)
     {
         return Task.CompletedTask;
     }
@@ -221,7 +221,7 @@ await ComputeAndReportCurrentDiagnosticsAsync(
             // Some implementations of the spec will re-open requests as soon as we close them, spamming the server.
             // In those cases, we wait for the implementation to indicate that changes have occurred, then we close the connection
             // so that the client asks us again.
-            await WaitForChangesAsync(context, cancellationToken).ConfigureAwait(false);
+            await WaitForChangesAsync(category, context, cancellationToken).ConfigureAwait(false);
 
             // If we had a progress object, then we will have been reporting to that.  Otherwise, take what we've been
             // collecting and return that.

src/LanguageServer/Protocol/Handler/Diagnostics/AbstractWorkspacePullDiagnosticsHandler.cs

@@ -3,7 +3,9 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Diagnostics;
@@ -22,12 +24,15 @@ internal abstract class AbstractWorkspacePullDiagnosticsHandler<TDiagnosticsPara
     protected readonly IDiagnosticSourceManager DiagnosticSourceManager;
 
     /// <summary>
-    /// Flag that represents whether the LSP view of the world has changed.
-    /// It is totally fine for this to somewhat over-report changes
-    /// as it is an optimization used to delay closing workspace diagnostic requests
-    /// until something has changed.
+    /// Gate to guard access to <see cref="_categoryToLspChanged"/>
     /// </summary>
-    private int _lspChanged = 0;
+    private readonly object _gate = new();
+
+    /// <summary>
+    /// Stores the LSP changed state on a per category basis.  This ensures that requests for different categories
+    /// are 'walled off' from each other and only reset state for their own category.
+    /// </summary>
+    private readonly Dictionary<string, bool> _categoryToLspChanged = new();
 
     protected AbstractWorkspacePullDiagnosticsHandler(
         LspWorkspaceManager workspaceManager,
@@ -76,23 +81,50 @@ private void OnLspTextChanged(object? sender, EventArgs e)
 
     private void UpdateLspChanged()
     {
-        Interlocked.Exchange(ref _lspChanged, 1);
+        lock (_gate)
+        {
+            // Loop through our map of source -> has changed and mark them as all having changed.
+            foreach (var category in _categoryToLspChanged.Keys.ToImmutableArray())
+            {
+                _categoryToLspChanged[category] = true;
+            }
+        }
     }
 
-    protected override async Task WaitForChangesAsync(RequestContext context, CancellationToken cancellationToken)
+    protected override async Task WaitForChangesAsync(string? category, RequestContext context, CancellationToken cancellationToken)
     {
+        // A null category counts a separate category and should track changes independently of other categories, so we'll add an empty entry in our map for it.
+        category ??= string.Empty;
+
         // Spin waiting until our LSP change flag has been set.  When the flag is set (meaning LSP has changed),
         // we reset the flag to false and exit out of the loop allowing the request to close.
         // The client will automatically trigger a new request as soon as we close it, bringing us up to date on diagnostics.
-        while (Interlocked.CompareExchange(ref _lspChanged, value: 0, comparand: 1) == 0)
+        while (!HasChanged())
         {
             // There have been no changes between now and when the last request finished - we will hold the connection open while we poll for changes.
             await Task.Delay(TimeSpan.FromMilliseconds(100), cancellationToken).ConfigureAwait(false);
         }
 
-        context.TraceInformation("Closing workspace/diagnostics request");
         // We've hit a change, so we close the current request to allow the client to open a new one.
+        context.TraceInformation("Closing workspace/diagnostics request");
         return;
+
+        bool HasChanged()
+        {
+            lock (_gate)
+            {
+                // Get the LSP changed value of this category.  If it doesn't exist we add it with a value of 'changed' since this is the first
+                // request for the category and we don't know if it has changed since the request started.
+                var changed = _categoryToLspChanged.GetOrAdd(category, true);
+                if (changed)
+                {
+                    // We've observed a change, so we reset the flag to false for this source and return true.
+                    _categoryToLspChanged[category] = false;
+                }
+
+                return changed;
+            }
+        }
     }
 
     internal abstract TestAccessor GetTestAccessor();

src/LanguageServer/ProtocolUnitTests/Diagnostics/PullDiagnosticTests.cs

@@ -1925,7 +1925,12 @@ public async Task TestWorkspaceDiagnosticsWaitsForLspTextChanges(bool useVSDiagn
         await using var testLspServer = await CreateTestWorkspaceWithDiagnosticsAsync(
             [markup1, markup2], mutatingLspWorkspace, BackgroundAnalysisScope.FullSolution, useVSDiagnostics);
 
+        // The very first request should return immediately (as we're have no prior state to tell if the sln changed).
         var resultTask = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics, useProgress: true, triggerConnectionClose: false);
+        await resultTask;
+
+        // The second request should wait for a solution change before returning.
+        resultTask = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics, useProgress: true, triggerConnectionClose: false);
 
         // Assert that the connection isn't closed and task doesn't complete even after some delay.
         await Task.Delay(TimeSpan.FromSeconds(5));
@@ -1949,7 +1954,12 @@ public async Task TestWorkspaceDiagnosticsWaitsForLspSolutionChanges(bool useVSD
         await using var testLspServer = await CreateTestWorkspaceWithDiagnosticsAsync(
             [markup1, markup2], mutatingLspWorkspace, BackgroundAnalysisScope.FullSolution, useVSDiagnostics);
 
+        // The very first request should return immediately (as we're have no prior state to tell if the sln changed).
         var resultTask = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics, useProgress: true, triggerConnectionClose: false);
+        await resultTask;
+
+        // The second request should wait for a solution change before returning.
+        resultTask = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics, useProgress: true, triggerConnectionClose: false);
 
         // Assert that the connection isn't closed and task doesn't complete even after some delay.
         await Task.Delay(TimeSpan.FromSeconds(5));
@@ -1964,6 +1974,58 @@ public async Task TestWorkspaceDiagnosticsWaitsForLspSolutionChanges(bool useVSD
         Assert.NotEmpty(results);
     }
 
+    [Theory, CombinatorialData]
+    public async Task TestWorkspaceDiagnosticsWaitsForLspTextChangesWithMultipleSources(bool useVSDiagnostics, bool mutatingLspWorkspace)
+    {
+        var markup1 =
+@"class A {";
+        var markup2 = "";
+        await using var testLspServer = await CreateTestWorkspaceWithDiagnosticsAsync(
+            [markup1, markup2], mutatingLspWorkspace, BackgroundAnalysisScope.FullSolution, useVSDiagnostics);
+
+        var resultTaskOne = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics, useProgress: true, category: PullDiagnosticCategories.WorkspaceDocumentsAndProject, triggerConnectionClose: false);
+        var resultTaskTwo = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics, useProgress: true, category: PullDiagnosticCategories.EditAndContinue, triggerConnectionClose: false);
+        // The very first requests should return immediately (as we're have no prior state to tell if the sln changed).
+        await Task.WhenAll(resultTaskOne, resultTaskTwo);
+
+        // The second request for each source should wait for a solution change before returning.
+        resultTaskOne = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics, useProgress: true, category: PullDiagnosticCategories.WorkspaceDocumentsAndProject, triggerConnectionClose: false);
+        resultTaskTwo = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics, useProgress: true, category: PullDiagnosticCategories.EditAndContinue, triggerConnectionClose: false);
+
+        // Assert that the connection isn't closed and task doesn't complete even after some delay.
+        await Task.Delay(TimeSpan.FromSeconds(5));
+        Assert.False(resultTaskOne.IsCompleted);
+        Assert.False(resultTaskTwo.IsCompleted);
+
+        // Make an LSP document change that will trigger connection close.
+        var uri = testLspServer.GetCurrentSolution().Projects.First().Documents.First().GetURI();
+        await testLspServer.OpenDocumentAsync(uri);
+
+        // Assert that both tasks completes after a change occurs
+        var resultsOne = await resultTaskOne;
+        var resultsTwo = await resultTaskTwo;
+        Assert.NotEmpty(resultsOne);
+        Assert.Empty(resultsTwo);
+
+        // Make new requests - these requests should again wait for new changes.
+        resultTaskOne = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics, useProgress: true, category: PullDiagnosticCategories.WorkspaceDocumentsAndProject, triggerConnectionClose: false);
+        resultTaskTwo = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics, useProgress: true, category: PullDiagnosticCategories.EditAndContinue, triggerConnectionClose: false);
+
+        // Assert that the new requests correctly wait for new changes and do not complete even after some delay.
+        await Task.Delay(TimeSpan.FromSeconds(5));
+        Assert.False(resultTaskOne.IsCompleted);
+        Assert.False(resultTaskTwo.IsCompleted);
+
+        // Make an LSP document change that will trigger connection close.
+        await testLspServer.CloseDocumentAsync(uri);
+
+        // Assert that both tasks again complete after a change occurs
+        resultsOne = await resultTaskOne;
+        resultsTwo = await resultTaskTwo;
+        Assert.NotEmpty(resultsOne);
+        Assert.Empty(resultsTwo);
+    }
+
     [Theory, CombinatorialData]
     public async Task TestWorkspaceDiagnosticsForClosedFilesSwitchFSAFromOnToOff(bool useVSDiagnostics, bool mutatingLspWorkspace)
     {