[anaysis_server] Cancel in-progress LSP completion requests if another arrives

Change-Id: Idf49b70750a3427529f71e952aa036fae4c7c6cb
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/252422
Commit-Queue: Brian Wilkerson <brianwilkerson@google.com>
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>

pkg/analysis_server/lib/src/lsp/handlers/commands/perform_refactor.dart

@@ -33,13 +33,10 @@ class PerformRefactorCommandHandler extends AbstractRefactorCommandHandler {
       final refactoring = await getRefactoring(
           RefactoringKind(kind), result, offset, length, options);
       return refactoring.mapResult((refactoring) async {
-        // If the token we were given is not cancellable, replace it with one that
-        // is for the rest of this request, as a future refactor may need to cancel
-        // this request.
-        // The original token should be kept and also checked for cancellation.
-        final cancelableToken = cancellationToken is CancelableToken
-            ? cancellationToken
-            : CancelableToken();
+        // If the token we were given is not cancelable, wrap it with one that
+        // is for the rest of this request as a future refactor may need to
+        // cancel this request.
+        final cancelableToken = cancellationToken.asCancelable();
         manager.begin(cancelableToken);
 
         try {
@@ -54,15 +51,13 @@ class PerformRefactorCommandHandler extends AbstractRefactorCommandHandler {
             return success(null);
           }
 
-          if (cancellationToken.isCancellationRequested ||
-              cancelableToken.isCancellationRequested) {
+          if (cancelableToken.isCancellationRequested) {
             return error(ErrorCodes.RequestCancelled, 'Request was cancelled');
           }
 
           final change = await refactoring.createChange();
 
-          if (cancellationToken.isCancellationRequested ||
-              cancelableToken.isCancellationRequested) {
+          if (cancelableToken.isCancellationRequested) {
             return error(ErrorCodes.RequestCancelled, 'Request was cancelled');
           }
 

pkg/analysis_server/lib/src/lsp/handlers/commands/validate_refactor.dart

@@ -42,13 +42,10 @@ class ValidateRefactorCommandHandler extends AbstractRefactorCommandHandler {
       final refactoring = await getRefactoring(
           RefactoringKind(kind), result, offset, length, options);
       return refactoring.mapResult((refactoring) async {
-        // If the token we were given is not cancellable, replace it with one that
-        // is for the rest of this request, as a future refactor may need to cancel
-        // this request.
-        // The original token should be kept and also checked for cancellation.
-        final cancelableToken = cancellationToken is CancelableToken
-            ? cancellationToken
-            : CancelableToken();
+        // If the token we were given is not cancelable, wrap it with one that
+        // is for the rest of this request as a future refactor may need to
+        // cancel this request.
+        final cancelableToken = cancellationToken.asCancelable();
         manager.begin(cancelableToken);
 
         try {

pkg/analysis_server/lib/src/lsp/handlers/handler_completion.dart

@@ -48,6 +48,15 @@ class CompletionHandler extends MessageHandler<CompletionParams, CompletionList>
   /// debugging).
   late final Duration completionBudgetDuration;
 
+  /// A cancellation token for the previous completion request.
+  ///
+  /// A new completion request will cancel the previous request. We do not allow
+  /// concurrent completion requests.
+  ///
+  /// `null` if there is no previous request. It the previous request has
+  /// already completed, cancelling this token will not do anything.
+  CancelableToken? previousRequestCancellationToken;
+
   CompletionHandler(super.server, LspInitializationOptions options)
       : suggestFromUnimportedLibraries = options.suggestFromUnimportedLibraries,
         previewNotImportedCompletions = options.previewNotImportedCompletions {
@@ -73,6 +82,12 @@ class CompletionHandler extends MessageHandler<CompletionParams, CompletionList>
       return serverNotInitializedError;
     }
 
+    // Cancel any existing in-progress completion request in case the client did
+    // not do it explicitly, because the results will not be useful and it may
+    // delay processing this one.
+    previousRequestCancellationToken?.cancel();
+    previousRequestCancellationToken = token.asCancelable();
+
     final requestLatency = message.timeSinceRequest;
     final triggerCharacter = params.context?.triggerCharacter;
     final pos = params.position;

pkg/analysis_server/test/lsp/completion_dart_test.dart

@@ -735,6 +735,27 @@ final a = Stri^
         request, throwsA(isResponseError(ErrorCodes.InvalidParams)));
   }
 
+  Future<void> test_concurrentRequestsCancellation() async {
+    // We expect a new completion request to cancel any in-flight request so
+    // send multiple without awaiting, then check only the last one completes.
+    final content = '^';
+
+    await initialize();
+    await openFile(mainFileUri, withoutMarkers(content));
+    final position = positionFromMarker(content);
+    final responseFutures = [
+      getCompletion(mainFileUri, position),
+      getCompletion(mainFileUri, position),
+      getCompletion(mainFileUri, position),
+    ];
+    expect(responseFutures[0],
+        throwsA(isResponseError(ErrorCodes.RequestCancelled)));
+    expect(responseFutures[1],
+        throwsA(isResponseError(ErrorCodes.RequestCancelled)));
+    final results = await responseFutures[2];
+    expect(results, isNotEmpty);
+  }
+
   Future<void> test_filterTextNotIncludeAdditionalText() async {
     // Some completions (eg. overrides) have additional text that is not part
     // of the label. That text should _not_ appear in filterText as it will

pkg/analyzer/lib/src/utilities/cancellation.dart

@@ -23,3 +23,37 @@ class NotCancelableToken extends CancellationToken {
   @override
   bool get isCancellationRequested => false;
 }
+
+/// A cancellable wrapper over another cancellation token.
+///
+/// This token will be considered cancelled if either it is itself cancelled,
+/// or if [child] is cancelled.
+///
+/// Cancelling this token will also cancel [child] if it is a cancelable
+/// token.
+class _WrappedCancelableToken extends CancelableToken {
+  final CancellationToken _child;
+
+  _WrappedCancelableToken(this._child);
+
+  @override
+  bool get isCancellationRequested =>
+      super.isCancellationRequested || _child.isCancellationRequested;
+
+  @override
+  void cancel() {
+    super.cancel();
+    final child = _child;
+    if (child is CancelableToken) {
+      child.cancel();
+    }
+  }
+}
+
+extension CancellationTokenExtension on CancellationToken {
+  /// Wraps this token to make it cancelable if it is not already.
+  CancelableToken asCancelable() {
+    final token = this;
+    return token is CancelableToken ? token : _WrappedCancelableToken(token);
+  }
+}