[analysis_server] Move support for locking requests to the scheduler

This removes the original `lockRequestsWhile` functionality of the LSP server which worked by pausing reading from stdin to instead pause the processing of messages in the scheduler.

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

Author: DanTup

pkg/analysis_server/lib/src/analysis_server.dart

@@ -273,6 +273,13 @@ abstract class AnalysisServer {
   /// A [TimingByteStore] that records timings for reads from the byte store.
   TimingByteStore? _timingByteStore;
 
+  /// Whether notifications caused by analysis should be suppressed.
+  ///
+  /// This is used when an operation is temporarily modifying overlays and does
+  /// not want the client to be notified of any analysis happening on the
+  /// temporary content.
+  bool suppressAnalysisResults = false;
+
   AnalysisServer(
     this.options,
     this.sdkManager,
@@ -884,6 +891,39 @@ abstract class AnalysisServer {
   /// given [path] was changed - added, updated, or removed.
   void notifyFlutterWidgetDescriptions(String path) {}
 
+  /// Prevents the scheduler from processing new messages until [operation]
+  /// completes.
+  ///
+  /// This can be used to obtain analysis results/resolved units consistent with
+  /// the state of a file at the time this method was called, preventing
+  /// changes by incoming file modifications.
+  ///
+  /// The contents of [operation] should be kept as short as possible and since
+  /// cancellation requests will also be blocked for the duration of this
+  /// operation, handlers should generally check the cancellation flag
+  /// immediately after this function returns.
+  Future<T> pauseSchedulerWhile<T>(FutureOr<T> Function() operation) async {
+    // TODO(dantup): Prevent this method from locking responses from the client
+    //  because this can lead to deadlocks if called during initialization where
+    //  the server may wait for something (configuration) from the client. This
+    //  might fit in with potential upcoming scheduler changes.
+    //
+    // This is currently used by Completion+FixAll (which are less likely, but
+    // possible to be called during init).
+    //
+    // https://github.com/dart-lang/sdk/issues/56311#issuecomment-2250089185
+
+    messageScheduler.pause();
+    try {
+      // `await` here is important to ensure `finally` doesn't execute until
+      // `operation()` completes (`whenComplete` is not available on
+      // `FutureOr`).
+      return await operation();
+    } finally {
+      messageScheduler.resume();
+    }
+  }
+
   /// Read all files, resolve all URIs, and perform required analysis in
   /// all current analysis drivers.
   Future<void> reanalyze() async {
@@ -1087,6 +1127,10 @@ abstract class CommonServerContextManagerCallbacks
   @override
   @mustCallSuper
   void handleFileResult(FileResult result) {
+    if (analysisServer.suppressAnalysisResults) {
+      return;
+    }
+
     var path = result.path;
     filesToFlush.add(path);
 

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

@@ -5,6 +5,7 @@
 import 'dart:async';
 
 import 'package:analysis_server/lsp_protocol/protocol.dart';
+import 'package:analysis_server/src/analysis_server.dart';
 import 'package:analysis_server/src/lsp/constants.dart';
 import 'package:analysis_server/src/lsp/error_or.dart';
 import 'package:analysis_server/src/lsp/handlers/commands/simple_edit_handler.dart';
@@ -89,22 +90,22 @@ class FixAllCommandHandler extends SimpleEditCommandHandler<LspAnalysisServer> {
 
 /// Computes edits for iterative fix-all using temporary overlays.
 class _FixAllOperation extends TemporaryOverlayOperation
-    with HandlerHelperMixin<LspAnalysisServer> {
+    with HandlerHelperMixin<AnalysisServer> {
   final MessageInfo message;
   final CancellationToken cancellationToken;
   final String path;
   final bool autoTriggered;
 
   _FixAllOperation({
-    required LspAnalysisServer server,
+    required AnalysisServer server,
     required this.message,
     required this.path,
     required this.cancellationToken,
     required this.autoTriggered,
   }) : super(server);
 
   Future<ErrorOr<WorkspaceEdit?>> computeEdits() async {
-    return await lockRequestsWithTemporaryOverlays(_computeEditsImpl);
+    return await pauseSchedulerWithTemporaryOverlays(_computeEditsImpl);
   }
 
   Future<ErrorOr<WorkspaceEdit?>> _computeEditsImpl() async {

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

@@ -119,7 +119,7 @@ class CompletionHandler
     // unit and LineInfo.
     late ErrorOr<LineInfo> lineInfo;
     late ErrorOr<ResolvedUnitResult> unit;
-    await server.lockRequestsWhile(() async {
+    await server.pauseSchedulerWhile(() async {
       unit = await path.mapResult(requireResolvedUnit);
       lineInfo = await unit.map(
         // If we don't have a unit, we can still try to obtain the line info from

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

@@ -106,9 +106,6 @@ class LspAnalysisServer extends AnalysisServer {
   @visibleForTesting
   int contextBuilds = 0;
 
-  /// The subscription to the stream of incoming messages from the client.
-  late final StreamSubscription<void> _channelSubscription;
-
   /// An optional manager to handle file systems which may not always be
   /// available.
   final DetachableFileSystemManager? detachableFileSystemManager;
@@ -117,13 +114,6 @@ class LspAnalysisServer extends AnalysisServer {
   /// `sendStatusNotification` was invoked.
   bool wasAnalyzing = false;
 
-  /// Whether notifications caused by analysis should be suppressed.
-  ///
-  /// This is used when an operation is temporarily modifying overlays and does
-  /// not want the client to be notified of any analysis happening on the
-  /// temporary content.
-  bool suppressAnalysisResults = false;
-
   /// Tracks files that have non-empty diagnostics on the client.
   ///
   /// This is an optimization to avoid sending empty diagnostics when they are
@@ -181,11 +171,7 @@ class LspAnalysisServer extends AnalysisServer {
         .listen(handleAnalysisEvent);
     analysisDriverScheduler.start();
 
-    _channelSubscription = channel.listen(
-      scheduleMessage,
-      onDone: done,
-      onError: socketError,
-    );
+    channel.listen(scheduleMessage, onDone: done, onError: socketError);
 
     if (AnalysisServer.supportsPlugins) {
       _pluginChangeSubscription = pluginManager.pluginsChanged.listen(
@@ -572,45 +558,6 @@ class LspAnalysisServer extends AnalysisServer {
     }, socketError);
   }
 
-  /// Locks the server from processing incoming messages until [operation]
-  /// completes.
-  ///
-  /// This can be used to obtain analysis results/resolved units consistent with
-  /// the state of a file at the time this method was called, preventing
-  /// changes by incoming file modifications.
-  ///
-  /// The contents of [operation] should be kept as short as possible and since
-  /// cancellation requests will also be blocked for the duration of this
-  /// operation, handles should generally check the cancellation flag
-  /// immediately after this function returns.
-  Future<T> lockRequestsWhile<T>(FutureOr<T> Function() operation) async {
-    // TODO(dantup): Prevent this method from locking responses from the client
-    //  because this can lead to deadlocks if called during initialization where
-    //  the server may wait for something (configuration) from the client. This
-    //  might fit in with potential upcoming scheduler changes.
-    //
-    // This is currently used by Completion+FixAll (which are less likely, but
-    // possible to be called during init).
-    //
-    // https://github.com/dart-lang/sdk/issues/56311#issuecomment-2250089185
-    var completer = Completer<void>();
-
-    // Pause handling incoming messages until `operation` completes.
-    //
-    // If this method is called multiple times, the pauses will stack, meaning
-    // the subscription will not resume until all operations complete.
-    _channelSubscription.pause(completer.future);
-
-    try {
-      // `await` here is important to ensure `finally` doesn't execute until
-      // `operation()` completes (`whenComplete` is not available on
-      // `FutureOr`).
-      return await operation();
-    } finally {
-      completer.complete();
-    }
-  }
-
   /// Logs the error on the client using window/logMessage.
   void logErrorToClient(String message) {
     channel.sendNotification(
@@ -1304,15 +1251,6 @@ class LspServerContextManagerCallbacks
     }
   }
 
-  @override
-  void handleFileResult(FileResult result) {
-    if (analysisServer.suppressAnalysisResults) {
-      return;
-    }
-
-    super.handleFileResult(result);
-  }
-
   @override
   void handleResolvedUnitResult(ResolvedUnitResult result) {
     var path = result.path;

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

@@ -4,8 +4,8 @@
 
 import 'dart:async';
 
+import 'package:analysis_server/src/analysis_server.dart';
 import 'package:analysis_server/src/context_manager.dart';
-import 'package:analysis_server/src/lsp/lsp_analysis_server.dart';
 import 'package:analysis_server/src/protocol_server.dart';
 import 'package:analyzer/dart/analysis/analysis_context.dart';
 import 'package:analyzer/file_system/overlay_file_system.dart';
@@ -18,7 +18,7 @@ import 'package:analyzer/src/dart/analysis/driver.dart';
 /// need to be merged together (to be mappable to LSP document changes) and then
 /// reverted to allow the client to apply the change.
 abstract class TemporaryOverlayOperation {
-  final LspAnalysisServer server;
+  final AnalysisServer server;
   final ContextManager contextManager;
   final OverlayResourceProvider resourceProvider;
 
@@ -78,16 +78,16 @@ abstract class TemporaryOverlayOperation {
   }
 
   /// Locks the server from processing incoming messages until [operation]
-  /// completes just like [LspAnalysisServer.lockRequestsWhile] but
+  /// completes just like [AnalysisServer.pauseSchedulerWhile] but
   /// additionally provides a function for writing temporary overlays that will
   /// be reverted when the operation completes.
   ///
   /// Additionally, sending diagnostics, outlines, etc. are suppressed by the
   /// temporary overlays and re-enabled after the overlays are restored.
-  Future<T> lockRequestsWithTemporaryOverlays<T>(
+  Future<T> pauseSchedulerWithTemporaryOverlays<T>(
     Future<T> Function() operation,
   ) {
-    return server.lockRequestsWhile(() async {
+    return server.pauseSchedulerWhile(() async {
       // Wait for any in-progress analysis to complete before we start
       // suppressing analysis results.
       server.contextManager.pauseWatchers();

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

@@ -45,6 +45,11 @@ final class MessageScheduler {
   /// Whether the [MessageScheduler] is currently processing messages.
   bool _isProcessing = false;
 
+  /// The number of times [pause] has been called without matching [resume]s.
+  ///
+  /// If zero, the queue is not paused.
+  int _pauseCount = 0;
+
   /// The completer used to indicate that message handling has been completed.
   Completer<void> _completer = Completer();
 
@@ -58,6 +63,9 @@ final class MessageScheduler {
   /// around.
   MessageScheduler({required this.listener});
 
+  /// Whether the queue is currently paused.
+  bool get isPaused => _pauseCount > 0;
+
   /// Add the [message] to the end of the pending messages queue.
   ///
   /// Some incoming messages are handled immediately rather than being added to
@@ -177,12 +185,28 @@ final class MessageScheduler {
     }
   }
 
+  /// Pauses processing messages.
+  ///
+  /// Any messages that are already being processed will continue until they
+  /// complete, but no new messages will be processed.
+  ///
+  /// If this method is called multiple times, [resume] will need to be called
+  /// an equal number of times for processing to continue.
+  void pause() {
+    _pauseCount++;
+    listener?.pauseProcessingMessages(_pauseCount);
+  }
+
   /// Dispatch the first message in the queue to be executed.
   void processMessages() async {
+    if (isPaused) {
+      return;
+    }
+
     _isProcessing = true;
     listener?.startProcessingMessages();
     try {
-      while (_pendingMessages.isNotEmpty) {
+      while (_pendingMessages.isNotEmpty && !isPaused) {
         var currentMessage = _pendingMessages.removeFirst();
         _activeMessages.addLast(currentMessage);
         listener?.addActiveMessage(currentMessage);
@@ -249,6 +273,21 @@ final class MessageScheduler {
     listener?.endProcessingMessages();
   }
 
+  /// Resumes processing messages.
+  void resume() {
+    if (!isPaused) {
+      throw StateError('Cannot resume if not paused');
+    }
+    _pauseCount--;
+    listener?.resumeProcessingMessages(_pauseCount);
+    if (!isPaused && !_isProcessing) {
+      // Process on the next tick so that the caller to resume() doesn't get
+      // messages in the queue attributed to their time (or run before they
+      // complete).
+      Future.delayed(Duration.zero, processMessages);
+    }
+  }
+
   /// Returns the parameters of a cancellation [message].
   lsp.CancelParams? _getCancelParams(lsp.NotificationMessage message) {
     try {
@@ -477,6 +516,14 @@ abstract class MessageSchedulerListener {
   /// This implies that the message was active and wasn't cancelled.
   void messageCompleted(ScheduledMessage message);
 
+  /// Report that the pause counter was increased to [newPauseCount], and that
+  /// processing will be paused.
+  void pauseProcessingMessages(int newPauseCount);
+
+  /// Report that the pause counter was decreased to [newPauseCount] which, if
+  /// zero, indicates processing will resume.
+  void resumeProcessingMessages(int newPauseCount);
+
   /// Report that the loop that processes messages has started to run.
   void startProcessingMessages();
 }

pkg/analysis_server/lib/src/scheduler/scheduler_tracking_listener.dart

@@ -113,23 +113,26 @@ class SchedulerTrackingListener extends MessageSchedulerListener {
 
   @override
   void cancelActiveMessage(ScheduledMessage message) {
-    var messageData = _messageDataMap.remove(message);
+    var messageData = _messageDataMap[message];
     if (messageData == null) {
       return;
     }
     messageData.completeTime = _now;
     messageData.wasCancelled = true;
-    _activeMessageCount--;
-    _reportMessageData(messageData);
+    // Don't decrement counts or report message data yet because
+    // cancelled messages still complete and call messageCompleted().
   }
 
   @override
   void cancelPendingMessage(ScheduledMessage message) {
-    var messageData = _messageDataMap.remove(message)!;
+    var messageData = _messageDataMap[message];
+    if (messageData == null) {
+      return;
+    }
     messageData.completeTime = _now;
     messageData.wasCancelled = true;
-    _pendingMessageCount--;
-    _reportMessageData(messageData);
+    // Don't decrement counts or report message data yet because
+    // cancelled messages still complete and call messageCompleted().
   }
 
   /// Report that the loop that processes messages has stopped running.
@@ -148,6 +151,16 @@ class SchedulerTrackingListener extends MessageSchedulerListener {
     _reportMessageData(messageData);
   }
 
+  @override
+  void pauseProcessingMessages(int newPauseCount) {
+    // TODO(dantup): Consider tracking the pause start time if newPauseCount=1.
+  }
+
+  @override
+  void resumeProcessingMessages(int newPauseCount) {
+    // TODO(dantup): Consider recording the pause duration if newPauseCount=0.
+  }
+
   @override
   void startProcessingMessages() {
     processingStartTime = _now;