Merge pull request dotnet#57617 from dibarbet/lsp_queue_task

Move the task completion source for a queue item into the queue item itself.

src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs

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

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

@@ -6,6 +6,7 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Linq;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.LanguageServer.Handler.DocumentChanges;
 using Microsoft.CodeAnalysis.Text;
 using Roslyn.Utilities;
@@ -59,6 +60,25 @@ public ImmutableArray<SourceText> GetTrackedTexts()
             public LspWorkspaceManager GetLspWorkspaceManager() => _queue._lspWorkspaceManager;
 
             public bool IsComplete() => _queue._queue.IsCompleted && _queue._queue.IsEmpty;
+
+            /// <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
+            /// are done before calling.
+            /// </summary>
+            public async Task<bool> AreAllItemsCancelledUnsafeAsync()
+            {
+                while (!_queue._queue.IsEmpty)
+                {
+                    var (_, cancellationToken) = await _queue._queue.DequeueAsync().ConfigureAwait(false);
+                    if (!cancellationToken.IsCancellationRequested)
+                    {
+                        return false;
+                    }
+                }
+
+                return true;
+            }
         }
     }
 }

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

@@ -5,59 +5,79 @@
 #nullable enable
 
 using System;
-using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
+using Microsoft.VisualStudio.Threading;
 
 namespace Microsoft.CodeAnalysis.LanguageServer.Handler
 {
     internal partial class RequestExecutionQueue
     {
-        private readonly struct QueueItem
+        private interface IQueueItem
         {
             /// <summary>
-            /// Callback to call into underlying <see cref="IRequestHandler"/> to perform the actual work of this item.
+            /// Begins executing the work specified by this queue item.
             /// </summary>
-            private readonly Func<RequestContext?, CancellationToken, Task> _callbackAsync;
+            Task CallbackAsync(RequestContext? context, CancellationToken cancellationToken);
 
-            /// <summary>
-            /// <see cref="CorrelationManager.ActivityId"/> used to properly correlate this work with the loghub
-            /// tracing/logging subsystem.
-            /// </summary>
-            public readonly Guid ActivityId;
-            private readonly ILspLogger _logger;
+            /// <inheritdoc cref="IRequestHandler.RequiresLSPSolution" />
+            bool RequiresLSPSolution { get; }
 
             /// <inheritdoc cref="IRequestHandler.MutatesSolutionState" />
-            public readonly bool MutatesSolutionState;
-
-            /// <inheritdoc cref="IRequestHandler.RequiresLSPSolution" />
-            public readonly bool RequiresLSPSolution;
+            bool MutatesSolutionState { get; }
 
             /// <inheritdoc cref="RequestContext.ClientName" />
-            public readonly string? ClientName;
-            public readonly string MethodName;
+            string? ClientName { get; }
 
-            /// <inheritdoc cref="RequestContext.ClientCapabilities" />
-            public readonly ClientCapabilities ClientCapabilities;
+            string MethodName { get; }
 
             /// <summary>
             /// The document identifier that will be used to find the solution and document for this request. This comes from the <see cref="TextDocumentIdentifier"/> returned from the handler itself via a call to <see cref="IRequestHandler{RequestType, ResponseType}.GetTextDocumentIdentifier(RequestType)"/>.
             /// </summary>
-            public readonly TextDocumentIdentifier? TextDocument;
+            TextDocumentIdentifier? TextDocument { get; }
+
+            /// <inheritdoc cref="RequestContext.ClientCapabilities" />
+            ClientCapabilities ClientCapabilities { get; }
 
             /// <summary>
-            /// A cancellation token that will cancel the handing of this request. The request could also be cancelled by the queue shutting down.
+            /// <see cref="CorrelationManager.ActivityId"/> used to properly correlate this work with the loghub
+            /// tracing/logging subsystem.
             /// </summary>
-            public readonly CancellationToken CancellationToken;
+            Guid ActivityId { get; }
+
+            RequestMetrics Metrics { get; }
+        }
+
+        private class QueueItem<TRequestType, TResponseType> : IQueueItem
+        {
+            private readonly ILspLogger _logger;
+
+            private readonly TRequestType _request;
+            private readonly IRequestHandler<TRequestType, TResponseType> _handler;
 
             /// <summary>
-            /// An action to be called when the queue fails to begin execution of this work item.
+            /// A task completion source representing the result of this queue item's work.
+            /// This is the task that the client is waiting on.
             /// </summary>
-            public readonly Action<Exception> HandleQueueFailure;
+            private readonly TaskCompletionSource<TResponseType?> _completionSource;
+
+            public bool RequiresLSPSolution { get; }
+
+            public bool MutatesSolutionState { get; }
+
+            public string? ClientName { get; }
 
-            public readonly RequestMetrics Metrics;
+            public string MethodName { get; }
+
+            public TextDocumentIdentifier? TextDocument { get; }
+
+            public ClientCapabilities ClientCapabilities { get; }
+
+            public Guid ActivityId { get; }
+
+            public RequestMetrics Metrics { get; }
 
             public QueueItem(
                 bool mutatesSolutionState,
@@ -66,17 +86,22 @@ public QueueItem(
                 string? clientName,
                 string methodName,
                 TextDocumentIdentifier? textDocument,
+                TRequestType request,
+                IRequestHandler<TRequestType, TResponseType> handler,
                 Guid activityId,
                 ILspLogger logger,
                 RequestTelemetryLogger telemetryLogger,
-                Action<Exception> handleQueueFailure,
-                Func<RequestContext?, CancellationToken, Task> callbackAsync,
                 CancellationToken cancellationToken)
             {
+                _completionSource = new TaskCompletionSource<TResponseType?>();
+                // Set the tcs state to cancelled if the token gets cancelled outside of our callback (for example the server shutting down).
+                cancellationToken.Register(() => _completionSource.TrySetCanceled(cancellationToken));
+
                 Metrics = new RequestMetrics(methodName, telemetryLogger);
 
-                _callbackAsync = callbackAsync;
+                _handler = handler;
                 _logger = logger;
+                _request = request;
 
                 ActivityId = activityId;
                 MutatesSolutionState = mutatesSolutionState;
@@ -85,12 +110,43 @@ public QueueItem(
                 ClientName = clientName;
                 MethodName = methodName;
                 TextDocument = textDocument;
-                HandleQueueFailure = handleQueueFailure;
-                CancellationToken = cancellationToken;
+            }
+
+            public static (IQueueItem, Task<TResponseType?>) Create(
+                bool mutatesSolutionState,
+                bool requiresLSPSolution,
+                ClientCapabilities clientCapabilities,
+                string? clientName,
+                string methodName,
+                TextDocumentIdentifier? textDocument,
+                TRequestType request,
+                IRequestHandler<TRequestType, TResponseType> handler,
+                Guid activityId,
+                ILspLogger logger,
+                RequestTelemetryLogger telemetryLogger,
+                CancellationToken cancellationToken)
+            {
+                var queueItem = new QueueItem<TRequestType, TResponseType>(
+                    mutatesSolutionState,
+                    requiresLSPSolution,
+                    clientCapabilities,
+                    clientName,
+                    methodName,
+                    textDocument,
+                    request,
+                    handler,
+                    activityId,
+                    logger,
+                    telemetryLogger,
+                    cancellationToken);
+
+                return (queueItem, queueItem._completionSource.Task);
             }
 
             /// <summary>
-            /// Processes the queued request. Exceptions that occur will be sent back to the requesting client, then re-thrown
+            /// Processes the queued request. Exceptions will be sent to the task completion source
+            /// representing the task that the client is waiting for, then re-thrown so that
+            /// the queue can correctly handle them depending on the type of request.
             /// </summary>
             public async Task CallbackAsync(RequestContext? context, CancellationToken cancellationToken)
             {
@@ -99,25 +155,54 @@ public async Task CallbackAsync(RequestContext? context, CancellationToken cance
                 _logger.TraceStart($"{MethodName} - Roslyn");
                 try
                 {
-                    await _callbackAsync(context, cancellationToken).ConfigureAwait(false);
-                    this.Metrics.RecordSuccess();
+                    cancellationToken.ThrowIfCancellationRequested();
+
+                    TResponseType? result;
+                    if (context == null)
+                    {
+                        // If we weren't able to get a corresponding context for this request (for example, we
+                        // couldn't map a doc request to a particular Document, or we couldn't find an appropriate
+                        // Workspace for a global operation), then just immediately complete the request with a
+                        // 'null' response.  Note: the lsp spec was checked to ensure that 'null' is valid for all
+                        // the requests this could happen for.  However, this assumption may not hold in the future.
+                        // If that turns out to be the case, we could defer to the individual handler to decide
+                        // what to do.
+                        _logger.TraceWarning($"Could not get request context for {MethodName}");
+                        this.Metrics.RecordFailure();
+                        result = default;
+                    }
+                    else
+                    {
+                        result = await _handler.HandleRequestAsync(_request, context.Value, cancellationToken).ConfigureAwait(false);
+                        this.Metrics.RecordSuccess();
+                    }
+
+                    _completionSource.TrySetResult(result);
                 }
-                catch (OperationCanceledException)
+                catch (OperationCanceledException ex)
                 {
+                    // Record logs + metrics on cancellation.
                     _logger.TraceInformation($"{MethodName} - Canceled");
                     this.Metrics.RecordCancellation();
-                    throw;
+
+                    _completionSource.TrySetCanceled(ex.CancellationToken);
                 }
                 catch (Exception ex)
                 {
+                    // Record logs and metrics on the exception.
                     _logger.TraceException(ex);
                     this.Metrics.RecordFailure();
-                    throw;
+
+                    _completionSource.TrySetException(ex);
                 }
                 finally
                 {
                     _logger.TraceStop($"{MethodName} - Roslyn");
                 }
+
+                // Return the result of this completion source to the caller
+                // so it can decide how to handle the result / exception.
+                await _completionSource.Task.ConfigureAwait(false);
             }
         }
     }