Improve performance of Java package parser by making it Bazel worker

aspect/intellij_info_impl.bzl

@@ -694,7 +694,15 @@ def build_java_package_manifest(ctx, target, source_files, suffix):
         join_with = ":",
         map_each = _package_manifest_file_argument,
     )
-    args.use_param_file("@%s")
+
+    # Bazel has an option to put your command line args in a file, and then pass the name of that file as the only
+    # argument to your executable. The PackageParser supports taking args in this way, we can pass in an args file
+    # as "@filename".
+    # Bazel Persistent Workers take their input as a file that contains the argument that will be parsed and turned
+    # into a WorkRequest proto and read on stdin. It also wants an argument of the form "@filename". We can use the
+    # params file as an arg file.
+    # Thus if we always use a params file, we can support both persistent worker mode and local mode (regular) mode.
+    args.use_param_file("@%s", use_always = True)
     args.set_param_file_format("multiline")
 
     ctx.actions.run(
@@ -704,6 +712,10 @@ def build_java_package_manifest(ctx, target, source_files, suffix):
         arguments = [args],
         mnemonic = "JavaPackageManifest",
         progress_message = "Parsing java package strings for " + str(target.label),
+        execution_requirements = {
+            "supports-workers": "1",
+            "requires-worker-protocol": "proto",
+        },
     )
     return output
 

aspect/tools/BUILD

@@ -22,6 +22,7 @@ java_library(
         ":guava",
         "//intellij_platform_sdk:jsr305",
         "//proto:proto_deps",
+        "//third_party/bazel/src/main/protobuf:worker_protocol_java_proto",
     ],
 )
 

aspect/tools/src/com/google/idea/blaze/aspect/PackageParser.java

@@ -22,6 +22,8 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
+import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
 import com.google.devtools.intellij.aspect.Common.ArtifactLocation;
 import com.google.devtools.intellij.ideinfo.IntellijIdeInfo.JavaSourcePackage;
 import com.google.devtools.intellij.ideinfo.IntellijIdeInfo.PackageManifest;
@@ -68,14 +70,61 @@ static PackageParserOptions parseArgs(String[] args) {
 
   private static final Pattern PACKAGE_PATTERN = Pattern.compile("^\\s*package\\s+([\\w\\.]+)");
 
+  private static boolean isWorkerMode(String [] args) {
+    return args.length == 1 && args[0].equals("--persistent_worker");
+  }
+
+  private static String[] getWorkRequestArgs(WorkRequest workRequest) {
+    return workRequest.getArgumentsList().toArray(new String[0]);
+  }
+
+  private static void parsePackagesAndWriteManifest(
+      PackageParser parser, PackageParserOptions options) throws Exception {
+    Preconditions.checkNotNull(options.outputManifest);
+
+    Map<ArtifactLocation, String> outputMap = parser.parsePackageStrings(options.sources);
+    parser.writeManifest(outputMap, options.outputManifest);
+  }
+
+  private static void runPersistentWorker(PackageParser parser) throws IOException {
+    while (true) {
+      WorkRequest workRequest = WorkRequest.parseDelimitedFrom(System.in);
+      if (workRequest == null) {
+        // parseDelimitedFrom returns null iff the stream is at EOF, thus we have no more work to
+        // do.
+        return;
+      }
+
+      WorkResponse.Builder builder =
+          WorkResponse.newBuilder().setRequestId(workRequest.getRequestId());
+
+      try {
+        PackageParserOptions workRequestOptions = parseArgs(getWorkRequestArgs(workRequest));
+
+        parsePackagesAndWriteManifest(parser, workRequestOptions);
+
+        builder.build().writeDelimitedTo(System.out);
+      } catch (Exception e) {
+        logger.log(Level.SEVERE, "Error parsing package strings ", e);
+        builder.setOutput("Error parsing package strings: " + e);
+        builder.setExitCode(1);
+      }
+    }
+  }
+
   public static void main(String[] args) throws Exception {
     PackageParserOptions options = parseArgs(args);
-    Preconditions.checkNotNull(options.outputManifest);
+    PackageParser parser = new PackageParser(PackageParserIoProvider.INSTANCE);
 
     try {
-      PackageParser parser = new PackageParser(PackageParserIoProvider.INSTANCE);
-      Map<ArtifactLocation, String> outputMap = parser.parsePackageStrings(options.sources);
-      parser.writeManifest(outputMap, options.outputManifest);
+      if (isWorkerMode(args)) {
+        // Bazel persistent workers are required to only write WorkResponses to stdout. The
+        // java.util.logging API defaults to writing to System.err, so we don't have to
+        // redirect log output away from stdout.
+        runPersistentWorker(parser);
+      } else {
+        parsePackagesAndWriteManifest(parser, options);
+      }
     } catch (Throwable e) {
       logger.log(Level.SEVERE, "Error parsing package strings", e);
       System.exit(1);