[coverage] Finish collection as soon as main isolate exits

pkgs/coverage/CHANGELOG.md

@@ -1,6 +1,9 @@
 ## 1.12.0-wip
 
 - Introduced support for specifying coverage flags through a YAML file.
+- Fix a [bug](https://github.com/dart-lang/tools/issues/2049) where coverage
+  collection could hang forever if the package under test used background
+  isolates that were never shut down.
 
 ## 1.11.1
 

pkgs/coverage/lib/src/isolate_paused_listener.dart

@@ -28,101 +28,93 @@ class IsolatePausedListener {
 
   final _isolateGroups = <String, IsolateGroupState>{};
 
-  int _numNonMainIsolates = 0;
-  final _allNonMainIsolatesExited = Completer<void>();
-  bool _finished = false;
+  int _numIsolates = 0;
+  bool _finishedListening = false;
 
   IsolateRef? _mainIsolate;
-  final _mainIsolatePaused = Completer<bool>();
+  bool _hasMainIsolate = false;
+
+  // Completes when either:
+  // - there is a main isolate, and it is paused or exited
+  // - there is no main isolate, and all isolates have exited
+  final _mainIsolatePausedOrAllIsolatesExited = Completer<void>();
 
   /// Starts listening and returns a future that completes when all isolates
   /// have exited.
   Future<void> waitUntilAllExited() async {
     await listenToIsolateLifecycleEvents(_service, _onStart, _onPause, _onExit);
 
-    await _allNonMainIsolatesExited.future;
-
-    // Resume the main isolate.
-    try {
-      if (_mainIsolate != null) {
-        if (await _mainIsolatePaused.future) {
-          await _runCallbackAndResume(
-              _mainIsolate!, !_getGroup(_mainIsolate!).collected);
+    await _mainIsolatePausedOrAllIsolatesExited.future;
+    _finishedListening = true;
+
+    // Collect all remaining uncollected groups.
+    final collectionTasks = <Future<void>>[];
+    for (final group in _isolateGroups.values) {
+      if (!group.collected) {
+        group.collected = true;
+        final iso = group.paused.firstOrNull ?? group.running.firstOrNull;
+        if (iso != null) {
+          collectionTasks.add(_onIsolatePaused(iso, true));
         }
       }
-    } finally {
-      _finished = true;
+    }
+    await Future.wait(collectionTasks);
+
+    // 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 (_finished) return;
+    if (_finishedListening) return;
     final group = _getGroup(isolateRef);
-    group.start(isolateRef.id!);
-    if (_mainIsolate == null && _isMainIsolate(isolateRef)) {
+    group.start(isolateRef);
+    ++_numIsolates;
+    if (!_hasMainIsolate && _isMainIsolate(isolateRef)) {
       _mainIsolate = isolateRef;
-    } else {
-      ++_numNonMainIsolates;
+      _hasMainIsolate = true;
     }
   }
 
   Future<void> _onPause(IsolateRef isolateRef) async {
-    if (_finished) return;
+    if (_finishedListening) return;
     final group = _getGroup(isolateRef);
-    group.pause(isolateRef.id!);
+    group.pause(isolateRef);
     if (isolateRef.id! == _mainIsolate?.id) {
-      _mainIsolatePaused.complete(true);
-
-      // If the main isolate is the only isolate, then _allNonMainIsolatesExited
-      // will never be completed. So check that case here.
-      _checkCompleted();
+      _mainIsolatePausedOrAllIsolatesExited.safeComplete();
     } else {
-      await _runCallbackAndResume(isolateRef, group.noRunningIsolates);
-    }
-  }
-
-  Future<void> _runCallbackAndResume(
-      IsolateRef isolateRef, bool isLastIsolateInGroup) async {
-    if (isLastIsolateInGroup) {
-      _getGroup(isolateRef).collected = true;
-    }
-    try {
-      await _onIsolatePaused(isolateRef, isLastIsolateInGroup);
-    } finally {
-      await _service.resume(isolateRef.id!);
+      final isLastIsolateInGroup = group.noRunningIsolates;
+      if (isLastIsolateInGroup) {
+        _getGroup(isolateRef).collected = true;
+      }
+      try {
+        await _onIsolatePaused(isolateRef, isLastIsolateInGroup);
+      } finally {
+        group.exit(isolateRef);
+        await _service.resume(isolateRef.id!);
+      }
     }
   }
 
   void _onExit(IsolateRef isolateRef) {
-    if (_finished) return;
+    if (_finishedListening) return;
     final group = _getGroup(isolateRef);
-    group.exit(isolateRef.id!);
+    group.exit(isolateRef);
+    --_numIsolates;
     if (group.noLiveIsolates && !group.collected) {
       _log('ERROR: An isolate exited without pausing, causing '
           'coverage data to be lost for group ${isolateRef.isolateGroupId!}.');
     }
     if (isolateRef.id! == _mainIsolate?.id) {
-      if (!_mainIsolatePaused.isCompleted) {
-        // Main isolate exited without pausing.
-        _mainIsolatePaused.complete(false);
-      }
-    } else {
-      --_numNonMainIsolates;
-      _checkCompleted();
-    }
-  }
-
-  bool get _mainRunning =>
-      _mainIsolate != null && !_mainIsolatePaused.isCompleted;
-
-  void _checkCompleted() {
-    if (_numNonMainIsolates == 0 &&
-        !_mainRunning &&
-        !_allNonMainIsolatesExited.isCompleted) {
-      _allNonMainIsolatesExited.complete();
+      // Main isolate exited without pausing.
+      _mainIsolate = null;
+      _mainIsolatePausedOrAllIsolatesExited.safeComplete();
+    } else if (!_hasMainIsolate && _numIsolates == 0) {
+      _mainIsolatePausedOrAllIsolatesExited.safeComplete();
     }
   }
 
@@ -139,6 +131,12 @@ class IsolatePausedListener {
   }
 }
 
+extension on Completer {
+  void safeComplete() {
+    if (!isCompleted) complete();
+  }
+}
+
 /// Listens to isolate start and pause events, and backfills events for isolates
 /// that existed before listening started.
 ///
@@ -220,30 +218,30 @@ Future<void> listenToIsolateLifecycleEvents(
 class IsolateGroupState {
   // IDs of the isolates running in this group.
   @visibleForTesting
-  final running = <String>{};
+  final running = <IsolateRef>{};
 
   // IDs of the isolates paused just before exiting in this group.
   @visibleForTesting
-  final paused = <String>{};
+  final paused = <IsolateRef>{};
 
   bool collected = false;
 
   bool get noRunningIsolates => running.isEmpty;
   bool get noLiveIsolates => running.isEmpty && paused.isEmpty;
 
-  void start(String id) {
-    paused.remove(id);
-    running.add(id);
+  void start(IsolateRef iso) {
+    paused.remove(iso);
+    running.add(iso);
   }
 
-  void pause(String id) {
-    running.remove(id);
-    paused.add(id);
+  void pause(IsolateRef iso) {
+    running.remove(iso);
+    paused.add(iso);
   }
 
-  void exit(String id) {
-    running.remove(id);
-    paused.remove(id);
+  void exit(IsolateRef iso) {
+    running.remove(iso);
+    paused.remove(iso);
   }
 
   @override

pkgs/coverage/test/collect_coverage_test.dart

@@ -133,24 +133,25 @@ void main() {
       39: 1,
       41: 1,
       42: 4,
-      43: 1,
-      44: 3,
-      45: 1,
-      48: 1,
+      43: 2,
+      44: 1,
+      45: 3,
+      46: 1,
       49: 1,
-      51: 1,
-      54: 1,
+      50: 1,
+      52: 1,
       55: 1,
       56: 1,
-      59: 1,
+      57: 1,
       60: 1,
-      62: 1,
+      61: 1,
       63: 1,
       64: 1,
-      66: 1,
+      65: 1,
       67: 1,
       68: 1,
-      71: 3,
+      69: 1,
+      72: 3,
     };
     expect(isolateFile?.lineHits, expectedHits);
     expect(isolateFile?.funcHits, {11: 1, 19: 1, 23: 1, 28: 1, 38: 1});
@@ -174,7 +175,7 @@ void main() {
         32: 0,
         38: 1,
         42: 1,
-        71: 1,
+        72: 1,
       },
     );
   });