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

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.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,23 @@ 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;
 
     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 +69,22 @@ 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 (!_projectManager.TryGetDocument(key, out var document))
+            {
+                continue;
+            }
+
             var codeDocument = await document.GetGeneratedOutputAsync(token).ConfigureAwait(false);
 
             foreach (var listener in _listeners)
@@ -109,7 +119,7 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                     {
                         if (newProject.TryGetDocument(documentFilePath, out var document))
                         {
-                            EnqueueIfNecessary(document);
+                            EnqueueIfNecessary(document.Key, document.Version);
                         }
                     }
 
@@ -127,11 +137,11 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
 
                     if (newProject.TryGetDocument(documentFilePath, out var document))
                     {
-                        EnqueueIfNecessary(document);
+                        EnqueueIfNecessary(document.Key, document.Version);
 
                         foreach (var relatedDocument in newProject.GetRelatedDocuments(document))
                         {
-                            EnqueueIfNecessary(relatedDocument);
+                            EnqueueIfNecessary(relatedDocument.Key, relatedDocument.Version);
                         }
                     }
 
@@ -152,7 +162,7 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
 
                             if (newProject.TryGetDocument(relatedDocumentFilePath, out var newRelatedDocument))
                             {
-                                EnqueueIfNecessary(newRelatedDocument);
+                                EnqueueIfNecessary(newRelatedDocument.Key, newRelatedDocument.Version);
                             }
                         }
                     }
@@ -167,17 +177,17 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                 }
         }
 
-        void EnqueueIfNecessary(DocumentSnapshot document)
+        void EnqueueIfNecessary(DocumentKey documentKey, int documentVersion)
         {
-            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}");
+            _logger.LogDebug($"Enqueuing generation of {documentKey.FilePath} in {documentKey.ProjectKey.Id} at version {documentVersion}");
 
-            _workQueue.AddWork(document);
+            _workQueue.AddWork(documentKey);
         }
     }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.ProjectEngineHost/ProjectSystem/DocumentKey.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/ProjectKey.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.CodeAnalysis.Razor.Workspaces/ProjectSystem/AbstractRazorProjectInfoDriver.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)
             {

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/DocumentSnapshot.cs

@@ -19,6 +19,7 @@ internal sealed class DocumentSnapshot(ProjectSnapshot project, DocumentState st
 
     public HostDocument HostDocument => _state.HostDocument;
 
+    public DocumentKey Key => new(Project.Key, FilePath);
     public string FileKind => _state.HostDocument.FileKind;
     public string FilePath => _state.HostDocument.FilePath;
     public string TargetPath => _state.HostDocument.TargetPath;

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/ProjectSnapshotManagerExtensions.cs

@@ -42,6 +42,21 @@ public static bool TryGetDocument(
     public static DocumentSnapshot GetRequiredDocument(this ProjectSnapshotManager projectManager, ProjectKey projectKey, string documentFilePath)
         => projectManager.GetDocument(projectKey, documentFilePath).AssumeNotNull();
 
+    public static bool ContainsDocument(this ProjectSnapshotManager projectManager, DocumentKey documentKey)
+        => projectManager.ContainsDocument(documentKey.ProjectKey, documentKey.FilePath);
+
+    public static bool TryGetDocument(
+        this ProjectSnapshotManager projectManager,
+        DocumentKey documentKey,
+        [NotNullWhen(true)] out DocumentSnapshot? document)
+        => projectManager.TryGetDocument(documentKey.ProjectKey, documentKey.FilePath, out document);
+
+    public static DocumentSnapshot? GetDocument(this ProjectSnapshotManager projectManager, DocumentKey documentKey)
+        => projectManager.GetDocument(documentKey.ProjectKey, documentKey.FilePath);
+
+    public static DocumentSnapshot GetRequiredDocument(this ProjectSnapshotManager projectManager, DocumentKey documentKey)
+        => projectManager.GetRequiredDocument(documentKey.ProjectKey, documentKey.FilePath);
+
     public static ProjectSnapshot? GetProject(this ProjectSnapshotManager.Updater updater, ProjectKey projectKey)
         => updater.TryGetProject(projectKey, out var result)
             ? result
@@ -74,4 +89,19 @@ public static bool TryGetDocument(
 
     public static DocumentSnapshot GetRequiredDocument(this ProjectSnapshotManager.Updater updater, ProjectKey projectKey, string documentFilePath)
         => updater.GetDocument(projectKey, documentFilePath).AssumeNotNull();
+
+    public static bool ContainsDocument(this ProjectSnapshotManager.Updater updater, DocumentKey documentKey)
+        => updater.ContainsDocument(documentKey.ProjectKey, documentKey.FilePath);
+
+    public static bool TryGetDocument(
+        this ProjectSnapshotManager.Updater updater,
+        DocumentKey documentKey,
+        [NotNullWhen(true)] out DocumentSnapshot? document)
+        => updater.TryGetDocument(documentKey.ProjectKey, documentKey.FilePath, out document);
+
+    public static DocumentSnapshot? GetDocument(this ProjectSnapshotManager.Updater updater, DocumentKey documentKey)
+        => updater.GetDocument(documentKey.ProjectKey, documentKey.FilePath);
+
+    public static DocumentSnapshot GetRequiredDocument(this ProjectSnapshotManager.Updater updater, DocumentKey documentKey)
+        => updater.GetRequiredDocument(documentKey.ProjectKey, documentKey.FilePath);
 }