Implement DesugarDexShardingAction binary in Java for increased speed

Wrapping all JVM invocations (desugar, dex) inside of a Java binary saves on JVM startups and reuses the JIT cache, thus improving overall performance, and enables persistent worker mode for Bazel.

#292

PiperOrigin-RevId: 745679107
Change-Id: I099e06115e2ca0105246c04fd135cda99ca8ea92

Author: ted-xie

mobile_install/dependency_map.bzl

@@ -17,6 +17,9 @@ versioned_deps = struct(
     android_kit = struct(
         head = "//src/tools/ak",
     ),
+    desugar_dex_sharding = struct(
+        head = "//src/tools/java/com/google/devtools/build/android:DesugarDexShardingAction_deploy.jar",
+    ),
     bootstraper = struct(
         head = "//tools/android:gen_fail",
     ),

mobile_install/tools_common.bzl

@@ -121,6 +121,12 @@ TOOL_ATTRS = dict(
         cfg = "exec",
         executable = True,
     ),
+    _desugar_dex_sharding = attr.label(
+        default = versioned_deps.desugar_dex_sharding.head,
+        allow_files = True,
+        cfg = "exec",
+        executable = True,
+    ),
     _deploy = attr.label(
         default = versioned_deps.deploy.head,
         allow_files = True,

mobile_install/transform.bzl

@@ -42,8 +42,7 @@ def dex(
         data,
         deps = constants.EMPTY_LIST,
         num_shards = None,
-        create_file = _declare_file,
-        desugar = True):
+        create_file = _declare_file):
     """Dex a list of Jars.
 
     Args:
@@ -65,7 +64,6 @@ def dex(
 
         Returns:
           A File.
-      desugar: A boolean that determines whether to apply desugaring.
 
     Returns:
       A list of tuples where each entry contains the originating Jar path and
@@ -88,7 +86,7 @@ def dex(
                 dirname + "/" + str(i) + ".zip",
                 sibling = jar,
             ))
-        utils.dex(ctx, jar, out_dex_shards, deps, desugar)
+        utils.dex(ctx, jar, out_dex_shards, deps)
         dex_files.append(out_dex_shards)
     return dex_files
 

mobile_install/utils.bzl

@@ -146,7 +146,7 @@ def host_jvm_path(ctx):
     """
     return str(ctx.attr._host_java_runtime[java_common.JavaRuntimeInfo].java_executable_exec_path)
 
-def dex(ctx, jar, out_dex_shards, deps = None, desugar = True):
+def dex(ctx, jar, out_dex_shards, deps = None):
     """Desugar, dex and shard a Jar.
 
     Args:
@@ -156,32 +156,41 @@ def dex(ctx, jar, out_dex_shards, deps = None, desugar = True):
         is given, will shard the Jar to Dex across all given files in a
         deterministic manner.
       deps: The list of dependencies for the Jar being desugared.
-      desugar: A boolean that determines whether to apply desugaring.
     """
     min_sdk = _min_sdk_version.get(ctx)
     args = ctx.actions.args()
     args.use_param_file(param_file_arg = "-flagfile=%s", use_always = True)
-    if desugar:
-        args.add("-desugar", ctx.executable._desugar_java8)
-        args.add("-android_jar", first(ctx.files._android_sdk))
-        if deps:
-            args.add_joined("-classpath", deps, join_with = ",")
-        args.add("-desugar_core_libs", "True")
+
+    args.add("-android_jar", first(ctx.files._android_sdk))
+    if deps:
+        args.add_joined("-classpath", deps, join_with = ",")
+    args.add("-desugar_core_libs", "True")
 
     # Unconditionally add -desugared_lib_config. This matches the behavior of tools/android/d8_desugar.sh.
     args.add("-desugared_lib_config", ctx.file._desugared_lib_config)
 
-    args.add("-dexbuilder", ctx.executable._dexbuilder)
     args.add("-min_sdk_version", min_sdk)
     args.add("-in", jar)
     args.add_joined("-out", out_dex_shards, join_with = ",")
 
+    java_runtime = ctx.attr._mi_host_javabase[java_common.JavaRuntimeInfo]
+    java = java_runtime.java_executable_exec_path
+
+    # Performance-related JVM flags for the desugar tool.
+    jvm_flags = [
+        # b/71513487
+        "-XX:+TieredCompilation",
+        "-XX:TieredStopAtLevel=1",
+        "-Xms8g",
+        "-Xmx8g",
+    ]
+
     ctx.actions.run(
-        executable = ctx.executable._android_kit,
-        arguments = ["dex", args],
-        tools = [ctx.executable._desugar_java8, ctx.executable._dexbuilder],
+        executable = java,
+        tools = [ctx.executable._desugar_dex_sharding],
+        arguments = jvm_flags + ["-jar", ctx.executable._desugar_dex_sharding.path, args],
         inputs = depset(
-            ctx.files._android_sdk + [jar, ctx.file._desugared_lib_config],
+            ctx.files._android_sdk + ctx.files._mi_host_javabase + [jar, ctx.file._desugared_lib_config],
             transitive = [deps] if deps else [],
         ),
         outputs = out_dex_shards,

src/tools/java/com/google/devtools/build/android/BUILD

@@ -106,6 +106,8 @@ java_library(
             "CompatShellQuotedParamsFileProcessor.java",
             "AndroidDataBindingWrapper.java",
             "ProcessXmlOptionsWrapper.java",
+            "DesugarDexShardingAction.java",
+            "Fnv1a32bHash.java",
         ],
     ),
     javacopts = [
@@ -141,6 +143,29 @@ java_library(
     ],
 )
 
+java_library(
+    name = "desugar_dex_sharding_lib",
+    srcs = [
+        "DesugarDexShardingAction.java",
+        "Fnv1a32bHash.java",
+    ],
+    deps = [
+        ":android_builder_lib",
+        ":zip_jar",
+        "//src/tools/java/com/google/devtools/build/android/r8",
+        "//src/tools/java/com/google/devtools/build/android/r8:constants",
+        "@rules_android_maven//:com_beust_jcommander",
+        "@rules_android_maven//:com_google_guava_guava",
+    ],
+)
+
+java_binary(
+    name = "DesugarDexShardingAction",
+    main_class = "com.google.devtools.build.android.DesugarDexShardingAction",
+    visibility = ["//visibility:public"],
+    runtime_deps = [":desugar_dex_sharding_lib"],
+)
+
 java_library(
     name = "dependency_info",
     srcs = ["DependencyInfo.java"],
@@ -155,8 +180,38 @@ run_singlejar(
     srcs = ["@android_tools//:all_android_tools_deploy.jar"],
     out = "android_common_exec.jar",
     include_prefixes = [
+        "com/android/AndroidProjectTypes",
+        "com/android/ProgressManagerAdapter",
+        "com/android/SdkConstants",
+        "com/android/Version",
+        "com/android/annotations",
+        "com/android/build/gradle",
+        "com/android/builder",
+        "com/android/ide",
+        "com/android/io",
+        "com/android/manifmerger",
+        "com/android/repository",
+        "com/android/resources",
+        "com/android/signflinger",
+        "com/android/tools/analytics/library/protos",
+        "com/android/tools/analytics/library/shared",
+        "com/android/tools/analytics/library/tracker",
+        "com/android/tools/annotations",
+        "com/android/tools/build/apksig",
+        "com/android/tools/build/apkzlib",
+        "com/android/tools/build/builder",
+        "com/android/tools/build/builder/model",
+        "com/android/tools/build/manifest/merger",
+        "com/android/tools/common",
+        "com/android/tools/ddms/ddmlib",
+        "com/android/tools/dvlib",
+        "com/android/tools/layoutlib/layoutlib/api",
+        "com/android/tools/repository",
+        "com/android/tools/sdk/common",
+        "com/android/tools/sdklib",
+        "com/android/utils",
+        "com/android/zipflinger",
         "com/google/wireless/",
-        "com/android/",
     ],
 )