Replace BatchingWorkQueue with AsyncBatchingWorkQueue from Roslyn

eng/Versions.props

@@ -151,6 +151,7 @@
     <NuGetSolutionRestoreManagerInteropVersion>4.8.0</NuGetSolutionRestoreManagerInteropVersion>
     <StreamJsonRpcPackageVersion>2.17.11</StreamJsonRpcPackageVersion>
     <SystemRuntimeInteropServicesRuntimePackageVersion>4.3.0</SystemRuntimeInteropServicesRuntimePackageVersion>
+    <SystemThreadingTasksExtensions>4.5.4</SystemThreadingTasksExtensions>
     <Tooling_MicrosoftCodeAnalysisAnalyzersPackageVersion>3.11.0-beta1.24170.2</Tooling_MicrosoftCodeAnalysisAnalyzersPackageVersion>
     <Tooling_MicrosoftCodeAnalysisBannedApiAnalyzersPackageVersion>$(Tooling_MicrosoftCodeAnalysisAnalyzersPackageVersion)</Tooling_MicrosoftCodeAnalysisBannedApiAnalyzersPackageVersion>
     <Tooling_RoslynDiagnosticsAnalyzersPackageVersion>$(Tooling_MicrosoftCodeAnalysisAnalyzersPackageVersion)</Tooling_RoslynDiagnosticsAnalyzersPackageVersion>

src/Razor/src/Microsoft.AspNetCore.Razor.ExternalAccess.RoslynWorkspace/TaskDelayScheduler.cs

@@ -1,6 +1,8 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
+using Microsoft.AspNetCore.Razor.Threading;
+
 namespace Microsoft.AspNetCore.Razor.ExternalAccess.RoslynWorkspace;
 
 // Copied from https://github.com/dotnet/project-system/blob/e4db47666e0a49f6c38e701f8630dbc31380fb64/src/Microsoft.VisualStudio.ProjectSystem.Managed/Threading/Tasks/TaskDelayScheduler.cs

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Extensions/IServiceCollectionExtensions.cs

@@ -139,7 +139,7 @@ public static void AddSemanticTokensServices(this IServiceCollection services, L
 
         services.AddHandler<RazorSemanticTokensRefreshEndpoint>();
 
-        services.AddSingleton<IWorkspaceSemanticTokensRefreshPublisher, WorkspaceSemanticTokensRefreshPublisher>();
+        services.AddSingleton<IWorkspaceSemanticTokensRefreshNotifier, WorkspaceSemanticTokensRefreshNotifier>();
         services.AddSingleton<IRazorStartupService, WorkspaceSemanticTokensRefreshTrigger>();
     }
 

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/IWorkspaceSemanticTokensRefreshPublisher.cs → src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/IWorkspaceSemanticTokensRefreshNotifier.cs

@@ -3,7 +3,7 @@
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer;
 
-internal interface IWorkspaceSemanticTokensRefreshPublisher
+internal interface IWorkspaceSemanticTokensRefreshNotifier
 {
-    void EnqueueWorkspaceSemanticTokensRefresh();
+    void NotifyWorkspaceSemanticTokensRefresh();
 }

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

@@ -0,0 +1,39 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.Collections.Generic;
+using Microsoft.CodeAnalysis.Razor;
+using Microsoft.CodeAnalysis.Razor.ProjectSystem;
+
+namespace Microsoft.AspNetCore.Razor.LanguageServer;
+
+internal partial class OpenDocumentGenerator
+{
+    private sealed class Comparer : IEqualityComparer<IDocumentSnapshot>
+    {
+        public static readonly Comparer Instance = new();
+
+        private Comparer()
+        {
+        }
+
+        public bool Equals(IDocumentSnapshot? x, IDocumentSnapshot? y)
+        {
+            if (x is null)
+            {
+                return y is null;
+            }
+            else if (y is null)
+            {
+                return false;
+            }
+
+            return FilePathComparer.Instance.Equals(x.FilePath, y.FilePath);
+        }
+
+        public int GetHashCode(IDocumentSnapshot obj)
+        {
+            return FilePathComparer.Instance.GetHashCode(obj);
+        }
+    }
+}

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

@@ -3,46 +3,54 @@
 
 using System;
 using System.Collections.Generic;
-using System.Linq;
+using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
+using Microsoft.CodeAnalysis.Razor.Utilities;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer;
 
-internal class OpenDocumentGenerator : IRazorStartupService, IDisposable
+internal partial class OpenDocumentGenerator : IRazorStartupService, IDisposable
 {
     // Using 10 milliseconds for the delay here because we want document synchronization to be very fast,
     // so that features like completion are not delayed, but at the same time we don't want to do more work
     // than necessary when both C# and HTML documents change at the same time, firing our event handler
     // twice. Through testing 10ms was a good balance towards providing some de-bouncing but having minimal
     // to no impact on results.
+    //
     // It's worth noting that the queue implementation means that this delay is not restarted with each new
-    // work item, so even in very high speed typing, with changings coming in at sub-10-millisecond speed,
+    // work item, so even in very high speed typing, with changes coming in at sub-10-millisecond speed,
     // the queue will still process documents even if the user doesn't pause at all, but also will not process
     // a document for each keystroke.
-    private static readonly TimeSpan s_batchingTimeSpan = TimeSpan.FromMilliseconds(10);
+    private static readonly TimeSpan s_delay = TimeSpan.FromMilliseconds(10);
 
+    private readonly ImmutableArray<DocumentProcessedListener> _listeners;
     private readonly IProjectSnapshotManager _projectManager;
     private readonly ProjectSnapshotManagerDispatcher _dispatcher;
     private readonly LanguageServerFeatureOptions _options;
-    private readonly IReadOnlyList<DocumentProcessedListener> _listeners;
-    private readonly BatchingWorkQueue _workQueue;
+
+    private readonly AsyncBatchingWorkQueue<IDocumentSnapshot> _workQueue;
+    private readonly CancellationTokenSource _disposeTokenSource;
 
     public OpenDocumentGenerator(
         IEnumerable<DocumentProcessedListener> listeners,
         IProjectSnapshotManager projectManager,
         ProjectSnapshotManagerDispatcher dispatcher,
-        LanguageServerFeatureOptions options,
-        IErrorReporter errorReporter)
+        LanguageServerFeatureOptions options)
     {
-        _listeners = listeners.ToArray();
+        _listeners = listeners.ToImmutableArray();
         _projectManager = projectManager;
         _dispatcher = dispatcher;
         _options = options;
-        _workQueue = new BatchingWorkQueue(s_batchingTimeSpan, FilePathComparer.Instance, errorReporter);
+
+        _disposeTokenSource = new();
+        _workQueue = new AsyncBatchingWorkQueue<IDocumentSnapshot>(
+            s_delay,
+            ProcessBatchAsync,
+            _disposeTokenSource.Token);
 
         _projectManager.Changed += ProjectManager_Changed;
 
@@ -54,7 +62,41 @@ public OpenDocumentGenerator(
 
     public void Dispose()
     {
-        _workQueue.Dispose();
+        _disposeTokenSource.Cancel();
+        _disposeTokenSource.Dispose();
+    }
+
+    private async ValueTask ProcessBatchAsync(ImmutableArray<IDocumentSnapshot> items, CancellationToken token)
+    {
+        foreach (var document in items.GetMostRecentUniqueItems(Comparer.Instance))
+        {
+            if (token.IsCancellationRequested)
+            {
+                return;
+            }
+
+            var codeDocument = await document.GetGeneratedOutputAsync().ConfigureAwait(false);
+
+            await _dispatcher
+                .RunAsync(
+                    static state =>
+                    {
+                        var (codeDocument, document, listeners, token) = state;
+
+                        foreach (var listener in listeners)
+                        {
+                            if (token.IsCancellationRequested)
+                            {
+                                return;
+                            }
+
+                            listener.DocumentProcessed(codeDocument, document);
+                        }
+                    },
+                    state: (codeDocument, document, _listeners, token),
+                    token)
+                .ConfigureAwait(false);
+        }
     }
 
     private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
@@ -65,8 +107,6 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
             return;
         }
 
-        _dispatcher.AssertRunningOnDispatcher();
-
         switch (args.Kind)
         {
             case ProjectChangeKind.ProjectChanged:
@@ -75,9 +115,9 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
 
                     foreach (var documentFilePath in newProject.DocumentFilePaths)
                     {
-                        if (newProject.GetDocument(documentFilePath) is { } document)
+                        if (newProject.TryGetDocument(documentFilePath, out var document))
                         {
-                            TryEnqueue(document);
+                            EnqueueIfNecessary(document);
                         }
                     }
 
@@ -89,13 +129,13 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                     var newProject = args.Newer.AssumeNotNull();
                     var documentFilePath = args.DocumentFilePath.AssumeNotNull();
 
-                    if (newProject.GetDocument(documentFilePath) is { } document)
+                    if (newProject.TryGetDocument(documentFilePath, out var document))
                     {
-                        // We don't enqueue the current document because added documents are by default closed.
+                        // We don't enqueue the current document because added documents are initially closed.
 
                         foreach (var relatedDocument in newProject.GetRelatedDocuments(document))
                         {
-                            TryEnqueue(relatedDocument);
+                            EnqueueIfNecessary(relatedDocument);
                         }
                     }
 
@@ -107,13 +147,13 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                     var newProject = args.Newer.AssumeNotNull();
                     var documentFilePath = args.DocumentFilePath.AssumeNotNull();
 
-                    if (newProject.GetDocument(documentFilePath) is { } document)
+                    if (newProject.TryGetDocument(documentFilePath, out var document))
                     {
-                        TryEnqueue(document);
+                        EnqueueIfNecessary(document);
 
                         foreach (var relatedDocument in newProject.GetRelatedDocuments(document))
                         {
-                            TryEnqueue(relatedDocument);
+                            EnqueueIfNecessary(relatedDocument);
                         }
                     }
 
@@ -126,15 +166,15 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                     var oldProject = args.Older.AssumeNotNull();
                     var documentFilePath = args.DocumentFilePath.AssumeNotNull();
 
-                    if (oldProject.GetDocument(documentFilePath) is { } document)
+                    if (oldProject.TryGetDocument(documentFilePath, out var document))
                     {
                         foreach (var relatedDocument in oldProject.GetRelatedDocuments(document))
                         {
                             var relatedDocumentFilePath = relatedDocument.FilePath.AssumeNotNull();
 
-                            if (newProject.GetDocument(relatedDocumentFilePath) is { } newRelatedDocument)
+                            if (newProject.TryGetDocument(relatedDocumentFilePath, out var newRelatedDocument))
                             {
-                                TryEnqueue(newRelatedDocument);
+                                EnqueueIfNecessary(newRelatedDocument);
                             }
                         }
                     }
@@ -149,49 +189,15 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                 }
         }
 
-        void TryEnqueue(IDocumentSnapshot document)
+        void EnqueueIfNecessary(IDocumentSnapshot document)
         {
-            var documentFilePath = document.FilePath.AssumeNotNull();
-
-            if (!_projectManager.IsDocumentOpen(documentFilePath) &&
+            if (!_projectManager.IsDocumentOpen(document.FilePath.AssumeNotNull()) &&
                 !_options.UpdateBuffersForClosedDocuments)
             {
                 return;
             }
 
-            var key = $"{document.Project.Key.Id}:{documentFilePath}";
-            var workItem = new ProcessWorkItem(document, _listeners, _dispatcher);
-            _workQueue.Enqueue(key, workItem);
-        }
-    }
-
-    private class ProcessWorkItem : BatchableWorkItem
-    {
-        private readonly IDocumentSnapshot _latestDocument;
-        private readonly IEnumerable<DocumentProcessedListener> _documentProcessedListeners;
-        private readonly ProjectSnapshotManagerDispatcher _dispatcher;
-
-        public ProcessWorkItem(
-            IDocumentSnapshot latestDocument,
-            IReadOnlyList<DocumentProcessedListener> documentProcessedListeners,
-            ProjectSnapshotManagerDispatcher dispatcher)
-        {
-            _latestDocument = latestDocument;
-            _documentProcessedListeners = documentProcessedListeners;
-            _dispatcher = dispatcher;
-        }
-
-        public override async ValueTask ProcessAsync(CancellationToken cancellationToken)
-        {
-            var codeDocument = await _latestDocument.GetGeneratedOutputAsync().ConfigureAwait(false);
-
-            await _dispatcher.RunAsync(() =>
-            {
-                foreach (var listener in _documentProcessedListeners)
-                {
-                    listener.DocumentProcessed(codeDocument, _latestDocument);
-                }
-            }, cancellationToken).ConfigureAwait(false);
+            _workQueue.AddWork(document);
         }
     }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Semantic/RazorSemanticTokensRefreshEndpoint.cs

@@ -9,16 +9,16 @@
 namespace Microsoft.AspNetCore.Razor.LanguageServer.Semantic;
 
 [RazorLanguageServerEndpoint(CustomMessageNames.RazorSemanticTokensRefreshEndpoint)]
-internal sealed class RazorSemanticTokensRefreshEndpoint(IWorkspaceSemanticTokensRefreshPublisher semanticTokensRefreshPublisher) : IRazorNotificationHandler<SemanticTokensRefreshParams>
+internal sealed class RazorSemanticTokensRefreshEndpoint(IWorkspaceSemanticTokensRefreshNotifier semanticTokensRefreshPublisher) : IRazorNotificationHandler<SemanticTokensRefreshParams>
 {
-    private readonly IWorkspaceSemanticTokensRefreshPublisher _semanticTokensRefreshPublisher = semanticTokensRefreshPublisher;
+    private readonly IWorkspaceSemanticTokensRefreshNotifier _semanticTokensRefreshPublisher = semanticTokensRefreshPublisher;
 
     public bool MutatesSolutionState { get; } = false;
 
     public Task HandleNotificationAsync(SemanticTokensRefreshParams request, RazorRequestContext requestContext, CancellationToken cancellationToken)
     {
         // We have to invalidate the tokens cache since it may no longer be up to date.
-        _semanticTokensRefreshPublisher.EnqueueWorkspaceSemanticTokensRefresh();
+        _semanticTokensRefreshPublisher.NotifyWorkspaceSemanticTokensRefresh();
 
         return Task.CompletedTask;
     }