Rework analysis_server to clear flakiness (#813)

diff --git a/lib/src/analysis_server.dart b/lib/src/analysis_server.dart
index f0f433b..984646a 100644
--- a/lib/src/analysis_server.dart
+++ b/lib/src/analysis_server.dart
@@ -113,14 +113,10 @@
       await analysisServer.server.setSubscriptions(<String>['STATUS']);
 
       listenForCompletions();
-      listenForAnalysisComplete();
-      listenForErrors();
 
-      final analysisComplete = getAnalysisCompleteCompleter();
-      await analysisServer.analysis
-          .setAnalysisRoots(<String>[_sourceDirPath], <String>[]);
-      await _sendAddOverlays(<String, String>{mainPath: _warmupSrc});
-      await analysisComplete.future;
+      await analysisServer.analysis.setAnalysisRoots([_sourceDirPath], []);
+      // Warmup.
+      await _sendAddOverlays({mainPath: _warmupSrc});
       await _sendRemoveOverlays();
     } catch (err, st) {
       _logger.severe('Error starting analysis server ($sdkPath): $err.\n$st');
@@ -242,9 +238,7 @@
     _logger.fine('dartdoc: Scheduler queue: ${serverScheduler.queueCount}');
 
     return serverScheduler.schedule(ClosureTask<Map<String, String>>(() async {
-      final analysisCompleter = getAnalysisCompleteCompleter();
       await _loadSources(<String, String>{mainPath: source});
-      await analysisCompleter.future;
 
       final result = await analysisServer.analysis.getHover(mainPath, offset);
       await _unloadSources();
@@ -280,15 +274,13 @@
 
     return serverScheduler
         .schedule(ClosureTask<proto.AnalysisResults>(() async {
-      clearErrors();
-
-      final analysisCompleter = getAnalysisCompleteCompleter();
       sources = _getOverlayMapWithPaths(sources);
       await _loadSources(sources);
-      await analysisCompleter.future;
+      final errors = (await analysisServer.analysis.getErrors(mainPath)).errors;
+      await _unloadSources();
 
-      // Calculate the issues.
-      final issues = getErrors().map((AnalysisError error) {
+      // Convert the issues to protos.
+      final issues = errors.map((error) {
         final issue = proto.AnalysisIssue()
           ..kind = error.severity.toLowerCase()
           ..line = error.location.startLine
@@ -344,13 +336,14 @@
     }
 
     return serverScheduler.schedule(ClosureTask<AssistsResult>(() async {
-      final analysisCompleter = getAnalysisCompleteCompleter();
       await _loadSources(sources);
-      await analysisCompleter.future;
-      const length = 1;
-      final assists =
-          await analysisServer.edit.getAssists(path, offset, length);
-      await _unloadSources();
+      final AssistsResult assists;
+      try {
+        assists =
+            await analysisServer.edit.getAssists(path, offset, 1 /* length */);
+      } finally {
+        await _unloadSources();
+      }
       return assists;
     }, timeoutDuration: _analysisServerTimeout));
   }
@@ -475,8 +468,12 @@
         _getPathFromName(sourceName),
         offset,
       );
-      final results = await getCompletionResults(id.id);
-      await _unloadSources();
+      final CompletionResults results;
+      try {
+        results = await getCompletionResults(id.id);
+      } finally {
+        await _unloadSources();
+      }
       return results;
     }, timeoutDuration: _analysisServerTimeout));
   }
@@ -492,11 +489,13 @@
     }
 
     return serverScheduler.schedule(ClosureTask<FixesResult>(() async {
-      final analysisCompleter = getAnalysisCompleteCompleter();
       await _loadSources(sources);
-      await analysisCompleter.future;
-      final fixes = await analysisServer.edit.getFixes(path, offset);
-      await _unloadSources();
+      final FixesResult fixes;
+      try {
+        fixes = await analysisServer.edit.getFixes(path, offset);
+      } finally {
+        await _unloadSources();
+      }
       return fixes;
     }, timeoutDuration: _analysisServerTimeout));
   }
@@ -505,9 +504,13 @@
     _logger.fine('FormatImpl: Scheduler queue: ${serverScheduler.queueCount}');
 
     return serverScheduler.schedule(ClosureTask<FormatResult>(() async {
-      await _loadSources(<String, String>{mainPath: src});
-      final result = await analysisServer.edit.format(mainPath, offset, 0);
-      await _unloadSources();
+      await _loadSources({mainPath: src});
+      final FormatResult result;
+      try {
+        result = await analysisServer.edit.format(mainPath, offset, 0);
+      } finally {
+        await _unloadSources();
+      }
       return result;
     }, timeoutDuration: _analysisServerTimeout));
   }
@@ -523,48 +526,50 @@
   String _getPathFromName(String sourceName) =>
       path.join(_sourceDirPath, sourceName);
 
-  /// Warm up the analysis server to be ready for use.
-  Future<void> warmup() => complete(_warmupSrc, 10);
-
   final Set<String> _overlayPaths = <String>{};
 
+  /// Loads [sources] as file system overlays to the analysis server.
+  ///
+  /// The analysis server then begins to analyze these as priority files.
   Future<void> _loadSources(Map<String, String> sources) async {
     if (_overlayPaths.isNotEmpty) {
-      await _sendRemoveOverlays();
+      throw StateError(
+          'There should be no overlay paths while loading sources, but we '
+          'have: $_overlayPaths');
     }
     await _sendAddOverlays(sources);
     await analysisServer.analysis.setPriorityFiles(sources.keys.toList());
   }
 
-  Future<dynamic> _unloadSources() {
-    return Future.wait(<Future<dynamic>>[
-      _sendRemoveOverlays(),
-      analysisServer.analysis.setPriorityFiles(<String>[]),
-    ]);
+  Future<void> _unloadSources() async {
+    await _sendRemoveOverlays();
+    await analysisServer.analysis.setPriorityFiles([]);
   }
 
-  Future<dynamic> _sendAddOverlays(Map<String, String> overlays) {
-    final params = overlays.map((overlayPath, content) =>
+  /// Sends [overlays] to the analysis server.
+  Future<void> _sendAddOverlays(Map<String, String> overlays) async {
+    final contentOverlays = overlays.map((overlayPath, content) =>
         MapEntry(overlayPath, AddContentOverlay(content)));
 
     _logger.fine('About to send analysis.updateContent');
-    _logger.fine('  ${params.keys}');
+    _logger.fine('  ${contentOverlays.keys}');
 
-    _overlayPaths.addAll(params.keys);
+    _overlayPaths.addAll(contentOverlays.keys);
 
-    return analysisServer.analysis.updateContent(params);
+    await analysisServer.analysis.updateContent(contentOverlays);
   }
 
-  Future<dynamic> _sendRemoveOverlays() {
+  Future<void> _sendRemoveOverlays() async {
     _logger.fine('About to send analysis.updateContent remove overlays:');
     _logger.fine('  $_overlayPaths');
 
-    final params = {
+    final contentOverlays = {
       for (final overlayPath in _overlayPaths)
         overlayPath: RemoveContentOverlay()
     };
     _overlayPaths.clear();
-    return analysisServer.analysis.updateContent(params);
+
+    await analysisServer.analysis.updateContent(contentOverlays);
   }
 
   final Map<String, Completer<CompletionResults>> _completionCompleters =
@@ -586,50 +591,6 @@
     _completionCompleters[id] = completer;
     return completer.future;
   }
-
-  final List<Completer<void>> _analysisCompleters = [];
-
-  void listenForAnalysisComplete() {
-    analysisServer.server.onStatus.listen((ServerStatus status) {
-      final analysis = status.analysis;
-      if (analysis == null) return;
-
-      if (!analysis.isAnalyzing) {
-        for (final completer in _analysisCompleters) {
-          completer.complete();
-        }
-
-        _analysisCompleters.clear();
-      }
-    });
-  }
-
-  Completer<void> getAnalysisCompleteCompleter() {
-    final completer = Completer<void>();
-    _analysisCompleters.add(completer);
-    return completer;
-  }
-
-  final Map<String, List<AnalysisError>> _errors =
-      <String, List<AnalysisError>>{};
-
-  void listenForErrors() {
-    analysisServer.analysis.onErrors.listen((AnalysisErrors result) {
-      if (result.errors.isEmpty) {
-        _errors.remove(result.file);
-      } else {
-        _errors[result.file] = result.errors;
-      }
-    });
-  }
-
-  void clearErrors() => _errors.clear();
-
-  List<AnalysisError> getErrors() {
-    return [
-      for (final errors in _errors.values) ...errors,
-    ];
-  }
 }
 
 class Location {
diff --git a/lib/src/analysis_servers.dart b/lib/src/analysis_servers.dart
index e351bdd..3952f36 100644
--- a/lib/src/analysis_servers.dart
+++ b/lib/src/analysis_servers.dart
@@ -39,7 +39,7 @@
   bool get isHealthy => (_restartingSince == null ||
       DateTime.now().difference(_restartingSince!).inMinutes < 30);
 
-  Future<List<void>> warmup() async {
+  Future<void> warmup() async {
     _logger.info('Beginning AnalysisServersWrapper init().');
     _dartAnalysisServer = DartAnalysisServerWrapper(dartSdkPath: _dartSdkPath);
     _flutterAnalysisServer =
@@ -66,11 +66,6 @@
     }));
 
     _restartingSince = null;
-
-    return Future.wait([
-      _flutterAnalysisServer.warmup(),
-      _dartAnalysisServer.warmup(),
-    ]);
   }
 
   Future<void> _restart() async {
diff --git a/test/common_server_api_protobuf_test.dart b/test/common_server_api_protobuf_test.dart
index e9414ee..66ccb1d 100644
--- a/test/common_server_api_protobuf_test.dart
+++ b/test/common_server_api_protobuf_test.dart
@@ -20,8 +20,6 @@
 import 'package:shelf/shelf_io.dart' as shelf_io;
 import 'package:test/test.dart';
 
-import 'utils.dart';
-
 const preFormattedCode = r'''
 void main()
 {
@@ -139,18 +137,14 @@
     test('analyze errors', () async {
       final request = proto.SourceRequest()..source = sampleCodeError;
       late proto.AnalysisResults reply;
-      await tryWithReruns(() async {
-        final response =
-            await sendPostRequest('dartservices/v2/analyze', request);
-        expect(response.statusCode, 200);
-        expect(response.headers['content-type'],
-            ['application/json; charset=utf-8']);
-        final data = json.decode(await response.transform(utf8.decoder).join());
-        reply = proto.AnalysisResults()..mergeFromProto3Json(data);
-        if (reply.issues.isEmpty) {
-          throw StateError('Flaky response');
-        }
-      });
+      final response =
+          await sendPostRequest('dartservices/v2/analyze', request);
+      expect(response.statusCode, 200);
+      expect(response.headers['content-type'],
+          ['application/json; charset=utf-8']);
+      final data = json.decode(await response.transform(utf8.decoder).join());
+      reply = proto.AnalysisResults()..mergeFromProto3Json(data);
+      expect(reply.issues, hasLength(1));
       expect(reply.issues[0].kind, 'error');
       expect(reply.issues[0].line, 2);
       expect(reply.issues[0].sourceName, 'main.dart');
diff --git a/test/common_server_api_test.dart b/test/common_server_api_test.dart
index 6e99307..26a3279 100644
--- a/test/common_server_api_test.dart
+++ b/test/common_server_api_test.dart
@@ -19,8 +19,6 @@
 import 'package:shelf/shelf_io.dart' as shelf_io;
 import 'package:test/test.dart';
 
-import 'utils.dart';
-
 const versions = ['v1', 'v2'];
 
 const preFormattedCode = r'''
@@ -189,19 +187,13 @@
     test('analyze errors', () async {
       for (final version in versions) {
         final jsonData = {'source': sampleCodeError};
-        late Map<String, Object> dataMap;
-        await tryWithReruns(() async {
-          final response =
-              await sendPostRequest('dartservices/$version/analyze', jsonData);
-          expect(response.statusCode, 200);
-          expect(response.headers['content-type'],
-              ['application/json; charset=utf-8']);
-          final data = await response.transform(utf8.decoder).join();
-          dataMap = (json.decode(data) as Map).cast<String, Object>();
-          if (dataMap.isEmpty) {
-            throw StateError('Flaky result');
-          }
-        });
+        final response =
+            await sendPostRequest('dartservices/$version/analyze', jsonData);
+        expect(response.statusCode, 200);
+        expect(response.headers['content-type'],
+            ['application/json; charset=utf-8']);
+        final data = await response.transform(utf8.decoder).join();
+        final dataMap = (json.decode(data) as Map).cast<String, Object>();
         expect(
           dataMap,
           {
diff --git a/test/flutter_analysis_server_test.dart b/test/flutter_analysis_server_test.dart
index be6974b..9bab401 100644
--- a/test/flutter_analysis_server_test.dart
+++ b/test/flutter_analysis_server_test.dart
@@ -15,8 +15,6 @@
 import 'package:dart_services/src/server_cache.dart';
 import 'package:test/test.dart';
 
-import 'utils.dart';
-
 final channel = Platform.environment['FLUTTER_CHANNEL'] ?? stableChannel;
 void main() => defineTests();
 
@@ -29,7 +27,6 @@
       analysisServer =
           FlutterAnalysisServerWrapper(dartSdkPath: sdk.dartSdkPath);
       await analysisServer.init();
-      await analysisServer.warmup();
     });
 
     tearDown(() async {
@@ -65,14 +62,12 @@
 
     test('reports errors with Flutter code', () async {
       late AnalysisResults results;
-      await tryWithReruns(() async {
-        results = await analysisServersWrapper.analyze('''
+      results = await analysisServersWrapper.analyze('''
 import 'package:flutter/material.dart';
 
 String x = 7;
 
 void main() async {
-
   runApp(MaterialApp(
       debugShowCheckedModeBanner: false, home: Scaffold(body: HelloWorld())));
 }
@@ -82,10 +77,6 @@
   Widget build(context) => const Center(child: Text('Hello world'));
 }
 ''');
-        if (results.issues.isEmpty) {
-          throw StateError('Flaky result');
-        }
-      });
       expect(results.issues, hasLength(1));
       final issue = results.issues[0];
       expect(issue.line, 3);
@@ -99,8 +90,7 @@
     // https://github.com/dart-lang/dart-pad/issues/2005
     test('reports lint with Flutter code', () async {
       late AnalysisResults results;
-      await tryWithReruns(() async {
-        results = await analysisServersWrapper.analyze('''
+      results = await analysisServersWrapper.analyze('''
 import 'package:flutter/material.dart';
 
 void main() async {
@@ -116,10 +106,6 @@
   Widget build(context) => const Center(child: Text('Hello world'));
 }
 ''');
-        if (results.issues.isEmpty) {
-          throw StateError('Flaky result');
-        }
-      });
       expect(results.issues, hasLength(1));
       final issue = results.issues[0];
       expect(issue.line, 4);
diff --git a/test/utils.dart b/test/utils.dart
deleted file mode 100644
index e8637c8..0000000
--- a/test/utils.dart
+++ /dev/null
@@ -1,24 +0,0 @@
-// Copyright (c) 2021, 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';
-
-/// Runs [fn] up to three times, catching thrown exceptions, until it returns
-/// normally.
-Future<void> tryWithReruns(FutureOr<void> Function() fn) async {
-  for (var i = 0; i < 3; i++) {
-    try {
-      await fn();
-      // If the function returns normally, continue.
-      return;
-    } catch (e) {
-      print('Setup function threw: $e;');
-      if (i < 2) {
-        print('Retrying...');
-      } else {
-        rethrow;
-      }
-    }
-  }
-}
diff --git a/tool/fuzz_driver.dart b/tool/fuzz_driver.dart
index df9cf67..f0e8270 100644
--- a/tool/fuzz_driver.dart
+++ b/tool/fuzz_driver.dart
@@ -131,9 +131,6 @@
       analysis_server.DartAnalysisServerWrapper(dartSdkPath: sdk.dartSdkPath);
   await analysisServer!.init();
 
-  print('Warming up analysis server');
-  await analysisServer!.warmup();
-
   print('Warming up compiler');
   compiler = comp.Compiler(sdk);
   await compiler.warmup();