Move cancellation token creation to MessageScheduler

Change-Id: I1def99245a7f063c5613cb9b82d75b686aff9ce7
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/391203
Commit-Queue: Keerti Parthasarathy <[email protected]>
Reviewed-by: Brian Wilkerson <[email protected]>

Author: keertip

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

@@ -417,12 +417,12 @@ abstract class ServerStateMessageHandler {
   final AnalysisServer server;
   final Map<Method, MessageHandler<Object?, Object?, AnalysisServer>>
       messageHandlers = {};
-  final CancelRequestHandler _cancelHandler;
+  final CancelRequestHandler cancelHandler;
   final NotCancelableToken _notCancelableToken = NotCancelableToken();
 
   ServerStateMessageHandler(this.server)
-      : _cancelHandler = CancelRequestHandler(server) {
-    registerHandler(_cancelHandler);
+      : cancelHandler = CancelRequestHandler(server) {
+    registerHandler(cancelHandler);
   }
 
   /// Handle the given [message]. If the [message] is a [RequestMessage], then the
@@ -446,7 +446,7 @@ abstract class ServerStateMessageHandler {
     // server), create a new cancellation token that will allow us to cancel
     // this request if requested. This saves some processing but the handler
     // will need to specifically check the token after `await`s.
-    cancellationToken ??= _cancelHandler.createToken(message);
+    cancellationToken ??= cancelHandler.createToken(message);
     try {
       var result =
           await handler.handleMessage(message, messageInfo, cancellationToken);
@@ -458,7 +458,7 @@ abstract class ServerStateMessageHandler {
       await Future.delayed(Duration.zero);
       return cancellationToken.isCancellationRequested ? cancelled() : result;
     } finally {
-      _cancelHandler.clearToken(message);
+      cancelHandler.clearToken(message);
     }
   }
 

pkg/analysis_server/lib/src/lsp/lsp_analysis_server.dart

@@ -477,7 +477,7 @@ class LspAnalysisServer extends AnalysisServer {
   }
 
   /// Handle a [message] that was read from the communication channel.
-  void handleMessage(Message message) {
+  void handleMessage(Message message, {CancelableToken? cancellationToken}) {
     var startTime = DateTime.now();
     performance.logRequestTiming(message.clientRequestTime);
     runZonedGuarded(() async {
@@ -506,7 +506,8 @@ class LspAnalysisServer extends AnalysisServer {
             if (message is RequestMessage) {
               analyticsManager.startedRequestMessage(
                   request: message, startTime: startTime);
-              await _handleRequestMessage(message, messageInfo);
+              await _handleRequestMessage(message, messageInfo,
+                  cancellationToken: cancellationToken);
             } else if (message is NotificationMessage) {
               await _handleNotificationMessage(message, messageInfo);
               analyticsManager.handledNotificationMessage(
@@ -760,7 +761,14 @@ class LspAnalysisServer extends AnalysisServer {
   }
 
   void scheduleMessage(Message message) {
-    messageScheduler.add(LspMessage(message: message));
+    CancelableToken? cancellationToken;
+    // Create a  cancellation token that will allow us to cancel
+    // this request if requested.
+    if (message is RequestMessage) {
+      cancellationToken = messageHandler.cancelHandler.createToken(message);
+    }
+    messageScheduler.add(
+        LspMessage(message: message, cancellationToken: cancellationToken));
     messageScheduler.notify();
   }
 
@@ -1062,10 +1070,10 @@ class LspAnalysisServer extends AnalysisServer {
   }
 
   Future<void> _handleRequestMessage(
-    RequestMessage message,
-    MessageInfo messageInfo,
-  ) async {
-    var result = await messageHandler.handleMessage(message, messageInfo);
+      RequestMessage message, MessageInfo messageInfo,
+      {CancelableToken? cancellationToken}) async {
+    var result = await messageHandler.handleMessage(message, messageInfo,
+        cancellationToken: cancellationToken);
     result.ifError((error) => sendErrorResponse(message, error));
     result.ifResult(
       (result) => sendResponse(ResponseMessage(

pkg/analysis_server/lib/src/server/message_scheduler.dart

@@ -11,6 +11,7 @@ import 'package:analysis_server/src/analysis_server.dart';
 import 'package:analysis_server/src/legacy_analysis_server.dart';
 import 'package:analysis_server/src/lsp/lsp_analysis_server.dart';
 import 'package:analyzer/src/util/performance/operation_performance.dart';
+import 'package:analyzer/src/utilities/cancellation.dart';
 import 'package:meta/meta.dart';
 
 /// Represents a message from DTD (Dart Tooling Daemon).
@@ -35,8 +36,9 @@ final class LegacyMessage extends MessageObject {
 /// Represents a message in the LSP protocol format.
 final class LspMessage extends MessageObject {
   final lsp.Message message;
+  CancelableToken? cancellationToken;
 
-  LspMessage({required this.message});
+  LspMessage({required this.message, this.cancellationToken});
 }
 
 /// Represents a message from a client, can be an IDE, DTD etc.
@@ -75,7 +77,8 @@ final class MessageScheduler {
     switch (message) {
       case LspMessage():
         var lspMessage = message.message;
-        (server as LspAnalysisServer).handleMessage(lspMessage);
+        (server as LspAnalysisServer).handleMessage(lspMessage,
+            cancellationToken: message.cancellationToken);
       case LegacyMessage():
         var request = message.request;
         (server as LegacyAnalysisServer).handleRequest(request);