From 72d27152d59a9f8357d6be51e3a8f686c0f79610 Mon Sep 17 00:00:00 2001
From: "xieran.sai" <xieran.sai@bytedance.com>
Date: Fri, 26 May 2023 14:58:15 +0800
Subject: [PATCH] support multiplex worker for jetify

---
 WORKSPACE                                     |  6 +++
 private/rules/jetifier.bzl                    |  6 +++
 third_party/jetifier/BUILD                    | 18 +++++++-
 .../JetifierWorkerWrapper.java                | 46 +++++++++++++++++++
 4 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 third_party/jetifier/jetifier-standalone/java/com/github/bazelbuild/rules_jvm_external/JetifierWorkerWrapper.java

diff --git a/WORKSPACE b/WORKSPACE
index 9f95d4093..841067aca 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -37,6 +37,12 @@ load("@io_bazel_stardoc//:setup.bzl", "stardoc_repositories")
 
 stardoc_repositories()
 
+http_archive(
+    name = "android_tools",
+    sha256 = "1afa4b7e13c82523c8b69e87f8d598c891ec7e2baa41d9e24e08becd723edb4d",
+    url = "https://mirror.bazel.build/bazel_android_tools/android_tools_pkg-0.27.0.tar.gz",
+)
+
 http_archive(
     name = "io_bazel_rules_kotlin",
     sha256 = "946747acdbeae799b085d12b240ec346f775ac65236dfcf18aa0cd7300f6de78",
diff --git a/private/rules/jetifier.bzl b/private/rules/jetifier.bzl
index 9f194749a..913e34e35 100644
--- a/private/rules/jetifier.bzl
+++ b/private/rules/jetifier.bzl
@@ -13,12 +13,18 @@ def _jetify_impl(ctx):
             jetify_args.add("-o", jetified_outfile)
             jetify_args.add("-i", artifact)
             jetify_args.add("-timestampsPolicy", "keepPrevious")
+            jetify_args.use_param_file("@%s", use_always = True)
+            jetify_args.set_param_file_format("multiline")
             ctx.actions.run(
                 mnemonic = "Jetify",
                 inputs = [artifact],
                 outputs = [jetified_outfile],
                 progress_message = "Jetifying {}".format(artifact.owner),
                 executable = ctx.executable._jetifier,
+                execution_requirements = {
+                    "supports-workers" : "1",
+                    "supports-multiplex-workers": "1",
+                },
                 arguments = [jetify_args],
             )
             outfiles.append(jetified_outfile)
diff --git a/third_party/jetifier/BUILD b/third_party/jetifier/BUILD
index 8e07b1273..bff25891b 100644
--- a/third_party/jetifier/BUILD
+++ b/third_party/jetifier/BUILD
@@ -3,11 +3,27 @@ java_import(
     jars = glob(["jetifier-standalone/lib/*.jar"]),
 )
 
+java_import(
+    name = "all_android_tools",
+    jars = ["@android_tools//:all_android_tools_deploy.jar"],
+)
+
+java_library(
+    name = "jetifier_worker_wrapper",
+    visibility = ["//visibility:public"],
+    srcs = glob(["**/*.java"]),
+    deps = [
+        "//third_party/jetifier:jetifier_jars",
+        ":all_android_tools",
+    ],
+)
+
 java_binary(
     name = "jetifier",
-    main_class = "com.android.tools.build.jetifier.standalone.Main",
+    main_class = "com.github.bazelbuild.rules_jvm_external.JetifierWorkerWrapper",
     visibility = ["//visibility:public"],
     runtime_deps = [
         ":jetifier_jars",
+        ":jetifier_worker_wrapper"
     ],
 )
diff --git a/third_party/jetifier/jetifier-standalone/java/com/github/bazelbuild/rules_jvm_external/JetifierWorkerWrapper.java b/third_party/jetifier/jetifier-standalone/java/com/github/bazelbuild/rules_jvm_external/JetifierWorkerWrapper.java
new file mode 100644
index 000000000..fbc1dde63
--- /dev/null
+++ b/third_party/jetifier/jetifier-standalone/java/com/github/bazelbuild/rules_jvm_external/JetifierWorkerWrapper.java
@@ -0,0 +1,46 @@
+package com.github.bazelbuild.rules_jvm_external;
+
+import com.android.tools.build.jetifier.standalone.Main;
+import com.google.devtools.build.lib.worker.ProtoWorkerMessageProcessor;
+import com.google.devtools.build.lib.worker.WorkRequestHandler;
+
+import java.io.IOException;
+import java.io.PrintWriter;
+import java.time.Duration;
+import java.util.List;
+
+/**
+ * A Wrapper for jetifier-standalone to support multiplex worker
+ */
+public class JetifierWorkerWrapper {
+
+  public static void main(String[] args) {
+    if (args.length == 1 && args[0].equals("--persistent_worker")) {
+      WorkRequestHandler workerHandler =
+          new WorkRequestHandler.WorkRequestHandlerBuilder(
+              JetifierWorkerWrapper::jetifier,
+              System.err,
+              new ProtoWorkerMessageProcessor(System.in, System.out))
+              .setCpuUsageBeforeGc(Duration.ofSeconds(10))
+              .build();
+      int exitCode = 1;
+      try {
+        workerHandler.processRequests();
+        exitCode = 0;
+      } catch (IOException e) {
+        System.err.println(e.getMessage());
+      } finally {
+        // Prevent hanging threads from keeping the worker alive.
+        System.exit(exitCode);
+      }
+    } else {
+      Main.main(args);
+    }
+  }
+
+  private static int jetifier(List<String> args, PrintWriter pw) {
+    Main.main(args.toArray(new String[0]));
+    return 0;
+  }
+
+}
\ No newline at end of file