Add server side support for refreshing source generated files

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

@@ -24,7 +24,7 @@ internal record struct DiagnosticsRequestState(Project Project, int GlobalStateV
     /// and works well for us in the normal case.  The latter still allows us to reuse diagnostics when changes happen that
     /// update the version stamp but not the content (for example, forking LSP text).
     /// </summary>
-    private sealed class DiagnosticsPullCache(string uniqueKey) : VersionedPullCache<(int globalStateVersion, VersionStamp? dependentVersion), (int globalStateVersion, Checksum dependentChecksum), DiagnosticsRequestState, DiagnosticData>(uniqueKey)
+    private sealed class DiagnosticsPullCache(string uniqueKey) : VersionedPullCache<(int globalStateVersion, VersionStamp? dependentVersion), (int globalStateVersion, Checksum dependentChecksum), DiagnosticsRequestState, ImmutableArray<DiagnosticData>>(uniqueKey)
     {
         public override async Task<(int globalStateVersion, VersionStamp? dependentVersion)> ComputeCheapVersionAsync(DiagnosticsRequestState state, CancellationToken cancellationToken)
         {

src/LanguageServer/Protocol/Handler/PullHandlers/VersionedPullCache.CacheItem.cs

@@ -52,7 +52,7 @@ private sealed class CacheItem(string uniqueKey)
         /// 
         /// Returns <see langword="null"/> if the previousPullResult can be re-used, otherwise returns a new resultId and the new data associated with it.
         /// </summary>
-        public async Task<(string, ImmutableArray<TComputedData>)?> UpdateCacheItemAsync(
+        public async Task<(string, TComputedData)?> UpdateCacheItemAsync(
             VersionedPullCache<TCheapVersion, TExpensiveVersion, TState, TComputedData> cache,
             PreviousPullResult? previousPullResult,
             bool isFullyLoaded,

src/LanguageServer/Protocol/Handler/PullHandlers/VersionedPullCache.cs

@@ -4,7 +4,6 @@
 
 using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Host;
@@ -19,7 +18,6 @@ namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
 /// with different computation costs to determine if the previous cached data is still valid.
 /// </summary>
 internal abstract partial class VersionedPullCache<TCheapVersion, TExpensiveVersion, TState, TComputedData>(string uniqueKey)
-    where TComputedData : notnull
 {
     /// <summary>
     /// Map of workspace and diagnostic source to the data used to make the last pull report.
@@ -59,9 +57,9 @@ internal abstract partial class VersionedPullCache<TCheapVersion, TExpensiveVers
     /// 
     /// Note - this will run under the semaphore in <see cref="CacheItem._gate"/>.
     /// </summary>
-    public abstract Task<ImmutableArray<TComputedData>> ComputeDataAsync(TState state, CancellationToken cancellationToken);
+    public abstract Task<TComputedData> ComputeDataAsync(TState state, CancellationToken cancellationToken);
 
-    public abstract Checksum ComputeChecksum(ImmutableArray<TComputedData> data);
+    public abstract Checksum ComputeChecksum(TComputedData data);
 
     /// <summary>
     /// If results have changed since the last request this calculates and returns a new
@@ -70,7 +68,7 @@ internal abstract partial class VersionedPullCache<TCheapVersion, TExpensiveVers
     /// <param name="idToClientLastResult">a map of roslyn document or project id to the previous result the client sent us for that doc.</param>
     /// <param name="projectOrDocumentId">the id of the project or document that we are checking to see if it has changed.</param>
     /// <returns>Null when results are unchanged, otherwise returns a non-null new resultId.</returns>
-    public async Task<(string ResultId, ImmutableArray<TComputedData> Data)?> GetOrComputeNewDataAsync(
+    public async Task<(string ResultId, TComputedData Data)?> GetOrComputeNewDataAsync(
         Dictionary<ProjectOrDocumentId, PreviousPullResult> idToClientLastResult,
         ProjectOrDocumentId projectOrDocumentId,
         Project project,

src/LanguageServer/Protocol/Handler/SourceGenerators/SourceGeneratedDocumentCache.cs

@@ -0,0 +1,60 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Composition;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.Text;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis.LanguageServer.Handler.SourceGenerators;
+
+internal record struct SourceGeneratedDocumentGetTextState(Document Document);
+
+internal class SourceGeneratedDocumentCache(string uniqueKey) : VersionedPullCache<(SourceGeneratorExecutionVersion, VersionStamp), object?, SourceGeneratedDocumentGetTextState, SourceText?>(uniqueKey), ILspService
+{
+    public override async Task<(SourceGeneratorExecutionVersion, VersionStamp)> ComputeCheapVersionAsync(SourceGeneratedDocumentGetTextState state, CancellationToken cancellationToken)
+    {
+        // The execution version and the dependent version must be considered as one version cached together -
+        // it is not correct to say that if the execution version is the same then we can re-use results (as in automatic mode the execution version never changes).
+        var executionVersion = state.Document.Project.Solution.GetSourceGeneratorExecutionVersion(state.Document.Project.Id);
+        var dependentVersion = await state.Document.Project.GetDependentVersionAsync(cancellationToken).ConfigureAwait(false);
+        return (executionVersion, dependentVersion);
+    }
+
+    public override Task<object?> ComputeExpensiveVersionAsync(SourceGeneratedDocumentGetTextState state, CancellationToken cancellationToken)
+    {
+        return SpecializedTasks.Null<object>();
+    }
+
+    public override Checksum ComputeChecksum(SourceText? data)
+    {
+        return data is null ? Checksum.Null : Checksum.From(data.GetChecksum());
+    }
+
+    public override async Task<SourceText?> ComputeDataAsync(SourceGeneratedDocumentGetTextState state, CancellationToken cancellationToken)
+    {
+        // When a user has a open source-generated file, we ensure that the contents in the LSP snapshot match the contents that we
+        // get through didOpen/didChanges, like any other file. That way operations in LSP file are in sync with the
+        // contents the user has. However in this case, we don't want to look at that frozen text, but look at what the
+        // generator would generate if we ran it again. Otherwise, we'll get "stuck" and never update the file with something new.
+        var unfrozenDocument = await state.Document.Project.Solution.WithoutFrozenSourceGeneratedDocuments().GetDocumentAsync(state.Document.Id, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+        return unfrozenDocument == null
+            ? null
+            : await unfrozenDocument.GetTextAsync(cancellationToken).ConfigureAwait(false);
+    }
+}
+
+[ExportCSharpVisualBasicLspServiceFactory(typeof(SourceGeneratedDocumentCache)), Shared]
+[method: ImportingConstructor]
+[method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+internal sealed class SourceGeneratedDocumentCacheFactory() : ILspServiceFactory
+{
+    public ILspService CreateILspService(LspServices lspServices, WellKnownLspServerKinds serverKind)
+    {
+        return new SourceGeneratedDocumentCache(this.GetType().Name);
+    }
+}

src/LanguageServer/Protocol/Handler/SourceGenerators/SourceGeneratedDocumentGetTextHandler.cs

@@ -3,14 +3,16 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Generic;
 using System.Composition;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.LanguageServer.Handler.Diagnostics;
 using Roslyn.Utilities;
 using LSP = Roslyn.LanguageServer.Protocol;
 
-namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
+namespace Microsoft.CodeAnalysis.LanguageServer.Handler.SourceGenerators;
 
 [ExportCSharpVisualBasicStatelessLspService(typeof(SourceGeneratedDocumentGetTextHandler)), Shared]
 [Method(MethodName)]
@@ -29,18 +31,45 @@ public async Task<SourceGeneratedDocumentText> HandleRequestAsync(SourceGenerato
     {
         var document = context.Document;
 
+        if (document is null)
+        {
+            // The source generated file being asked about is not present.
+            // This is a rare case the request queue always gives us a frozen, non-null document for any opened sg document,
+            // even if the generator itself was removed and the document no longer exists in the host solution.
+            //
+            // We can only get a null document here if the sg document has not been opened and
+            // the source generated document does not exist in the workspace.
+            //
+            // Return a value indicating that the document is removed.
+            return new SourceGeneratedDocumentText(ResultId: null, Text: null);
+        }
+
         // Nothing here strictly prevents this from working on any other document, but we'll assert we got a source-generated file, since
         // it wouldn't really make sense for the server to be asked for the contents of a regular file. Since this endpoint is intended for
         // source-generated files only, this would indicate that something else has gone wrong.
         Contract.ThrowIfFalse(document is SourceGeneratedDocument);
 
-        // When a user has a open source-generated file, we ensure that the contents in the LSP snapshot match the contents that we
-        // get through didOpen/didChanges, like any other file. That way operations in LSP file are in sync with the
-        // contents the user has. However in this case, we don't want to look at that frozen text, but look at what the
-        // generator would generate if we ran it again. Otherwise, we'll get "stuck" and never update the file with something new.
-        document = await document.Project.Solution.WithoutFrozenSourceGeneratedDocuments().GetDocumentAsync(document.Id, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+        var cache = context.GetRequiredLspService<SourceGeneratedDocumentCache>();
+        var projectOrDocument = new ProjectOrDocumentId(document.Id);
+
+        var previousPullResults = new Dictionary<ProjectOrDocumentId, PreviousPullResult>();
+        if (request.ResultId is not null)
+        {
+            previousPullResults.Add(projectOrDocument, new PreviousPullResult(request.ResultId, request.TextDocument));
+        }
+
+        var newResult = await cache.GetOrComputeNewDataAsync(previousPullResults, projectOrDocument, document.Project, new SourceGeneratedDocumentGetTextState(document), cancellationToken).ConfigureAwait(false);
 
-        var text = document != null ? await document.GetTextAsync(cancellationToken).ConfigureAwait(false) : null;
-        return new SourceGeneratedDocumentText(text?.ToString());
+        if (newResult is null)
+        {
+            Contract.ThrowIfNull(request.ResultId, "Attempted to reuse cache entry but given no resultId");
+            // The generated document is the same, we can return the same resultId.
+            return new SourceGeneratedDocumentText(request.ResultId, null);
+        }
+        else
+        {
+            var data = newResult.Value.Data?.ToString();
+            return new SourceGeneratedDocumentText(newResult.Value.ResultId, data);
+        }
     }
 }

src/LanguageServer/Protocol/Handler/SourceGenerators/SourceGeneratedDocumentText.cs

@@ -1,9 +1,11 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
 using System.Text.Json.Serialization;
 
 namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
 
-internal sealed record SourceGeneratedDocumentText([property: JsonPropertyName("text")] string? Text);
+internal sealed record SourceGeneratedDocumentText(
+    [property: JsonPropertyName("resultId")] string? ResultId,
+    [property: JsonPropertyName("text")] string? Text);

src/LanguageServer/Protocol/Handler/SourceGenerators/SourceGeneratorGetTextParams.cs

@@ -1,10 +1,12 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
 using System.Text.Json.Serialization;
 using Roslyn.LanguageServer.Protocol;
 
-namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
+namespace Microsoft.CodeAnalysis.LanguageServer.Handler.SourceGenerators;
 
-internal sealed record SourceGeneratorGetTextParams([property: JsonPropertyName("textDocument")] TextDocumentIdentifier TextDocument) : ITextDocumentParams;
+internal sealed record SourceGeneratorGetTextParams(
+    [property: JsonPropertyName("textDocument")] TextDocumentIdentifier TextDocument,
+    [property: JsonPropertyName("resultId")] string? ResultId) : ITextDocumentParams;

src/LanguageServer/Protocol/Handler/SourceGenerators/SourceGeneratorRefreshQueue.cs

@@ -0,0 +1,140 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Shared.TestHooks;
+using Roslyn.LanguageServer.Protocol;
+using Roslyn.Utilities;
+using StreamJsonRpc;
+
+namespace Microsoft.CodeAnalysis.LanguageServer.Handler.SourceGenerators;
+
+internal class SourceGeneratorRefreshQueue :
+        IOnInitialized,
+        ILspService,
+        IDisposable
+{
+    private const string RefreshSourceGeneratedDocumentName = "workspace/refreshSourceGeneratedDocument";
+
+    private readonly IAsynchronousOperationListener _asyncListener;
+    private readonly CancellationTokenSource _disposalTokenSource = new();
+    private readonly LspWorkspaceRegistrationService _lspWorkspaceRegistrationService;
+    private readonly LspWorkspaceManager _lspWorkspaceManager;
+    private readonly IClientLanguageServerManager _notificationManager;
+    private readonly AsyncBatchingWorkQueue _refreshQueue;
+
+    public SourceGeneratorRefreshQueue(
+                IAsynchronousOperationListenerProvider asynchronousOperationListenerProvider,
+                LspWorkspaceRegistrationService lspWorkspaceRegistrationService,
+                LspWorkspaceManager lspWorkspaceManager,
+                IClientLanguageServerManager notificationManager)
+    {
+        _lspWorkspaceRegistrationService = lspWorkspaceRegistrationService;
+        _lspWorkspaceManager = lspWorkspaceManager;
+        _notificationManager = notificationManager;
+        _asyncListener = asynchronousOperationListenerProvider.GetListener(FeatureAttribute.SourceGenerators);
+
+        // Batch up workspace notifications so that we only send a notification to refresh source generated files
+        // every 2 seconds - long enough to avoid spamming the client with notifications, but short enough to refresh
+        // the source generated files relatively frequently.
+        _refreshQueue = _refreshQueue = new AsyncBatchingWorkQueue(
+            delay: TimeSpan.FromMilliseconds(2000),
+            processBatchAsync: RefreshSourceGeneratedDocumentsAsync,
+            asyncListener: _asyncListener,
+            _disposalTokenSource.Token);
+    }
+
+    public Task OnInitializedAsync(ClientCapabilities clientCapabilities, RequestContext context, CancellationToken cancellationToken)
+    {
+        if (clientCapabilities.HasVisualStudioLspCapability())
+        {
+            // VS source generated document content is not provided by LSP.
+            return Task.CompletedTask;
+        }
+
+        // After we have initialized we can start listening for workspace changes.
+        _lspWorkspaceRegistrationService.LspSolutionChanged += OnLspSolutionChanged;
+        return Task.CompletedTask;
+    }
+
+    private void OnLspSolutionChanged(object? sender, WorkspaceChangeEventArgs e)
+    {
+        var projectId = e.ProjectId ?? e.DocumentId?.ProjectId;
+        if (projectId is not null)
+        {
+            // We have a specific changed project - do some additional checks to see if
+            // source generators possibly changed.  Note that this overreports actual
+            // changes to the source generated text; we rely on resultIds in the text retrieval to avoid unnecessary serialization.
+
+            // Trivial check.  see if the SG version of these projects changed.  If so, we definitely want to update
+            // this generated file.
+            if (e.OldSolution.GetSourceGeneratorExecutionVersion(projectId) !=
+                e.NewSolution.GetSourceGeneratorExecutionVersion(projectId))
+            {
+                _refreshQueue.AddWork();
+                return;
+            }
+
+            var oldProject = e.OldSolution.GetProject(projectId);
+            var newProject = e.NewSolution.GetProject(projectId);
+
+            if (oldProject != null && newProject != null)
+            {
+                var asyncToken = _asyncListener.BeginAsyncOperation($"{nameof(SourceGeneratorRefreshQueue)}.{nameof(OnLspSolutionChanged)}");
+                CheckDependentVersionsAsync(oldProject, newProject, _disposalTokenSource.Token).CompletesAsyncOperation(asyncToken);
+            }
+        }
+        else
+        {
+            // We don't have a specific project change - if this is a solution change we need to queue a refresh anyway.
+            if (e.Kind is WorkspaceChangeKind.SolutionChanged or WorkspaceChangeKind.SolutionAdded or WorkspaceChangeKind.SolutionRemoved or WorkspaceChangeKind.SolutionReloaded or WorkspaceChangeKind.SolutionCleared)
+            {
+                _refreshQueue.AddWork();
+            }
+        }
+
+        async Task CheckDependentVersionsAsync(Project oldProject, Project newProject, CancellationToken cancellationToken)
+        {
+            if (await oldProject.GetDependentVersionAsync(cancellationToken).ConfigureAwait(false) !=
+                await newProject.GetDependentVersionAsync(cancellationToken).ConfigureAwait(false))
+            {
+                _refreshQueue.AddWork();
+            }
+        }
+    }
+
+    private ValueTask RefreshSourceGeneratedDocumentsAsync(
+        CancellationToken cancellationToken)
+    {
+        var hasOpenSourceGeneratedDocuments = _lspWorkspaceManager.GetTrackedLspText().Keys.Any(uri => uri.Scheme == SourceGeneratedDocumentUri.Scheme);
+        if (!hasOpenSourceGeneratedDocuments)
+        {
+            // There are no opened source generated documents - we don't need to bother asking the client to refresh anything.
+            return ValueTaskFactory.CompletedTask;
+        }
+
+        try
+        {
+            return _notificationManager.SendNotificationAsync(RefreshSourceGeneratedDocumentName, cancellationToken);
+        }
+        catch (Exception ex) when (ex is ObjectDisposedException or ConnectionLostException)
+        {
+            // It is entirely possible that we're shutting down and the connection is lost while we're trying to send a notification
+            // as this runs outside of the guaranteed ordering in the queue. We can safely ignore this exception.
+        }
+
+        return ValueTaskFactory.CompletedTask;
+    }
+
+    public void Dispose()
+    {
+        _lspWorkspaceRegistrationService.LspSolutionChanged -= OnLspSolutionChanged;
+        _disposalTokenSource.Cancel();
+        _disposalTokenSource.Dispose();
+    }
+}