Fix workspace pull diagnostics

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

@@ -113,6 +113,15 @@ protected abstract ValueTask<ImmutableArray<IDiagnosticSource>> GetOrderedDiagno
 
         protected abstract string? GetDiagnosticCategory(TDiagnosticsParams diagnosticsParams);
 
+        /// <summary>
+        /// 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)
+        {
+            return Task.CompletedTask;
+        }
+
         public async Task<TReturn?> HandleRequestAsync(
             TDiagnosticsParams diagnosticsParams, RequestContext context, CancellationToken cancellationToken)
         {
@@ -182,6 +191,11 @@ 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).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.
             context.TraceInformation($"{this.GetType()} finished getting diagnostics");

src/Features/LanguageServer/Protocol/Handler/Diagnostics/Public/PublicWorkspacePullDiagnosticHandlerFactory.cs

@@ -8,29 +8,34 @@
 using Microsoft.CodeAnalysis.EditAndContinue;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Options;
-using Microsoft.VisualStudio.LanguageServer.Protocol;
 
 namespace Microsoft.CodeAnalysis.LanguageServer.Handler.Diagnostics.Public;
 
 [ExportCSharpVisualBasicLspServiceFactory(typeof(PublicWorkspacePullDiagnosticsHandler)), Shared]
 internal class PublicWorkspacePullDiagnosticHandlerFactory : ILspServiceFactory
 {
+    private readonly LspWorkspaceRegistrationService _registrationService;
     private readonly IDiagnosticAnalyzerService _analyzerService;
     private readonly EditAndContinueDiagnosticUpdateSource _editAndContinueDiagnosticUpdateSource;
     private readonly IGlobalOptionService _globalOptions;
 
     [ImportingConstructor]
     [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
     public PublicWorkspacePullDiagnosticHandlerFactory(
+        LspWorkspaceRegistrationService registrationService,
         IDiagnosticAnalyzerService analyzerService,
         EditAndContinueDiagnosticUpdateSource editAndContinueDiagnosticUpdateSource,
         IGlobalOptionService globalOptions)
     {
+        _registrationService = registrationService;
         _analyzerService = analyzerService;
         _editAndContinueDiagnosticUpdateSource = editAndContinueDiagnosticUpdateSource;
         _globalOptions = globalOptions;
     }
 
     public ILspService CreateILspService(LspServices lspServices, WellKnownLspServerKinds serverKind)
-        => new PublicWorkspacePullDiagnosticsHandler(_analyzerService, _editAndContinueDiagnosticUpdateSource, _globalOptions);
+    {
+        var workspaceManager = lspServices.GetRequiredService<LspWorkspaceManager>();
+        return new PublicWorkspacePullDiagnosticsHandler(workspaceManager, _registrationService, _analyzerService, _editAndContinueDiagnosticUpdateSource, _globalOptions);
+    }
 }

src/Features/LanguageServer/Protocol/Handler/Diagnostics/Public/PublicWorkspacePullDiagnosticsHandler.cs

@@ -12,6 +12,7 @@
 using Microsoft.CodeAnalysis.EditAndContinue;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.LanguageServer.Handler.Diagnostics.Public;
 
@@ -21,14 +22,32 @@ namespace Microsoft.CodeAnalysis.LanguageServer.Handler.Diagnostics.Public;
 using WorkspaceDiagnosticPartialReport = SumType<WorkspaceDiagnosticReport, WorkspaceDiagnosticReportPartialResult>;
 
 [Method(Methods.WorkspaceDiagnosticName)]
-internal class PublicWorkspacePullDiagnosticsHandler : AbstractPullDiagnosticHandler<WorkspaceDiagnosticParams, WorkspaceDiagnosticPartialReport, WorkspaceDiagnosticReport?>
+internal class PublicWorkspacePullDiagnosticsHandler : AbstractPullDiagnosticHandler<WorkspaceDiagnosticParams, WorkspaceDiagnosticPartialReport, WorkspaceDiagnosticReport?>, IDisposable
 {
+    private readonly LspWorkspaceRegistrationService _workspaceRegistrationService;
+    private readonly LspWorkspaceManager _workspaceManager;
+
+    /// <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.
+    /// </summary>
+    private int _lspChanged = 0;
+
     public PublicWorkspacePullDiagnosticsHandler(
+        LspWorkspaceManager workspaceManager,
+        LspWorkspaceRegistrationService registrationService,
         IDiagnosticAnalyzerService analyzerService,
         EditAndContinueDiagnosticUpdateSource editAndContinueDiagnosticUpdateSource,
         IGlobalOptionService globalOptions)
         : base(analyzerService, editAndContinueDiagnosticUpdateSource, globalOptions)
     {
+        _workspaceRegistrationService = registrationService;
+        _workspaceRegistrationService.LspSolutionChanged += OnLspSolutionChanged;
+
+        _workspaceManager = workspaceManager;
+        _workspaceManager.LspTextChanged += OnLspTextChanged;
     }
 
     /// <summary>
@@ -118,4 +137,41 @@ protected override ValueTask<ImmutableArray<IDiagnosticSource>> GetOrderedDiagno
             }
         }).ToImmutableArray();
     }
+
+    private void OnLspSolutionChanged(object? sender, WorkspaceChangeEventArgs e)
+    {
+        UpdateLspChanged();
+    }
+
+    private void OnLspTextChanged(object? sender, EventArgs e)
+    {
+        UpdateLspChanged();
+    }
+
+    private void UpdateLspChanged()
+    {
+        Interlocked.Exchange(ref _lspChanged, 1);
+    }
+
+    protected override async Task WaitForChangesAsync(RequestContext context)
+    {
+        // 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)
+        {
+            // 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)).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.
+        return;
+    }
+
+    public void Dispose()
+    {
+        _workspaceManager.LspTextChanged -= OnLspTextChanged;
+        _workspaceRegistrationService.LspSolutionChanged -= OnLspSolutionChanged;
+    }
 }

src/Features/LanguageServer/Protocol/Workspaces/LspWorkspaceManager.cs

@@ -87,6 +87,8 @@ public LspWorkspaceManager(
         _lspWorkspaceRegistrationService = lspWorkspaceRegistrationService;
     }
 
+    public EventHandler<EventArgs>? LspTextChanged;
+
     #region Implementation of IDocumentChangeTracker
 
     /// <summary>
@@ -102,6 +104,8 @@ public void StartTracking(Uri uri, SourceText documentText)
 
         // If LSP changed, we need to compare against the workspace again to get the updated solution.
         _cachedLspSolutions.Clear();
+
+        LspTextChanged?.Invoke(this, EventArgs.Empty);
     }
 
     /// <summary>
@@ -120,6 +124,8 @@ public void StopTracking(Uri uri)
 
         // Also remove it from our loose files workspace if it is still there.
         _lspMiscellaneousFilesWorkspace.TryRemoveMiscellaneousDocument(uri);
+
+        LspTextChanged?.Invoke(this, EventArgs.Empty);
     }
 
     /// <summary>
@@ -135,6 +141,8 @@ public void UpdateTrackedDocument(Uri uri, SourceText newSourceText)
 
         // If LSP changed, we need to compare against the workspace again to get the updated solution.
         _cachedLspSolutions.Clear();
+
+        LspTextChanged?.Invoke(this, EventArgs.Empty);
     }
 
     public ImmutableDictionary<Uri, SourceText> GetTrackedLspText() => _trackedDocuments;

src/Features/LanguageServer/ProtocolUnitTests/Diagnostics/PullDiagnosticTests.cs

@@ -1658,6 +1658,52 @@ public async Task TestWorkspaceDiagnosticsDoesNotThrowIfProjectWithoutFilePathEx
             Assert.Equal(@"C:/C2.cs", results[2].TextDocument.Uri.AbsolutePath);
         }
 
+        [Fact]
+        public async Task TestPublicWorkspaceDiagnosticsWaitsForLspTextChanges()
+        {
+            var markup1 =
+@"class A {";
+            var markup2 = "";
+            await using var testLspServer = await CreateTestWorkspaceWithDiagnosticsAsync(
+                new[] { markup1, markup2 }, BackgroundAnalysisScope.FullSolution, useVSDiagnostics: false);
+
+            var resultTask = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics: false, useProgress: true);
+
+            // Assert that the connection isn't closed and task doesn't complete even after some delay.
+            await Task.Delay(TimeSpan.FromSeconds(5));
+            Assert.False(resultTask.IsCompleted);
+
+            var uri = testLspServer.GetCurrentSolution().Projects.First().Documents.First().GetURI();
+            await testLspServer.OpenDocumentAsync(uri);
+
+            // Assert the task completes after a change occurs
+            var results = await resultTask;
+            Assert.NotEmpty(results);
+        }
+
+        [Fact]
+        public async Task TestPublicWorkspaceDiagnosticsWaitsForLspSolutionChanges()
+        {
+            var markup1 =
+@"class A {";
+            var markup2 = "";
+            await using var testLspServer = await CreateTestWorkspaceWithDiagnosticsAsync(
+                new[] { markup1, markup2 }, BackgroundAnalysisScope.FullSolution, useVSDiagnostics: false);
+
+            var resultTask = RunGetWorkspacePullDiagnosticsAsync(testLspServer, useVSDiagnostics: false, useProgress: true);
+
+            // Assert that the connection isn't closed and task doesn't complete even after some delay.
+            await Task.Delay(TimeSpan.FromSeconds(5));
+            Assert.False(resultTask.IsCompleted);
+
+            var projectInfo = testLspServer.TestWorkspace.Projects.Single().ToProjectInfo();
+            testLspServer.TestWorkspace.OnProjectReloaded(projectInfo);
+
+            // Assert the task completes after a change occurs
+            var results = await resultTask;
+            Assert.NotEmpty(results);
+        }
+
         #endregion
     }
 }