Use DocumentKey to avoid holding Razor project/document snapshots when the most recent will be used. (#11644)

The goal of this change is to move `DocumentKey` close to `ProjectKey`
and use it various async batching work queues rather than holding onto
project or document snapshots. I recommend reviewing commit-by-commit.
Each commit describes its changes.

CI Build:
https://dev.azure.com/dnceng/internal/_build/results?buildId=2668867&view=results
Test Insertion:
https://dev.azure.com/devdiv/DevDiv/_git/VS/pullrequest/621273

Author: DustinCampbell

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/GeneratedDocumentPublisher.cs

@@ -180,7 +180,7 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                     {
                         if (_publishedCSharpData.Remove(key))
                         {
-                            _logger.LogDebug($"Removing previous C# publish data for {key.ProjectKey}/{key.DocumentFilePath}");
+                            _logger.LogDebug($"Removing previous C# publish data for {key.ProjectKey}/{key.FilePath}");
                         }
                     }
 
@@ -208,15 +208,15 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                     {
                         if (_publishedCSharpData.Remove(documentKey))
                         {
-                            _logger.LogDebug($"Removing previous C# publish data for {documentKey.ProjectKey}/{documentKey.DocumentFilePath}");
+                            _logger.LogDebug($"Removing previous C# publish data for {documentKey.ProjectKey}/{documentKey.FilePath}");
                         }
                     }
 
                     lock (_publishedHtmlData)
                     {
                         if (_publishedHtmlData.Remove(documentFilePath))
                         {
-                            _logger.LogDebug($"Removing previous Html publish data for {documentKey.ProjectKey}/{documentKey.DocumentFilePath}");
+                            _logger.LogDebug($"Removing previous Html publish data for {documentKey.ProjectKey}/{documentKey.FilePath}");
                         }
                     }
                 }
@@ -249,7 +249,7 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                         {
                             if (_publishedCSharpData.Remove(key))
                             {
-                                _logger.LogDebug($"Removing previous C# publish data for {key.ProjectKey}/{key.DocumentFilePath}");
+                                _logger.LogDebug($"Removing previous C# publish data for {key.ProjectKey}/{key.FilePath}");
                             }
                         }
                     }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/OpenDocumentGenerator.Comparer.cs

@@ -6,6 +6,7 @@
 using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.ProjectSystem;
 using Microsoft.AspNetCore.Razor.Utilities;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
@@ -32,21 +33,29 @@ internal partial class OpenDocumentGenerator : IRazorStartupService, IDisposable
     private readonly LanguageServerFeatureOptions _options;
     private readonly ILogger _logger;
 
-    private readonly AsyncBatchingWorkQueue<DocumentSnapshot> _workQueue;
+    private readonly AsyncBatchingWorkQueue<DocumentKey> _workQueue;
     private readonly CancellationTokenSource _disposeTokenSource;
+    private readonly HashSet<DocumentKey> _workerSet;
+
+    // Note: This is likely to always be false. Only the Visual Studio ProjectSnapshotManager
+    // is notified of the solution opening and closing, so the language server shouldn't
+    // update this value. However, this may change at some point and keeping the check here means
+    // that the logic between this class and the Visual Studio BackgroundDocumentGenerator are in sync.
+    private bool _solutionIsClosing;
 
     public OpenDocumentGenerator(
         IEnumerable<IDocumentProcessedListener> listeners,
         ProjectSnapshotManager projectManager,
         LanguageServerFeatureOptions options,
         ILoggerFactory loggerFactory)
     {
-        _listeners = listeners.ToImmutableArray();
+        _listeners = [.. listeners];
         _projectManager = projectManager;
         _options = options;
 
+        _workerSet = [];
         _disposeTokenSource = new();
-        _workQueue = new AsyncBatchingWorkQueue<DocumentSnapshot>(
+        _workQueue = new AsyncBatchingWorkQueue<DocumentKey>(
             s_delay,
             ProcessBatchAsync,
             _disposeTokenSource.Token);
@@ -66,15 +75,30 @@ public void Dispose()
         _disposeTokenSource.Dispose();
     }
 
-    private async ValueTask ProcessBatchAsync(ImmutableArray<DocumentSnapshot> items, CancellationToken token)
+    private async ValueTask ProcessBatchAsync(ImmutableArray<DocumentKey> items, CancellationToken token)
     {
-        foreach (var document in items.GetMostRecentUniqueItems(Comparer.Instance))
+        _workerSet.Clear();
+
+        foreach (var key in items.GetMostRecentUniqueItems(_workerSet))
         {
             if (token.IsCancellationRequested)
             {
                 return;
             }
 
+            // If the solution is closing, avoid any in-progress work.
+            if (_solutionIsClosing)
+            {
+                return;
+            }
+
+            if (!_projectManager.TryGetDocument(key, out var document))
+            {
+                continue;
+            }
+
+            _logger.LogDebug($"Generating {key} at version {document.Version}");
+
             var codeDocument = await document.GetGeneratedOutputAsync(token).ConfigureAwait(false);
 
             foreach (var listener in _listeners)
@@ -91,26 +115,27 @@ private async ValueTask ProcessBatchAsync(ImmutableArray<DocumentSnapshot> items
 
     private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
     {
-        // Don't do any work if the solution is closing
+        // We don't want to do any work on solution close.
         if (args.IsSolutionClosing)
         {
+            _solutionIsClosing = true;
             return;
         }
 
+        _solutionIsClosing = false;
+
         _logger.LogDebug($"Got a project change of type {args.Kind} for {args.ProjectKey.Id}");
 
         switch (args.Kind)
         {
+            case ProjectChangeKind.ProjectAdded:
             case ProjectChangeKind.ProjectChanged:
                 {
                     var newProject = args.Newer.AssumeNotNull();
 
                     foreach (var documentFilePath in newProject.DocumentFilePaths)
                     {
-                        if (newProject.TryGetDocument(documentFilePath, out var document))
-                        {
-                            EnqueueIfNecessary(document);
-                        }
+                        EnqueueIfNecessary(new(newProject.Key, documentFilePath));
                     }
 
                     break;
@@ -125,14 +150,11 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                     var newProject = args.Newer.AssumeNotNull();
                     var documentFilePath = args.DocumentFilePath.AssumeNotNull();
 
-                    if (newProject.TryGetDocument(documentFilePath, out var document))
-                    {
-                        EnqueueIfNecessary(document);
+                    EnqueueIfNecessary(new(newProject.Key, documentFilePath));
 
-                        foreach (var relatedDocument in newProject.GetRelatedDocuments(document))
-                        {
-                            EnqueueIfNecessary(relatedDocument);
-                        }
+                    foreach (var relatedDocumentFilePath in newProject.GetRelatedDocumentFilePaths(documentFilePath))
+                    {
+                        EnqueueIfNecessary(new(newProject.Key, relatedDocumentFilePath));
                     }
 
                     break;
@@ -144,16 +166,14 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                     var oldProject = args.Older.AssumeNotNull();
                     var documentFilePath = args.DocumentFilePath.AssumeNotNull();
 
-                    if (oldProject.TryGetDocument(documentFilePath, out var document))
+                    // For removals use the old snapshot to find related documents to update if they exist
+                    // in the new snapshot.
+
+                    foreach (var relatedDocumentFilePath in oldProject.GetRelatedDocumentFilePaths(documentFilePath))
                     {
-                        foreach (var relatedDocument in oldProject.GetRelatedDocuments(document))
+                        if (newProject.ContainsDocument(relatedDocumentFilePath))
                         {
-                            var relatedDocumentFilePath = relatedDocument.FilePath;
-
-                            if (newProject.TryGetDocument(relatedDocumentFilePath, out var newRelatedDocument))
-                            {
-                                EnqueueIfNecessary(newRelatedDocument);
-                            }
+                            EnqueueIfNecessary(new(newProject.Key, relatedDocumentFilePath));
                         }
                     }
 
@@ -162,22 +182,24 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
 
             case ProjectChangeKind.ProjectRemoved:
                 {
-                    // No-op. We don't need to enqueue recompilations if the project is being removed
+                    // No-op. We don't need to compile anything if the project is being removed
                     break;
                 }
+
+            default:
+                Assumed.Unreachable($"Unknown {nameof(ProjectChangeKind)}: {args.Kind}");
+                break;
         }
 
-        void EnqueueIfNecessary(DocumentSnapshot document)
+        void EnqueueIfNecessary(DocumentKey documentKey)
         {
-            if (!_projectManager.IsDocumentOpen(document.FilePath) &&
-                !_options.UpdateBuffersForClosedDocuments)
+            if (!_options.UpdateBuffersForClosedDocuments &&
+                !_projectManager.IsDocumentOpen(documentKey.FilePath))
             {
                 return;
             }
 
-            _logger.LogDebug($"Enqueuing generation of {document.FilePath} in {document.Project.Key.Id} at version {document.Version}");
-
-            _workQueue.AddWork(document);
+            _workQueue.AddWork(documentKey);
         }
     }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/OpenDocumentGenerator.cs

@@ -0,0 +1,44 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System;
+using Microsoft.CodeAnalysis.Razor;
+using Microsoft.Extensions.Internal;
+
+namespace Microsoft.AspNetCore.Razor.ProjectSystem;
+
+internal readonly record struct DocumentKey : IComparable<DocumentKey>
+{
+    public ProjectKey ProjectKey { get; }
+    public string FilePath { get; }
+
+    public DocumentKey(ProjectKey projectKey, string filePath)
+    {
+        ProjectKey = projectKey;
+        FilePath = filePath;
+    }
+
+    public bool Equals(DocumentKey other)
+        => ProjectKey.Equals(other.ProjectKey) &&
+           FilePathComparer.Instance.Equals(FilePath, other.FilePath);
+
+    public override int GetHashCode()
+    {
+        var hash = HashCodeCombiner.Start();
+        hash.Add(ProjectKey);
+        hash.Add(FilePath, FilePathComparer.Instance);
+
+        return hash;
+    }
+
+    public int CompareTo(DocumentKey other)
+    {
+        var comparison = ProjectKey.CompareTo(other.ProjectKey);
+        if (comparison != 0)
+        {
+            return comparison;
+        }
+
+        return FilePathComparer.Instance.Compare(FilePath, other.FilePath);
+    }
+}

src/Razor/src/Microsoft.AspNetCore.Razor.ProjectEngineHost/ProjectSystem/DocumentKey.cs

@@ -14,7 +14,7 @@ namespace Microsoft.AspNetCore.Razor.ProjectSystem;
 /// identifier for a project.
 /// </summary>
 [DebuggerDisplay("id: {Id}")]
-internal readonly record struct ProjectKey
+internal readonly record struct ProjectKey : IComparable<ProjectKey>
 {
     public static ProjectKey Unknown { get; } = default;
 
@@ -38,4 +38,19 @@ public override int GetHashCode()
 
     public override string ToString()
         => IsUnknown ? "<Unknown Project>" : Id;
+
+    public int CompareTo(ProjectKey other)
+    {
+        // Sort "unknown" project keys after other project keys.
+        if (IsUnknown)
+        {
+            return other.IsUnknown ? 0 : 1;
+        }
+        else if (other.IsUnknown)
+        {
+            return -1;
+        }
+
+        return FilePathComparer.Instance.Compare(Id, other.Id);
+    }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.ProjectEngineHost/ProjectSystem/ProjectKey.cs

@@ -27,6 +27,7 @@ private sealed record Remove(ProjectKey ProjectKey) : Work(ProjectKey);
 
     private readonly CancellationTokenSource _disposeTokenSource;
     private readonly AsyncBatchingWorkQueue<Work> _workQueue;
+    private readonly HashSet<Work> _workerSet;
     private readonly Dictionary<ProjectKey, RazorProjectInfo> _latestProjectInfoMap;
     private ImmutableArray<IRazorProjectInfoListener> _listeners;
     private readonly TaskCompletionSource<bool> _initializationTaskSource;
@@ -37,6 +38,7 @@ protected AbstractRazorProjectInfoDriver(ILoggerFactory loggerFactory, TimeSpan?
     {
         Logger = loggerFactory.GetOrCreateLogger(GetType());
 
+        _workerSet = new(Comparer.Instance);
         _disposeTokenSource = new();
         _workQueue = new AsyncBatchingWorkQueue<Work>(delay ?? DefaultDelay, ProcessBatchAsync, _disposeTokenSource.Token);
         _latestProjectInfoMap = [];
@@ -89,7 +91,9 @@ protected void StartInitialization()
 
     private async ValueTask ProcessBatchAsync(ImmutableArray<Work> items, CancellationToken token)
     {
-        foreach (var work in items.GetMostRecentUniqueItems(Comparer.Instance))
+        _workerSet.Clear();
+
+        foreach (var work in items.GetMostRecentUniqueItems(_workerSet))
         {
             if (token.IsCancellationRequested)
             {