Handle concurrent requests to update workspace contents and workspace SG info. (#79628)

Author: CyrusNajmabadi

src/EditorFeatures/Core/EditAndContinue/PdbMatchingSourceTextProvider.cs

@@ -29,7 +29,7 @@ internal sealed class PdbMatchingSourceTextProvider() : IEventListener, IPdbMatc
     private readonly object _guard = new();
 
     private bool _isActive;
-    private int _baselineSolutionVersion;
+    private int _baselineSolutionContentVersion;
     private readonly Dictionary<string, (DocumentState state, int solutionVersion)> _documentsWithChangedLoaderByPath = [];
     private WorkspaceEventRegistration? _workspaceChangedDisposer;
 
@@ -84,12 +84,12 @@ private void WorkspaceChanged(WorkspaceChangeEventArgs e)
         // The file checksum is no longer available from the latter, so capture it at this moment.
         if (oldDocument.State.TextAndVersionSource.CanReloadText && !newDocument.State.TextAndVersionSource.CanReloadText)
         {
-            var oldSolutionVersion = oldDocument.Project.Solution.WorkspaceVersion;
+            var oldSolutionVersion = oldDocument.Project.Solution.SolutionStateContentVersion;
 
             lock (_guard)
             {
                 // ignore updates to a document that we have already seen this session:
-                if (_isActive && oldSolutionVersion >= _baselineSolutionVersion && !_documentsWithChangedLoaderByPath.ContainsKey(oldDocument.FilePath))
+                if (_isActive && oldSolutionVersion >= _baselineSolutionContentVersion && !_documentsWithChangedLoaderByPath.ContainsKey(oldDocument.FilePath))
                 {
                     _documentsWithChangedLoaderByPath.Add(oldDocument.FilePath, (oldDocument.DocumentState, oldSolutionVersion));
                 }
@@ -104,7 +104,7 @@ public void SetBaseline(Solution solution)
     {
         lock (_guard)
         {
-            _baselineSolutionVersion = solution.WorkspaceVersion;
+            _baselineSolutionContentVersion = solution.SolutionStateContentVersion;
         }
     }
 

src/EditorFeatures/Test/EditAndContinue/EditAndContinueLanguageServiceTests.cs

@@ -395,15 +395,15 @@ public async Task DefaultPdbMatchingSourceTextProvider()
 
         var (key, (documentState, version)) = sourceTextProvider.GetTestAccessor().GetDocumentsWithChangedLoaderByPath().Single();
         Assert.Equal(sourceFile.Path, key);
-        Assert.Equal(solution.WorkspaceVersion, version);
-        Assert.Equal(source1, (await documentState.GetTextAsync(CancellationToken.None)).ToString());
+        Assert.Equal(solution.SolutionStateContentVersion, version);
+        Assert.Equal(source1, documentState.GetTextSynchronously(CancellationToken.None).ToString());
 
         // check committed document status:
         var debuggingSession = service.GetTestAccessor().GetActiveDebuggingSessions().Single();
         var (document, state) = await debuggingSession.LastCommittedSolution.GetDocumentAndStateAsync(document1, CancellationToken.None);
         var text = await document.GetTextAsync();
         Assert.Equal(CommittedSolution.DocumentState.MatchesBuildOutput, state);
-        Assert.Equal(source1, (await document.GetTextAsync(CancellationToken.None)).ToString());
+        Assert.Equal(source1, document.GetTextSynchronously(CancellationToken.None).ToString());
 
         await languageService.EndSessionAsync(CancellationToken.None);
     }

src/Features/Core/Portable/Completion/Providers/CompletionUtilities.cs

@@ -38,7 +38,7 @@ public static ImmutableArray<Project> GetDistinctProjectsFromLatestSolutionSnaps
         foreach (var project in projects)
         {
             projectIds.Add(project.Id);
-            if (solution is null || project.Solution.WorkspaceVersion > solution.WorkspaceVersion)
+            if (solution is null || project.Solution.SolutionStateContentVersion > solution.SolutionStateContentVersion)
             {
                 solution = project.Solution;
             }

src/LanguageServer/Protocol/Workspaces/LspWorkspaceManager.cs

@@ -13,7 +13,6 @@
 using Microsoft.CodeAnalysis.LanguageServer.Handler;
 using Microsoft.CodeAnalysis.LanguageServer.Handler.DocumentChanges;
 using Microsoft.CodeAnalysis.PooledObjects;
-using Microsoft.CodeAnalysis.Shared.Collections;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.CommonLanguageServerProtocol.Framework;
@@ -53,7 +52,7 @@ internal sealed class LspWorkspaceManager : IDocumentChangeTracker, ILspService
     /// workspace).
     /// <para/> Access to this is guaranteed to be serial by the <see cref="RequestExecutionQueue{RequestContextType}"/>
     /// </summary>
-    private readonly Dictionary<Workspace, (int? forkedFromVersion, Solution solution)> _cachedLspSolutions = [];
+    private readonly Dictionary<Workspace, (int? forkedFromVersion, Checksum? sourceGeneratorChecksum, Solution solution)> _cachedLspSolutions = [];
 
     /// <summary>
     /// Stores the current source text for each URI that is being tracked by LSP. Each time an LSP text sync
@@ -393,13 +392,20 @@ .. registeredWorkspaces.Where(workspace => workspace.Kind == WorkspaceKind.Misce
             if (doesAllTextMatch && doesAllSourceGeneratedTextMatch)
             {
                 // Remember that the current LSP text matches the text in this workspace solution.
-                _cachedLspSolutions[workspace] = (forkedFromVersion: null, workspaceCurrentSolution);
+                _cachedLspSolutions[workspace] = (forkedFromVersion: null, sourceGeneratorChecksum: null, workspaceCurrentSolution);
                 return (workspaceCurrentSolution, IsForked: false);
             }
 
+            var forkedFromVersion = workspaceCurrentSolution.SolutionStateContentVersion;
+            var sourceGeneratorChecksum = workspaceCurrentSolution.CompilationState.SourceGeneratorExecutionVersionMap.GetChecksum();
+
             // Step 4: See if we can reuse a previously forked solution.
-            if (cachedSolution != default && cachedSolution.forkedFromVersion == workspaceCurrentSolution.WorkspaceVersion)
+            if (cachedSolution != default &&
+                cachedSolution.forkedFromVersion == forkedFromVersion &&
+                cachedSolution.sourceGeneratorChecksum == sourceGeneratorChecksum)
+            {
                 return (cachedSolution.solution, IsForked: true);
+            }
 
             // Step 5: Fork a new solution from the workspace with the LSP text applied.
             var lspSolution = workspaceCurrentSolution;
@@ -417,7 +423,7 @@ .. registeredWorkspaces.Where(workspace => workspace.Kind == WorkspaceKind.Misce
             }
 
             // Remember this forked solution and the workspace version it was forked from.
-            _cachedLspSolutions[workspace] = (workspaceCurrentSolution.WorkspaceVersion, lspSolution);
+            _cachedLspSolutions[workspace] = (forkedFromVersion, sourceGeneratorChecksum, lspSolution);
             return (lspSolution, IsForked: true);
         }
 

src/LanguageServer/ProtocolUnitTests/Workspaces/LspWorkspaceManagerTests.cs

@@ -20,12 +20,9 @@
 
 namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests.Workspaces;
 
-public sealed class LspWorkspaceManagerTests : AbstractLanguageServerProtocolTests
+public sealed class LspWorkspaceManagerTests(ITestOutputHelper testOutputHelper)
+    : AbstractLanguageServerProtocolTests(testOutputHelper)
 {
-    public LspWorkspaceManagerTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper)
-    {
-    }
-
     [Theory, CombinatorialData]
     public async Task TestUsesLspTextOnOpenCloseAsync(bool mutatingLspWorkspace)
     {

src/LanguageServer/ProtocolUnitTests/Workspaces/SourceGeneratedDocumentTests.cs

@@ -2,6 +2,7 @@
 // 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.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -11,7 +12,6 @@
 using Microsoft.CodeAnalysis.Text;
 using Roslyn.Test.Utilities;
 using Roslyn.Test.Utilities.TestGenerators;
-using Roslyn.Utilities;
 using Xunit;
 using Xunit.Abstractions;
 using LSP = Roslyn.LanguageServer.Protocol;
@@ -223,6 +223,103 @@ internal async Task TestReturnsGeneratedSourceWhenManuallyRefreshed(bool mutatin
         Assert.Equal("// callCount: 1", secondRequest.Text);
     }
 
+    [Theory, CombinatorialData]
+    internal async Task TestCanRunSourceGeneratorAndApplyChangesConcurrently(
+        bool mutatingLspWorkspace,
+        bool majorVersionUpdate,
+        SourceGeneratorExecutionPreference sourceGeneratorExecution)
+    {
+        await using var testLspServer = await CreateTestLspServerAsync("""
+            class C
+            {
+            }
+            """, mutatingLspWorkspace);
+
+        var configService = testLspServer.TestWorkspace.ExportProvider.GetExportedValue<TestWorkspaceConfigurationService>();
+        configService.Options = new WorkspaceConfigurationOptions(SourceGeneratorExecution: sourceGeneratorExecution);
+
+        var callCount = 0;
+        var generatorReference = await AddGeneratorAsync(new CallbackGenerator(() => ("hintName.cs", "// callCount: " + callCount++)), testLspServer.TestWorkspace);
+
+        var sourceGeneratedDocuments = await testLspServer.GetCurrentSolution().Projects.Single().GetSourceGeneratedDocumentsAsync();
+        var sourceGeneratedDocumentIdentity = sourceGeneratedDocuments.Single().Identity;
+        var sourceGeneratorDocumentUri = SourceGeneratedDocumentUri.Create(sourceGeneratedDocumentIdentity);
+
+        var text = await testLspServer.ExecuteRequestAsync<SourceGeneratorGetTextParams, SourceGeneratedDocumentText>(SourceGeneratedDocumentGetTextHandler.MethodName,
+            new SourceGeneratorGetTextParams(new LSP.TextDocumentIdentifier { DocumentUri = sourceGeneratorDocumentUri }, ResultId: null), CancellationToken.None);
+
+        AssertEx.NotNull(text);
+        Assert.Equal("// callCount: 0", text.Text);
+
+        var initialSolution = testLspServer.GetCurrentSolution();
+        var initialExecutionMap = initialSolution.CompilationState.SourceGeneratorExecutionVersionMap.Map;
+
+        // Updating the execution version should trigger source generators to run in both automatic and balanced mode.
+        var forceRegeneration = majorVersionUpdate;
+        testLspServer.TestWorkspace.EnqueueUpdateSourceGeneratorVersion(projectId: null, forceRegeneration);
+        await testLspServer.WaitForSourceGeneratorsAsync();
+
+        var solutionWithChangedExecutionVersion = testLspServer.GetCurrentSolution();
+
+        var secondRequest = await testLspServer.ExecuteRequestAsync<SourceGeneratorGetTextParams, SourceGeneratedDocumentText>(SourceGeneratedDocumentGetTextHandler.MethodName,
+            new SourceGeneratorGetTextParams(new LSP.TextDocumentIdentifier { DocumentUri = sourceGeneratorDocumentUri }, ResultId: text.ResultId), CancellationToken.None);
+        AssertEx.NotNull(secondRequest);
+
+        if (forceRegeneration)
+        {
+            Assert.NotEqual(text.ResultId, secondRequest.ResultId);
+            Assert.Equal("// callCount: 1", secondRequest.Text);
+        }
+        else
+        {
+            Assert.Equal(text.ResultId, secondRequest.ResultId);
+            Assert.Null(secondRequest.Text);
+        }
+
+        var projectId1 = initialSolution.ProjectIds.Single();
+        var solutionWithDocumentChanged = initialSolution.WithDocumentText(
+            initialSolution.Projects.Single().Documents.Single().Id,
+            SourceText.From("class D { }"));
+
+        var expectVersionChange = sourceGeneratorExecution is SourceGeneratorExecutionPreference.Balanced || forceRegeneration;
+
+        // The content forked solution should have an SG execution version *less than* the one we just changed.
+        // Note: this will be patched up once we call TryApplyChanges.
+        if (expectVersionChange)
+        {
+            Assert.True(
+                solutionWithChangedExecutionVersion.CompilationState.SourceGeneratorExecutionVersionMap[projectId1]
+                > solutionWithDocumentChanged.CompilationState.SourceGeneratorExecutionVersionMap[projectId1]);
+        }
+        else
+        {
+            Assert.Equal(
+                solutionWithChangedExecutionVersion.CompilationState.SourceGeneratorExecutionVersionMap[projectId1],
+                solutionWithDocumentChanged.CompilationState.SourceGeneratorExecutionVersionMap[projectId1]);
+        }
+
+        Assert.True(testLspServer.TestWorkspace.TryApplyChanges(solutionWithDocumentChanged));
+
+        var finalSolution = testLspServer.GetCurrentSolution();
+
+        if (expectVersionChange)
+        {
+            // In balanced (or if we forced regen) mode, the execution version should have been updated to the new value.
+            Assert.NotEqual(initialExecutionMap[projectId1], solutionWithChangedExecutionVersion.CompilationState.SourceGeneratorExecutionVersionMap[projectId1]);
+            Assert.NotEqual(initialExecutionMap[projectId1], finalSolution.CompilationState.SourceGeneratorExecutionVersionMap[projectId1]);
+        }
+        else
+        {
+            // In automatic mode, nothing should change wrt to execution versions (unless we specified force-regenerate).
+            Assert.Equal(initialExecutionMap[projectId1], solutionWithChangedExecutionVersion.CompilationState.SourceGeneratorExecutionVersionMap[projectId1]);
+            Assert.Equal(initialExecutionMap[projectId1], finalSolution.CompilationState.SourceGeneratorExecutionVersionMap[projectId1]);
+        }
+
+        // The final execution version for the project should match the changed execution version, no matter what.
+        // Proving that the content change happened, but didn't drop the execution version change.
+        Assert.Equal(solutionWithChangedExecutionVersion.CompilationState.SourceGeneratorExecutionVersionMap[projectId1], finalSolution.CompilationState.SourceGeneratorExecutionVersionMap[projectId1]);
+    }
+
     [Theory, CombinatorialData]
     public async Task TestReturnsNullForRemovedClosedGeneratedFile(bool mutatingLspWorkspace)
     {
@@ -279,7 +376,9 @@ public async Task TestReturnsNullForRemovedOpenedGeneratedFile(bool mutatingLspW
         Assert.Null(secondRequest.Text);
     }
 
-    private async Task<TestLspServer> CreateTestLspServerWithGeneratorAsync(bool mutatingLspWorkspace, string generatedDocumentText)
+    private async Task<TestLspServer> CreateTestLspServerWithGeneratorAsync(
+        bool mutatingLspWorkspace,
+        [StringSyntax(PredefinedEmbeddedLanguageNames.CSharpTest)] string generatedDocumentText)
     {
         var testLspServer = await CreateTestLspServerAsync(string.Empty, mutatingLspWorkspace);
         await AddGeneratorAsync(new SingleFileTestGenerator(generatedDocumentText), testLspServer.TestWorkspace);

src/Tools/ExternalAccess/Razor/Features/SolutionExtensions.cs

@@ -16,5 +16,5 @@ public static ImmutableArray<DocumentId> GetDocumentIds(this Solution solution,
         => LanguageServer.Extensions.GetDocumentIds(solution, new(documentUri));
 
     public static int GetWorkspaceVersion(this Solution solution)
-        => solution.WorkspaceVersion;
+        => solution.SolutionStateContentVersion;
 }

src/VisualStudio/IntegrationTest/New.IntegrationTests/InProcess/ITextViewWindowVerifierInProcessExtensions.cs

@@ -115,11 +115,11 @@ await textViewWindowVerifier.TestServices.Workspace.WaitForAllAsyncOperationsAsy
                     ],
                     cancellationToken);
 
-                if (codeActionLogger.Messages.Any())
+                if (codeActionLogger.Messages.Count != 0)
                 {
                     foreach (var e in events)
                     {
-                        codeActionLogger.Messages.Add($"{e.OldSolution.WorkspaceVersion} to {e.NewSolution.WorkspaceVersion}: {e.Kind} {e.DocumentId}");
+                        codeActionLogger.Messages.Add($"{e.OldSolution.SolutionStateContentVersion} to {e.NewSolution.SolutionStateContentVersion}: {e.Kind} {e.DocumentId}");
                     }
                 }
 

src/Workspaces/Core/Portable/CodeActions/Operations/ApplyChangesOperation.cs

@@ -53,7 +53,7 @@ internal static bool ApplyOrMergeChanges(
         var currentSolution = workspace.CurrentSolution;
 
         // if there was no intermediary edit, just apply the change fully.
-        if (changedSolution.WorkspaceVersion == currentSolution.WorkspaceVersion)
+        if (changedSolution.SolutionStateContentVersion == currentSolution.SolutionStateContentVersion)
         {
             var result = workspace.TryApplyChanges(changedSolution, progressTracker);
 

src/Workspaces/Core/Portable/ExternalAccess/UnitTesting/Api/UnitTestingSolutionExtensions.cs

@@ -10,7 +10,7 @@ namespace Microsoft.CodeAnalysis.ExternalAccess.UnitTesting.Api;
 internal static class UnitTestingSolutionExtensions
 {
     public static int GetWorkspaceVersion(this Solution solution)
-        => solution.WorkspaceVersion;
+        => solution.SolutionStateContentVersion;
 
     public static async Task<UnitTestingChecksumWrapper> GetChecksumAsync(this Solution solution, CancellationToken cancellationToken)
         => new UnitTestingChecksumWrapper(await solution.CompilationState.GetChecksumAsync(cancellationToken).ConfigureAwait(false));