Always sync text changes and active doc info to remote workspace

src/EditorFeatures/Core/Remote/SolutionChecksumUpdater.cs

@@ -3,13 +3,16 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Immutable;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Collections;
 using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.CodeAnalysis.Notification;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
+using Microsoft.CodeAnalysis.Text;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Remote;
@@ -48,7 +51,7 @@ internal sealed class SolutionChecksumUpdater
     private readonly AsyncBatchingWorkQueue _synchronizeActiveDocumentQueue;
 
     private readonly object _gate = new();
-    private bool _isPaused;
+    private bool _isSynchronizeWorkspacePaused;
 
     public SolutionChecksumUpdater(
         Workspace workspace,
@@ -62,20 +65,22 @@ public SolutionChecksumUpdater(
         _workspace = workspace;
         _documentTrackingService = workspace.Services.GetRequiredService<IDocumentTrackingService>();
 
-        _textChangeQueue = new AsyncBatchingWorkQueue<(Document oldDocument, Document newDocument)>(
+        _synchronizeWorkspaceQueue = new AsyncBatchingWorkQueue(
             DelayTimeSpan.NearImmediate,
-            SynchronizeTextChangesAsync,
+            SynchronizePrimaryWorkspaceAsync,
             listener,
             shutdownToken);
 
-        _synchronizeWorkspaceQueue = new AsyncBatchingWorkQueue(
-            DelayTimeSpan.NearImmediate,
-            SynchronizePrimaryWorkspaceAsync,
+        // Text changes and active doc info are tiny messages.  So attempt to send them immediately.  Just batching
+        // things up if we get a flurry of notifications.
+        _textChangeQueue = new AsyncBatchingWorkQueue<(Document oldDocument, Document newDocument)>(
+            TimeSpan.Zero,
+            SynchronizeTextChangesAsync,
             listener,
             shutdownToken);
 
         _synchronizeActiveDocumentQueue = new AsyncBatchingWorkQueue(
-            DelayTimeSpan.NearImmediate,
+            TimeSpan.Zero,
             SynchronizeActiveDocumentAsync,
             listener,
             shutdownToken);
@@ -90,14 +95,15 @@ public SolutionChecksumUpdater(
             _globalOperationService.Stopped += OnGlobalOperationStopped;
         }
 
-        // Enqueue the work to sync the initial solution.
-        ResumeWork();
+        // Enqueue the work to sync the initial data over.
+        _synchronizeActiveDocumentQueue.AddWork();
+        _synchronizeWorkspaceQueue.AddWork();
     }
 
     public void Shutdown()
     {
         // Try to stop any work that is in progress.
-        PauseWork();
+        PauseSynchronizingPrimaryWorkspace();
 
         _documentTrackingService.ActiveDocumentChanged -= OnActiveDocumentChanged;
         _workspace.WorkspaceChanged -= OnWorkspaceChanged;
@@ -110,43 +116,33 @@ public void Shutdown()
     }
 
     private void OnGlobalOperationStarted(object? sender, EventArgs e)
-        => PauseWork();
+        => PauseSynchronizingPrimaryWorkspace();
 
     private void OnGlobalOperationStopped(object? sender, EventArgs e)
-        => ResumeWork();
+        => ResumeSynchronizingPrimaryWorkspace();
 
-    private void PauseWork()
+    private void PauseSynchronizingPrimaryWorkspace()
     {
         // An expensive global operation started (like a build).  Pause ourselves and cancel any outstanding work in
         // progress to synchronize the solution.
         lock (_gate)
         {
             _synchronizeWorkspaceQueue.CancelExistingWork();
-            _synchronizeActiveDocumentQueue.CancelExistingWork();
-            _isPaused = true;
+            _isSynchronizeWorkspacePaused = true;
         }
     }
 
-    private void ResumeWork()
+    private void ResumeSynchronizingPrimaryWorkspace()
     {
         lock (_gate)
         {
-            _isPaused = false;
-            _synchronizeActiveDocumentQueue.AddWork();
+            _isSynchronizeWorkspacePaused = false;
             _synchronizeWorkspaceQueue.AddWork();
         }
     }
 
     private void OnWorkspaceChanged(object? sender, WorkspaceChangeEventArgs e)
     {
-        // Check if we're currently paused.  If so ignore this notification.  We don't want to any work in response
-        // to whatever the workspace is doing.
-        lock (_gate)
-        {
-            if (_isPaused)
-                return;
-        }
-
         if (e.Kind == WorkspaceChangeKind.DocumentChanged)
         {
             var oldDocument = e.OldSolution.GetDocument(e.DocumentId);
@@ -155,7 +151,13 @@ private void OnWorkspaceChanged(object? sender, WorkspaceChangeEventArgs e)
                 _textChangeQueue.AddWork((oldDocument, newDocument));
         }
 
-        _synchronizeWorkspaceQueue.AddWork();
+        // Check if we're currently paused.  If so ignore this notification.  We don't want to any work in response
+        // to whatever the workspace is doing.
+        lock (_gate)
+        {
+            if (!_isSynchronizeWorkspacePaused)
+                _synchronizeWorkspaceQueue.AddWork();
+        }
     }
 
     private void OnActiveDocumentChanged(object? sender, DocumentId? e)
@@ -204,55 +206,52 @@ private async ValueTask SynchronizeTextChangesAsync(
         ImmutableSegmentedList<(Document oldDocument, Document newDocument)> values,
         CancellationToken cancellationToken)
     {
-        foreach (var (oldDocument, newDocument) in values)
-        {
-            cancellationToken.ThrowIfCancellationRequested();
-            await SynchronizeTextChangesAsync(oldDocument, newDocument, cancellationToken).ConfigureAwait(false);
-        }
+        var client = await RemoteHostClient.TryGetClientAsync(_workspace, cancellationToken).ConfigureAwait(false);
+        if (client == null)
+            return;
 
-        return;
+        // this pushes text changes to the remote side if it can. this is purely perf optimization. whether this
+        // pushing text change worked or not doesn't affect feature's functionality.
+        //
+        // this basically see whether it can cheaply find out text changes between 2 snapshots, if it can, it will
+        // send out that text changes to remote side.
+        //
+        // the remote side, once got the text change, will again see whether it can use that text change information
+        // without any high cost and create new snapshot from it.
+        //
+        // otherwise, it will do the normal behavior of getting full text from VS side. this optimization saves
+        // times we need to do full text synchronization for typing scenario.
+        using var _ = ArrayBuilder<(DocumentId id, Checksum textChecksum, ImmutableArray<TextChange> changes)>.GetInstance(out var builder);
 
-        async ValueTask SynchronizeTextChangesAsync(Document oldDocument, Document newDocument, CancellationToken cancellationToken)
+        foreach (var (oldDocument, newDocument) in values)
         {
-            // this pushes text changes to the remote side if it can. this is purely perf optimization. whether this
-            // pushing text change worked or not doesn't affect feature's functionality.
-            //
-            // this basically see whether it can cheaply find out text changes between 2 snapshots, if it can, it will
-            // send out that text changes to remote side.
-            //
-            // the remote side, once got the text change, will again see whether it can use that text change information
-            // without any high cost and create new snapshot from it.
-            //
-            // otherwise, it will do the normal behavior of getting full text from VS side. this optimization saves
-            // times we need to do full text synchronization for typing scenario.
-
             if (!oldDocument.TryGetText(out var oldText) ||
                 !newDocument.TryGetText(out var newText))
             {
                 // we only support case where text already exist
-                return;
+                continue;
             }
 
             // Avoid allocating text before seeing if we can bail out.
             var changeRanges = newText.GetChangeRanges(oldText).AsImmutable();
             if (changeRanges.Length == 0)
-                return;
+                continue;
 
             // no benefit here. pulling from remote host is more efficient
             if (changeRanges is [{ Span.Length: var singleChangeLength }] && singleChangeLength == oldText.Length)
-                return;
-
-            var textChanges = newText.GetTextChanges(oldText).AsImmutable();
-
-            var client = await RemoteHostClient.TryGetClientAsync(_workspace, cancellationToken).ConfigureAwait(false);
-            if (client == null)
-                return;
+                continue;
 
             var state = await oldDocument.State.GetStateChecksumsAsync(cancellationToken).ConfigureAwait(false);
 
-            await client.TryInvokeAsync<IRemoteAssetSynchronizationService>(
-                (service, cancellationToken) => service.SynchronizeTextAsync(oldDocument.Id, state.Text, textChanges, cancellationToken),
-                cancellationToken).ConfigureAwait(false);
+            var textChanges = newText.GetTextChanges(oldText).AsImmutable();
+            builder.Add((oldDocument.Id, state.Text, textChanges));
         }
+
+        if (builder.Count == 0)
+            return;
+
+        await client.TryInvokeAsync<IRemoteAssetSynchronizationService>(
+            (service, cancellationToken) => service.SynchronizeTextChangesAsync(builder.ToImmutableAndClear(), cancellationToken),
+            cancellationToken).ConfigureAwait(false);
     }
 }

src/VisualStudio/Core/Test.Next/Services/ServiceHubServicesTests.cs

@@ -91,7 +91,7 @@ public async Task TestRemoteHostTextSynchronize()
 
             // sync
             await client.TryInvokeAsync<IRemoteAssetSynchronizationService>(
-                (service, cancellationToken) => service.SynchronizeTextAsync(oldDocument.Id, oldState.Text, newText.GetTextChanges(oldText).AsImmutable(), cancellationToken),
+                (service, cancellationToken) => service.SynchronizeTextChangesAsync([(oldDocument.Id, oldState.Text, newText.GetTextChanges(oldText).AsImmutable())], cancellationToken),
                 CancellationToken.None);
 
             // apply change to solution

src/VisualStudio/Core/Test.Next/Services/SolutionServiceTests.cs

@@ -1048,6 +1048,48 @@ class Program2
         objectReference2.AssertReleased();
     }
 
+    [Fact]
+    public void TestActiveAndRelatedDocumentSemanticModelCached()
+    {
+        using var workspace = TestWorkspace.Create("""
+            <Workspace>
+                <Project Language="C#" AssemblyName="Assembly1" CommonReferences="true">
+                    <Document FilePath="File1.cs">
+                        class Program1
+                        {
+                        }
+                    </Document>
+                    <Document FilePath="File2.cs">
+                        class Program2
+                        {
+                        }
+                    </Document>
+                </Project>
+                <Project Language="C#" AssemblyName="Assembly2" CommonReferences="true">
+                    <Document IsLinkFile="true" LinkAssemblyName="Assembly1" LinkFilePath="File1.cs" />
+                </Project>
+            </Workspace>
+            """, composition: s_compositionWithFirstDocumentIsActiveAndVisible);
+        using var remoteWorkspace = CreateRemoteWorkspace();
+
+        var solution = workspace.CurrentSolution;
+
+        var project1 = solution.Projects.Single(p => p.AssemblyName == "Assembly1");
+        var project2 = solution.Projects.Single(p => p.AssemblyName == "Assembly2");
+        var document1 = project1.Documents.First();
+        var document2 = project1.Documents.Last();
+        var document3 = project2.Documents.Single();
+
+        // Only the semantic model for the active document should be cached.
+        var objectReference1 = ObjectReference.CreateFromFactory(() => document1.GetSemanticModelAsync().GetAwaiter().GetResult());
+        var objectReference2 = ObjectReference.CreateFromFactory(() => document2.GetSemanticModelAsync().GetAwaiter().GetResult());
+        var objectReference3 = ObjectReference.CreateFromFactory(() => document3.GetSemanticModelAsync().GetAwaiter().GetResult());
+
+        objectReference1.AssertHeld();
+        objectReference2.AssertReleased();
+        objectReference3.AssertHeld();
+    }
+
     [Fact]
     public async Task TestRemoteWorkspaceCachesNothingIfActiveDocumentNotSynced()
     {

src/Workspaces/Core/Portable/Workspace/Solution/Solution_SemanticModelCaching.cs

@@ -2,58 +2,50 @@
 // 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.Collections.Immutable;
+using Roslyn.Utilities;
+
 namespace Microsoft.CodeAnalysis;
 
 public partial class Solution
 {
     /// <summary>
-    /// Strongly held reference to the semantic model for the active document.  By strongly holding onto it, we ensure
-    /// that it won't be GC'ed between feature requests from multiple features that care about it.  As the active
-    /// document has the most features running on it continuously, we definitely do not want to drop this.  Note: this
-    /// cached value is only to help with performance.  Not with correctness.  Importantly, the concept of 'active
-    /// document' is itself fundamentally racy.  That's ok though as we simply want to settle on these semantic models
-    /// settling into a stable state over time.  We don't need to be perfect about it.  They are intentionally not
-    /// locked either as we would only have contention right when switching to a new active document, and we would still
-    /// latch onto the new document very quickly.
+    /// Strongly held reference to the semantic models for the active document (and its related documents linked into
+    /// other projects).  By strongly holding onto them, we ensure that they won't be GC'ed between feature requests
+    /// from multiple features that care about it.  As the active document has the most features running on it
+    /// continuously, we definitely do not want to drop this.  Note: this cached value is only to help with performance.
+    /// Not with correctness.  Importantly, the concept of 'active document' is itself fundamentally racy.  That's ok
+    /// though as we simply want to settle on these semantic models settling into a stable state over time.  We don't
+    /// need to be perfect about it.
     /// </summary>
-    /// <remarks>
-    /// It is fine for these fields to never be read.  The purpose is simply to keep a strong reference around so that
-    /// they will not be GC'ed as long as the active document stays the same.
-    /// </remarks>
-#pragma warning disable IDE0052 // Remove unread private members
-    private SemanticModel? _activeDocumentSemanticModel;
-
-    /// <inheritdoc cref="_activeDocumentSemanticModel"/>
-    private SemanticModel? _activeDocumentNullableDisabledSemanticModel;
-#pragma warning restore IDE0052 // Remove unread private members
-
-    internal void OnSemanticModelObtained(DocumentId documentId, SemanticModel semanticModel)
+    private ImmutableArray<(DocumentId documentId, SemanticModel semanticModel)> _activeSemanticModels = [];
+
+    internal void OnSemanticModelObtained(
+        DocumentId documentId, SemanticModel semanticModel)
     {
         var service = this.Services.GetRequiredService<IDocumentTrackingService>();
 
-        var activeDocumentId = service.TryGetActiveDocument();
-        if (activeDocumentId is null)
-        {
-            // no active document?  then clear out any caches we have.
-            _activeDocumentSemanticModel = null;
-            _activeDocumentNullableDisabledSemanticModel = null;
-        }
-        else if (activeDocumentId != documentId)
-        {
-            // We have an active document, but we just obtained the semantic model for some other doc.  Nothing to do
-            // here, we don't want to cache this.
+        // Operate on a local reference to the immutable array to ensure a consistent view of it.
+        var localArray = _activeSemanticModels;
+
+        // No need to do anything if we're already caching this pair.
+        if (localArray.Contains((documentId, semanticModel)))
             return;
-        }
-        else
-        {
-            // Ok.  We just obtained the semantic model for the active document.  Make a strong reference to it so that
-            // other features that wake up for this active document are sure to be able to reuse the same one.
-#pragma warning disable RSEXPERIMENTAL001 // sym-shipped usage of experimental API
-            if (semanticModel.NullableAnalysisIsDisabled)
-                _activeDocumentNullableDisabledSemanticModel = semanticModel;
-            else
-                _activeDocumentSemanticModel = semanticModel;
-#pragma warning restore RSEXPERIMENTAL001
-        }
+
+        var activeDocumentId = service.TryGetActiveDocument();
+        var relatedDocumentIds = activeDocumentId is null ? [] : this.GetRelatedDocumentIds(activeDocumentId);
+
+        // Remove any cached values for documents that are no longer the active document.
+        localArray = localArray.RemoveAll(
+            tuple => !relatedDocumentIds.Contains(tuple.documentId));
+
+        // Now cache this doc/semantic model pair if it's in the related document set.
+        if (relatedDocumentIds.Contains(documentId))
+            localArray = localArray.Add((documentId, semanticModel));
+
+        // Note: this code is racy. We could have two threads executing the code above, while only one thread will win
+        // here.  We accept that as this code is just intended to help just by making some strong references to semantic
+        // models to prevent them from being GC'ed.  We don't need to be perfect about it.
+        _activeSemanticModels = localArray;
     }
 }

src/Workspaces/Remote/Core/IRemoteAssetSynchronizationService.cs

@@ -21,11 +21,13 @@ internal interface IRemoteAssetSynchronizationService
     /// Synchronize the text changes made by a user to a particular document as they are editing it.  By sending over
     /// the text changes as they happen, we can attempt to 'prime' the remote asset cache with a final <see
     /// cref="SourceText"/> that is built based off of retrieving the remote source text with a checksum corresponding
-    /// to <paramref name="baseTextChecksum"/> and then applying the <paramref name="textChanges"/> to it.  Then, when
-    /// the next remote call comes in for the new solution snapshot, it can hopefully just pluck that text out of the
-    /// cache without having to sync the <em>entire</em> contents of the file over.
+    /// to baseTextChecksum and then applying the textChanges to it.  Then, when the next remote call comes in for the
+    /// new solution snapshot, it can hopefully just pluck that text out of the cache without having to sync the
+    /// <em>entire</em> contents of the file over.
     /// </summary>
-    ValueTask SynchronizeTextAsync(DocumentId documentId, Checksum baseTextChecksum, ImmutableArray<TextChange> textChanges, CancellationToken cancellationToken);
+    ValueTask SynchronizeTextChangesAsync(
+        ImmutableArray<(DocumentId documentId, Checksum baseTextChecksum, ImmutableArray<TextChange> textChanges)> changes,
+        CancellationToken cancellationToken);
 
     /// <summary>
     /// Synchronize over what the user's current active document is that they're editing.  This can then be used by the