[coverage] Proactively collect isolates as they pause

pkgs/coverage/CHANGELOG.md

@@ -1,3 +1,7 @@
+## 1.10.0-wip
+
+- Fix bug where tests involving multiple isolates never finish (#520).
+
 ## 1.9.2
 
 - Fix repository link in pubspec.

pkgs/coverage/lib/src/collect.dart

@@ -8,6 +8,7 @@ import 'dart:io';
 import 'package:vm_service/vm_service.dart';
 
 import 'hitmap.dart';
+import 'isolate_paused_listener.dart';
 import 'util.dart';
 
 const _retryInterval = Duration(milliseconds: 200);
@@ -25,8 +26,8 @@ const _debugTokenPositions = bool.fromEnvironment('DEBUG_COVERAGE');
 /// If [resume] is true, all isolates will be resumed once coverage collection
 /// is complete.
 ///
-/// If [waitPaused] is true, collection will not begin until all isolates are
-/// in the paused state.
+/// If [waitPaused] is true, collection will not begin for an isolate until it
+/// is in the paused state.
 ///
 /// If [includeDart] is true, code coverage for core `dart:*` libraries will be
 /// collected.
@@ -93,14 +94,17 @@ Future<Map<String, dynamic>> collect(Uri serviceUri, bool resume,
   }
 
   try {
-    if (waitPaused) {
-      await _waitIsolatesPaused(service, timeout: timeout);
-    }
-
-    return await _getAllCoverage(service, includeDart, functionCoverage,
-        branchCoverage, scopedOutput, isolateIds, coverableLineCache);
+    return await _getAllCoverage(
+        service,
+        includeDart,
+        functionCoverage,
+        branchCoverage,
+        scopedOutput,
+        isolateIds,
+        coverableLineCache,
+        waitPaused);
   } finally {
-    if (resume) {
+    if (resume && !waitPaused) {
       await _resumeIsolates(service);
     }
     // The signature changed in vm_service version 6.0.0.
@@ -114,11 +118,10 @@ Future<Map<String, dynamic>> _getAllCoverage(
     bool includeDart,
     bool functionCoverage,
     bool branchCoverage,
-    Set<String>? scopedOutput,
+    Set<String> scopedOutput,
     Set<String>? isolateIds,
-    Map<String, Set<int>>? coverableLineCache) async {
-  scopedOutput ??= <String>{};
-  final vm = await service.getVM();
+    Map<String, Set<int>>? coverableLineCache,
+    bool waitPaused) async {
   final allCoverage = <Map<String, dynamic>>[];
 
   final sourceReportKinds = [
@@ -133,11 +136,15 @@ Future<Map<String, dynamic>> _getAllCoverage(
   // group, otherwise we'll double count the hits.
   final coveredIsolateGroups = <String>{};
 
-  for (var isolateRef in vm.isolates!) {
-    if (isolateIds != null && !isolateIds.contains(isolateRef.id)) continue;
+  Future<void> collectIsolate(IsolateRef isolateRef) async {
+    if (!(isolateIds?.contains(isolateRef.id) ?? true)) return;
+
+    // coveredIsolateGroups is only relevant for the !waitPaused flow. The
+    // waitPaused flow achieves the same once-per-group behavior using the
+    // isLastIsolateInGroup flag.
     final isolateGroupId = isolateRef.isolateGroupId;
     if (isolateGroupId != null) {
-      if (coveredIsolateGroups.contains(isolateGroupId)) continue;
+      if (coveredIsolateGroups.contains(isolateGroupId)) return;
       coveredIsolateGroups.add(isolateGroupId);
     }
 
@@ -154,8 +161,9 @@ Future<Map<String, dynamic>> _getAllCoverage(
         librariesAlreadyCompiled: librariesAlreadyCompiled,
       );
     } on SentinelException {
-      continue;
+      return;
     }
+
     final coverage = await _processSourceReport(
         service,
         isolateRef,
@@ -166,6 +174,20 @@ Future<Map<String, dynamic>> _getAllCoverage(
         scopedOutput);
     allCoverage.addAll(coverage);
   }
+
+  if (waitPaused) {
+    await IsolatePausedListener(service,
+        (IsolateRef isolateRef, bool isLastIsolateInGroup) async {
+      if (isLastIsolateInGroup) {
+        await collectIsolate(isolateRef);
+      }
+    }).waitUntilAllExited();
+  } else {
+    for (final isolateRef in await getAllIsolates(service)) {
+      await collectIsolate(isolateRef);
+    }
+  }
+
   return <String, dynamic>{'type': 'CodeCoverage', 'coverage': allCoverage};
 }
 
@@ -190,29 +212,6 @@ Future _resumeIsolates(VmService service) async {
   }
 }
 
-Future _waitIsolatesPaused(VmService service, {Duration? timeout}) async {
-  final pauseEvents = <String>{
-    EventKind.kPauseStart,
-    EventKind.kPauseException,
-    EventKind.kPauseExit,
-    EventKind.kPauseInterrupted,
-    EventKind.kPauseBreakpoint
-  };
-
-  Future allPaused() async {
-    final vm = await service.getVM();
-    if (vm.isolates!.isEmpty) throw StateError('No isolates.');
-    for (var isolateRef in vm.isolates!) {
-      final isolate = await service.getIsolate(isolateRef.id!);
-      if (!pauseEvents.contains(isolate.pauseEvent!.kind)) {
-        throw StateError('Unpaused isolates remaining.');
-      }
-    }
-  }
-
-  return retry(allPaused, _retryInterval, timeout: timeout);
-}
-
 /// Returns the line number to which the specified token position maps.
 ///
 /// Performs a binary search within the script's token position table to locate
@@ -278,6 +277,7 @@ Future<List<Map<String, dynamic>>> _processSourceReport(
       return;
     }
     final funcName = await _getFuncName(service, isolateRef, func);
+    // TODO(liama): Is this still necessary, or is location.line valid?
     final tokenPos = location.tokenPos!;
     final line = _getLineFromTokenPos(script, tokenPos);
     if (line == null) {
@@ -299,7 +299,7 @@ Future<List<Map<String, dynamic>>> _processSourceReport(
     if (!scopedOutput.includesScript(scriptUriString)) {
       // Sometimes a range's script can be different to the function's script
       // (eg mixins), so we have to re-check the scope filter.
-      // See https://github.com/dart-lang/coverage/issues/495
+      // See https://github.com/dart-lang/tools/issues/530
       continue;
     }
     final scriptUri = Uri.parse(scriptUriString!);
@@ -308,7 +308,7 @@ Future<List<Map<String, dynamic>>> _processSourceReport(
     // SourceReportCoverage.misses: to add zeros to the coverage result for all
     // the lines that don't have a hit. Afterwards, add all the lines that were
     // hit or missed to the cache, so that the next coverage collection won't
-    // need to compile this libarry.
+    // need to compile this library.
     final coverableLines =
         coverableLineCache?.putIfAbsent(scriptUriString, () => <int>{});
 

pkgs/coverage/lib/src/isolate_paused_listener.dart

@@ -0,0 +1,218 @@
+// Copyright (c) 2024, the Dart project authors.  Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'dart:async';
+import 'dart:collection';
+
+import 'package:meta/meta.dart';
+import 'package:vm_service/vm_service.dart';
+
+import 'util.dart';
+
+/// Calls onIsolatePaused whenever an isolate reaches the pause-on-exit state,
+/// and passes a flag stating whether that isolate is the last one in the group.
+class IsolatePausedListener {
+  IsolatePausedListener(this._service, this._onIsolatePaused);
+
+  final VmService _service;
+  final Future<void> Function(IsolateRef isolate, bool isLastIsolateInGroup)
+      _onIsolatePaused;
+  final _allExitedCompleter = Completer<void>();
+  IsolateRef? _mainIsolate;
+
+  @visibleForTesting
+  final isolateGroups = <String, IsolateGroupState>{};
+
+  /// Starts listening and returns a future that completes when all isolates
+  /// have exited.
+  Future<void> waitUntilAllExited() async {
+    await listenToIsolateLifecycleEvents(_service, _onStart, _onPause, _onExit);
+
+    await _allExitedCompleter.future;
+
+    // Resume the main isolate.
+    if (_mainIsolate != null) {
+      await _service.resume(_mainIsolate!.id!);
+    }
+  }
+
+  IsolateGroupState _getGroup(IsolateRef isolateRef) =>
+      isolateGroups[isolateRef.isolateGroupId!] ??= IsolateGroupState();
+
+  void _onStart(IsolateRef isolateRef) {
+    if (_allExitedCompleter.isCompleted) return;
+    _getGroup(isolateRef).start(isolateRef.id!);
+  }
+
+  Future<void> _onPause(IsolateRef isolateRef) async {
+    if (_allExitedCompleter.isCompleted) return;
+    final group = _getGroup(isolateRef);
+    group.pause(isolateRef.id!);
+    try {
+      await _onIsolatePaused(isolateRef, group.noRunningIsolates);
+    } finally {
+      await _maybeResumeIsolate(isolateRef);
+    }
+  }
+
+  Future<void> _maybeResumeIsolate(IsolateRef isolateRef) async {
+    if (_mainIsolate == null && _isMainIsolate(isolateRef)) {
+      _mainIsolate = isolateRef;
+      // Pretend this isolate has exited so _allExitedCompleter can complete.
+      _onExit(isolateRef);
+    } else {
+      await _service.resume(isolateRef.id!);
+    }
+  }
+
+  void _onExit(IsolateRef isolateRef) {
+    if (_allExitedCompleter.isCompleted) return;
+    _getGroup(isolateRef).exit(isolateRef.id!);
+    if (isolateGroups.values.every((group) => group.noLiveIsolates)) {
+      _allExitedCompleter.complete();
+    }
+  }
+
+  static bool _isMainIsolate(IsolateRef isolateRef) {
+    // HACK: This should pretty reliably detect the main isolate, but it's not
+    // foolproof and relies on unstable features. The Dart standalone embedder
+    // and Flutter both call the main isolate "main", and they both also list
+    // this isolate first when querying isolates from the VM service. So
+    // selecting the first isolate named "main" combines these conditions and
+    // should be reliable enough for now, while we wait for a better test.
+    // TODO(https://github.com/dart-lang/sdk/issues/56732): Switch to more
+    // reliable test when it's available.
+    return isolateRef.name == 'main';
+  }
+}
+
+/// Listens to isolate start and pause events, and backfills events for isolates
+/// that existed before listening started.
+///
+/// Ensures that:
+///  - Every [onIsolatePaused] and [onIsolateExited] call will be preceeded by
+///    an [onIsolateStarted] call for the same isolate.
+///  - Not every [onIsolateExited] call will be preceeded by a [onIsolatePaused]
+///    call, but a [onIsolatePaused] will never follow a [onIsolateExited].
+///  - [onIsolateExited] will always run after [onIsolatePaused] completes, even
+///    if an exit event arrives while [onIsolatePaused] is being awaited.
+///  - Each callback will only be called once per isolate.
+Future<void> listenToIsolateLifecycleEvents(
+    VmService service,
+    void Function(IsolateRef isolate) onIsolateStarted,
+    Future<void> Function(IsolateRef isolate) onIsolatePaused,
+    void Function(IsolateRef isolate) onIsolateExited) async {
+  final started = <String>{};
+  void onStart(IsolateRef isolateRef) {
+    if (started.add(isolateRef.id!)) onIsolateStarted(isolateRef);
+  }
+
+  final paused = <String, Future<void>>{};
+  Future<void> onPause(IsolateRef isolateRef) async {
+    onStart(isolateRef);
+    await (paused[isolateRef.id!] ??= onIsolatePaused(isolateRef));
+  }
+
+  final exited = <String>{};
+  Future<void> onExit(IsolateRef isolateRef) async {
+    onStart(isolateRef);
+    if (exited.add(isolateRef.id!)) {
+      // Wait for in-progress pause callbacks. Otherwise prevent future pause
+      // callbacks from running.
+      await (paused[isolateRef.id!] ??= Future<void>.value());
+      onIsolateExited(isolateRef);
+    }
+  }
+
+  final eventBuffer = IsolateEventBuffer((Event event) async {
+    switch (event.kind) {
+      case EventKind.kIsolateStart:
+        return onStart(event.isolate!);
+      case EventKind.kPauseExit:
+        return await onPause(event.isolate!);
+      case EventKind.kIsolateExit:
+        return await onExit(event.isolate!);
+    }
+  });
+
+  // Listen for isolate start/exit events.
+  service.onIsolateEvent.listen(eventBuffer.add);
+  await service.streamListen(EventStreams.kIsolate);
+
+  // Listen for isolate paused events.
+  service.onDebugEvent.listen(eventBuffer.add);
+  await service.streamListen(EventStreams.kDebug);
+
+  // Backfill. Add/pause isolates that existed before we subscribed.
+  for (final isolateRef in await getAllIsolates(service)) {
+    onStart(isolateRef);
+    final isolate = await service.getIsolate(isolateRef.id!);
+    if (isolate.pauseEvent?.kind == EventKind.kPauseExit) {
+      await onPause(isolateRef);
+    }
+  }
+
+  // Flush the buffered stream events, and the start processing them as they
+  // arrive.
+  await eventBuffer.flush();
+}
+
+/// Keeps track of isolates in an isolate group.
+class IsolateGroupState {
+  // IDs of the isolates running in this group.
+  @visibleForTesting
+  final running = <String>{};
+
+  // IDs of the isolates paused just before exiting in this group.
+  @visibleForTesting
+  final paused = <String>{};
+
+  bool get noRunningIsolates => running.isEmpty;
+  bool get noLiveIsolates => running.isEmpty && paused.isEmpty;
+
+  void start(String id) {
+    paused.remove(id);
+    running.add(id);
+  }
+
+  void pause(String id) {
+    running.remove(id);
+    paused.add(id);
+  }
+
+  void exit(String id) {
+    running.remove(id);
+    paused.remove(id);
+  }
+}
+
+/// Buffers VM service isolate [Event]s until [flush] is called.
+///
+/// [flush] passes each buffered event to the handler function. After that, any
+/// further events are immediately passed to the handler. [flush] returns a
+/// future that completes when all the events in the queue have been handled (as
+/// well as any events that arrive while flush is in progress).
+class IsolateEventBuffer {
+  IsolateEventBuffer(this._handler);
+
+  final Future<void> Function(Event event) _handler;
+  final _buffer = Queue<Event>();
+  var _flushed = false;
+
+  Future<void> add(Event event) async {
+    if (_flushed) {
+      await _handler(event);
+    } else {
+      _buffer.add(event);
+    }
+  }
+
+  Future<void> flush() async {
+    while (_buffer.isNotEmpty) {
+      final event = _buffer.removeFirst();
+      await _handler(event);
+    }
+    _flushed = true;
+  }
+}