Merge pull request #57824 from dibarbet/fix_lsp_rps (#57871)

[LSP] Use test accessor task instead of async operation to wait for LSP server shutdown in tests

Co-authored-by: David Barbet <dabarbet@microsoft.com>

Author: dotnet-bot

src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs

@@ -404,8 +404,7 @@ private static RequestExecutionQueue CreateRequestQueue(TestWorkspace workspace)
             var registrationService = workspace.GetService<LspWorkspaceRegistrationService>();
             var globalOptions = workspace.GetService<IGlobalOptionService>();
             var lspMiscFilesWorkspace = new LspMiscellaneousFilesWorkspace(NoOpLspLogger.Instance);
-            var listenerProvider = workspace.ExportProvider.GetExportedValue<IAsynchronousOperationListenerProvider>();
-            return new RequestExecutionQueue(NoOpLspLogger.Instance, registrationService, lspMiscFilesWorkspace, globalOptions, listenerProvider, ProtocolConstants.RoslynLspLanguages, serverName: "Tests", "TestClient");
+            return new RequestExecutionQueue(NoOpLspLogger.Instance, registrationService, lspMiscFilesWorkspace, globalOptions, ProtocolConstants.RoslynLspLanguages, serverName: "Tests", "TestClient");
         }
 
         private static string GetDocumentFilePathFromName(string documentName)

src/Features/LanguageServer/Protocol/Handler/RequestExecutionQueue.DocumentChangeTracker.cs

@@ -61,6 +61,11 @@ public ImmutableArray<SourceText> GetTrackedTexts()
 
             public bool IsComplete() => _queue._queue.IsCompleted && _queue._queue.IsEmpty;
 
+            public async Task WaitForProcessingToStopAsync()
+            {
+                await _queue._queueProcessingTask.ConfigureAwait(false);
+            }
+
             /// <summary>
             /// Test only method to validate that remaining items in the queue are cancelled.
             /// This directly mutates the queue in an unsafe way, so ensure that all relevant queue operations

src/Features/LanguageServer/Protocol/Handler/RequestExecutionQueue.cs

@@ -65,11 +65,16 @@ internal partial class RequestExecutionQueue
         private readonly CancellationTokenSource _cancelSource = new CancellationTokenSource();
         private readonly RequestTelemetryLogger _requestTelemetryLogger;
         private readonly IGlobalOptionService _globalOptions;
-        private readonly IAsynchronousOperationListener _asynchronousOperationListener;
 
         private readonly ILspLogger _logger;
         private readonly LspWorkspaceManager _lspWorkspaceManager;
 
+        /// <summary>
+        /// For test purposes only.
+        /// A task that completes when the queue processing stops.
+        /// </summary>
+        private readonly Task _queueProcessingTask;
+
         public CancellationToken CancellationToken => _cancelSource.Token;
 
         /// <summary>
@@ -87,7 +92,6 @@ public RequestExecutionQueue(
             LspWorkspaceRegistrationService lspWorkspaceRegistrationService,
             LspMiscellaneousFilesWorkspace? lspMiscellaneousFilesWorkspace,
             IGlobalOptionService globalOptions,
-            IAsynchronousOperationListenerProvider listenerProvider,
             ImmutableArray<string> supportedLanguages,
             string serverName,
             string serverTypeName)
@@ -106,9 +110,7 @@ public RequestExecutionQueue(
             _lspWorkspaceManager = new LspWorkspaceManager(logger, lspMiscellaneousFilesWorkspace, lspWorkspaceRegistrationService, _requestTelemetryLogger);
 
             // Start the queue processing
-            _asynchronousOperationListener = listenerProvider.GetListener(FeatureAttribute.LanguageServer);
-            var token = _asynchronousOperationListener.BeginAsyncOperation($"{nameof(ProcessQueueAsync)}_{serverTypeName}");
-            _ = ProcessQueueAsync().CompletesAsyncOperation(token);
+            _queueProcessingTask = ProcessQueueAsync();
         }
 
         /// <summary>

src/Features/LanguageServer/Protocol/LanguageServerTarget.cs

@@ -87,7 +87,6 @@ internal LanguageServerTarget(
                 workspaceRegistrationService,
                 lspMiscellaneousFilesWorkspace,
                 globalOptions,
-                listenerProvider,
                 supportedLanguages,
                 userVisibleServerName,
                 TelemetryServerName);

src/Features/LanguageServer/ProtocolUnitTests/LanguageServerTargetTests.cs

@@ -27,34 +27,34 @@ public class LanguageServerTargetTests : AbstractLanguageServerProtocolTests
         [Fact]
         public async Task LanguageServerQueueEmptyOnShutdownMessage()
         {
-            await using var languageServerTarget = CreateLanguageServer(out var jsonRpc, out var listenerProvider);
+            await using var languageServerTarget = CreateLanguageServer(out var jsonRpc);
             AssertServerAlive(languageServerTarget);
 
             await languageServerTarget.ShutdownAsync(CancellationToken.None).ConfigureAwait(false);
-            await AssertServerQueueClosed(languageServerTarget, listenerProvider).ConfigureAwait(false);
+            await AssertServerQueueClosed(languageServerTarget).ConfigureAwait(false);
             Assert.False(jsonRpc.IsDisposed);
         }
 
         [Fact]
         public async Task LanguageServerCleansUpOnExitMessage()
         {
-            await using var languageServerTarget = CreateLanguageServer(out var jsonRpc, out var listenerProvider);
+            await using var languageServerTarget = CreateLanguageServer(out var jsonRpc);
             AssertServerAlive(languageServerTarget);
 
             await languageServerTarget.ShutdownAsync(CancellationToken.None).ConfigureAwait(false);
             await languageServerTarget.ExitAsync(CancellationToken.None).ConfigureAwait(false);
-            await AssertServerQueueClosed(languageServerTarget, listenerProvider).ConfigureAwait(false);
+            await AssertServerQueueClosed(languageServerTarget).ConfigureAwait(false);
             Assert.True(jsonRpc.IsDisposed);
         }
 
         [Fact]
         public async Task LanguageServerCleansUpOnUnexpectedJsonRpcDisconnectAsync()
         {
-            await using var languageServerTarget = CreateLanguageServer(out var jsonRpc, out var listenerProvider);
+            await using var languageServerTarget = CreateLanguageServer(out var jsonRpc);
             AssertServerAlive(languageServerTarget);
 
             jsonRpc.Dispose();
-            await AssertServerQueueClosed(languageServerTarget, listenerProvider).ConfigureAwait(false);
+            await AssertServerQueueClosed(languageServerTarget).ConfigureAwait(false);
             Assert.True(jsonRpc.IsDisposed);
         }
 
@@ -64,14 +64,14 @@ private static void AssertServerAlive(LanguageServerTarget server)
             Assert.False(server.GetTestAccessor().GetQueueAccessor().IsComplete());
         }
 
-        private static async Task AssertServerQueueClosed(LanguageServerTarget server, IAsynchronousOperationListenerProvider listenerProvider)
+        private static async Task AssertServerQueueClosed(LanguageServerTarget server)
         {
-            await listenerProvider.GetWaiter(FeatureAttribute.LanguageServer).ExpeditedWaitAsync();
+            await server.GetTestAccessor().GetQueueAccessor().WaitForProcessingToStopAsync().ConfigureAwait(false);
             Assert.True(server.HasShutdownStarted);
             Assert.True(server.GetTestAccessor().GetQueueAccessor().IsComplete());
         }
 
-        private LanguageServerTarget CreateLanguageServer(out JsonRpc serverJsonRpc, out IAsynchronousOperationListenerProvider listenerProvider)
+        private LanguageServerTarget CreateLanguageServer(out JsonRpc serverJsonRpc)
         {
             using var workspace = TestWorkspace.CreateCSharp("", composition: Composition);
 
@@ -80,7 +80,7 @@ private LanguageServerTarget CreateLanguageServer(out JsonRpc serverJsonRpc, out
             var lspWorkspaceRegistrationService = workspace.GetService<LspWorkspaceRegistrationService>();
             var capabilitiesProvider = workspace.GetService<DefaultCapabilitiesProvider>();
             var globalOptions = workspace.GetService<IGlobalOptionService>();
-            listenerProvider = workspace.GetService<IAsynchronousOperationListenerProvider>();
+            var listenerProvider = workspace.GetService<IAsynchronousOperationListenerProvider>();
 
             serverJsonRpc = new JsonRpc(new HeaderDelimitedMessageHandler(serverStream, serverStream))
             {

src/Features/LanguageServer/ProtocolUnitTests/Ordering/RequestOrderingTests.cs

@@ -173,9 +173,7 @@ public async Task FailingMutableTaskShutsDownQueue()
 
             // The failed request returns to the client before the shutdown completes.
             // Wait for the queue to finish handling the failed request and shutdown.
-            var operations = testLspServer.TestWorkspace.ExportProvider.GetExportedValue<AsynchronousOperationListenerProvider>();
-            var waiter = operations.GetWaiter(FeatureAttribute.LanguageServer);
-            await waiter.ExpeditedWaitAsync();
+            await testLspServer.GetQueueAccessor().WaitForProcessingToStopAsync().ConfigureAwait(false);
 
             // remaining tasks should be canceled
             var areAllItemsCancelled = await testLspServer.GetQueueAccessor().AreAllItemsCancelledUnsafeAsync();