Improve search isolate testing by loading an actual index. (#8821)

Author: isoos

app/bin/tools/search_benchmark.dart

@@ -2,26 +2,14 @@
 // 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:convert';
-import 'dart:io';
-
 import 'package:_pub_shared/search/search_form.dart';
-import 'package:pub_dev/package/overrides.dart';
-import 'package:pub_dev/search/mem_index.dart';
-import 'package:pub_dev/search/models.dart';
 import 'package:pub_dev/search/search_service.dart';
+import 'package:pub_dev/search/updater.dart';
 
 /// Loads a search snapshot and executes queries on it, benchmarking their total time to complete.
 Future<void> main(List<String> args) async {
   // Assumes that the first argument is a search snapshot file.
-  final file = File(args.first);
-  final content =
-      json.decode(utf8.decode(gzip.decode(await file.readAsBytes())))
-          as Map<String, Object?>;
-  final snapshot = SearchSnapshot.fromJson(content);
-  final index = InMemoryPackageIndex(
-      documents:
-          snapshot.documents!.values.where((d) => !isSdkPackage(d.package)));
+  final index = await loadInMemoryPackageIndexFromFile(args.first);
 
   // NOTE: please add more queries to this list, especially if there is a performance bottleneck.
   final queries = [

app/lib/search/updater.dart

@@ -3,10 +3,14 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import 'dart:async';
+import 'dart:convert';
+import 'dart:io';
 
 import 'package:gcloud/service_scope.dart' as ss;
 import 'package:logging/logging.dart';
 import 'package:meta/meta.dart';
+import 'package:pub_dev/package/overrides.dart';
+import 'package:pub_dev/search/models.dart';
 import 'package:pub_dev/search/search_service.dart';
 
 import '../package/models.dart' show Package;
@@ -24,6 +28,31 @@ void registerIndexUpdater(IndexUpdater updater) =>
 /// The active index updater.
 IndexUpdater get indexUpdater => ss.lookup(#_indexUpdater) as IndexUpdater;
 
+/// Loads a local search snapshot file and builds an in-memory package index from it.
+Future<InMemoryPackageIndex> loadInMemoryPackageIndexFromFile(
+    String path) async {
+  final file = File(path);
+  final content =
+      json.decode(utf8.decode(gzip.decode(await file.readAsBytes())))
+          as Map<String, Object?>;
+  final snapshot = SearchSnapshot.fromJson(content);
+  return InMemoryPackageIndex(
+      documents:
+          snapshot.documents!.values.where((d) => !isSdkPackage(d.package)));
+}
+
+/// Saves the provided [documents] into a local search snapshot file.
+Future<void> saveInMemoryPackageIndexToFile(
+    Iterable<PackageDocument> documents, String path) async {
+  final file = File(path);
+  final snapshot = SearchSnapshot();
+  for (final doc in documents) {
+    snapshot.add(doc);
+  }
+  await file
+      .writeAsBytes(gzip.encode(utf8.encode(json.encode(snapshot.toJson()))));
+}
+
 class IndexUpdater {
   final DatastoreDB _db;
 

app/lib/service/entrypoint/_isolate.dart

@@ -33,6 +33,7 @@ class IsolateRunner {
   final ServicesWrapperFn? servicesWrapperFn;
   final EntryPointFn? entryPoint;
   final Uri? spawnUri;
+  final List<String>? spawnArgs;
 
   int started = 0;
   final _isolates = <_Isolate>[];
@@ -43,14 +44,17 @@ class IsolateRunner {
     required this.kind,
     required ServicesWrapperFn this.servicesWrapperFn,
     required EntryPointFn this.entryPoint,
-  }) : spawnUri = null;
+  })  : spawnUri = null,
+        spawnArgs = null;
 
   IsolateRunner.uri({
     required this.logger,
     required this.kind,
     required Uri this.spawnUri,
+    List<String>? spawnArgs,
   })  : entryPoint = null,
-        servicesWrapperFn = null;
+        servicesWrapperFn = null,
+        spawnArgs = spawnArgs ?? <String>[];
 
   /// Starts [count] new isolates.
   Future<void> start(int count) async {
@@ -138,7 +142,7 @@ class IsolateRunner {
         entryPoint: entryPoint!,
       );
     } else {
-      await isolate.spawnUri(spawnUri: spawnUri!);
+      await isolate.spawnUri(spawnUri: spawnUri!, args: spawnArgs);
     }
     if (_closing) {
       await isolate.close();
@@ -176,12 +180,14 @@ Future<IsolateRunner> startWorkerIsolate({
 Future<IsolateRunner> startQueryIsolate({
   required Logger logger,
   required Uri spawnUri,
+  List<String>? spawnArgs,
   required String kind,
 }) async {
   final worker = IsolateRunner.uri(
     logger: logger,
     kind: kind,
     spawnUri: spawnUri,
+    spawnArgs: spawnArgs ?? [],
   );
   await worker.start(1);
   return worker;
@@ -242,10 +248,11 @@ class _Isolate {
 
   Future<void> spawnUri({
     required Uri spawnUri,
+    required List<String>? args,
   }) async {
     _isolate = await Isolate.spawnUri(
       spawnUri,
-      [],
+      args ?? [],
       EntryMessage(
         protocolSendPort: _protocolReceivePort.sendPort,
       ).encodeAsJson(),

app/lib/service/entrypoint/search_index.dart

@@ -5,6 +5,7 @@
 import 'dart:async';
 import 'dart:convert';
 
+import 'package:args/args.dart';
 import 'package:gcloud/service_scope.dart';
 import 'package:logging/logging.dart';
 import 'package:pub_dev/search/backend.dart';
@@ -20,10 +21,20 @@ import 'package:pub_dev/shared/utils.dart';
 
 final _logger = Logger('search_index');
 
+final _argParser = ArgParser()
+  ..addOption(
+    'snapshot',
+    help:
+        'If specified, the snapshot will be loaded from the local path instead of cloud storage.',
+  );
+
 /// Entry point for the search index isolate.
 Future<void> main(List<String> args, var message) async {
   final timer = Timer.periodic(Duration(milliseconds: 250), (_) {});
 
+  final argv = _argParser.parse(args);
+  final snapshot = argv['snapshot'] as String?;
+
   final ServicesWrapperFn servicesWrapperFn;
   if (envConfig.isRunningInAppengine) {
     servicesWrapperFn = withServices;
@@ -34,7 +45,11 @@ Future<void> main(List<String> args, var message) async {
   }
   await fork(() async {
     await servicesWrapperFn(() async {
-      await indexUpdater.init();
+      if (snapshot == null) {
+        await indexUpdater.init();
+      } else {
+        updatePackageIndex(await loadInMemoryPackageIndexFromFile(snapshot));
+      }
 
       await runIsolateFunctions(
         message: message,

app/test/service/entrypoint/search_index_test.dart

@@ -3,39 +3,66 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import 'package:logging/logging.dart';
+import 'package:path/path.dart' as p;
 import 'package:pub_dev/search/search_service.dart';
+import 'package:pub_dev/search/updater.dart';
 import 'package:pub_dev/service/entrypoint/_isolate.dart';
 import 'package:pub_dev/service/entrypoint/search_index.dart';
+import 'package:pub_dev/shared/utils.dart';
 import 'package:test/test.dart';
 
 final _logger = Logger('search_index_test');
 
 void main() {
   group('Search index inside an isolate', () {
-    final indexRunner = IsolateRunner.uri(
-      kind: 'index',
-      logger: _logger,
-      spawnUri:
-          Uri.parse('package:pub_dev/service/entrypoint/search_index.dart'),
-    );
-
-    tearDownAll(() async {
+    late IsolateRunner indexRunner;
+
+    tearDown(() async {
       await indexRunner.close();
     });
 
-    test('start and work with index', () async {
-      await indexRunner.start(1);
+    test('start and work with local index', () async {
+      await withTempDirectory((tempDir) async {
+        final snapshotPath = p.join(tempDir.path, 'index.json.gz');
+        await saveInMemoryPackageIndexToFile(
+          [
+            PackageDocument(
+              package: 'json_annotation',
+              description: 'Annotation metadata for JSON serialization.',
+            ),
+          ],
+          snapshotPath,
+        );
+
+        indexRunner = IsolateRunner.uri(
+          kind: 'index',
+          logger: _logger,
+          spawnUri:
+              Uri.parse('package:pub_dev/service/entrypoint/search_index.dart'),
+          spawnArgs: ['--snapshot', snapshotPath],
+        );
+
+        await indexRunner.start(1);
 
-      // index calling the sendport
-      final searchIndex = IsolateSearchIndex(indexRunner);
-      expect(await searchIndex.isReady(), true);
+        // index calling the sendport
+        final searchIndex = IsolateSearchIndex(indexRunner);
+        expect(await searchIndex.isReady(), true);
 
-      // working search only with SDK results (no packages in the isolate)
-      final rs =
-          await searchIndex.search(ServiceSearchQuery.parse(query: 'json'));
-      expect(rs.errorMessage, isNull);
-      expect(rs.sdkLibraryHits, isEmpty);
-      expect(rs.packageHits, isEmpty);
+        // returns package hit
+        final rs =
+            await searchIndex.search(ServiceSearchQuery.parse(query: 'json'));
+        expect(rs.toJson(), {
+          'timestamp': isNotEmpty,
+          'totalCount': 1,
+          'sdkLibraryHits': [],
+          'packageHits': [
+            {
+              'package': 'json_annotation',
+              'score': greaterThan(0.5),
+            },
+          ],
+        });
+      });
     }, timeout: Timeout(Duration(minutes: 5)));
   });
 }