Fix worker and multiplex workers for DexBuilder and Desugar actions

Author: Bencodes

src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java

@@ -905,10 +905,11 @@ public static class Options extends FragmentOptions {
         },
         help = "Enable persistent Android dex and desugar actions by using workers.",
         expansion = {
+          "--internal_persistent_android_dex_desugar",
           "--strategy=Desugar=worker",
           "--strategy=DexBuilder=worker",
         })
-    public Void persistentDexDesugar;
+    public Void persistentAndroidDexDesugar;
 
     @Option(
         name = "persistent_multiplex_android_dex_desugar",
@@ -921,10 +922,9 @@ public static class Options extends FragmentOptions {
         help = "Enable persistent multiplexed Android dex and desugar actions by using workers.",
         expansion = {
           "--persistent_android_dex_desugar",
-          "--modify_execution_info=Desugar=+supports-multiplex-workers",
-          "--modify_execution_info=DexBuilder=+supports-multiplex-workers",
+          "--internal_persistent_multiplex_android_dex_desugar",
         })
-    public Void persistentMultiplexDexDesugar;
+    public Void persistentMultiplexAndroidDexDesugar;
 
     @Option(
         name = "persistent_multiplex_android_tools",
@@ -974,6 +974,36 @@ public static class Options extends FragmentOptions {
         help = "Tracking flag for when multiplexed busybox workers are enabled.")
     public boolean persistentMultiplexBusyboxTools;
 
+    /**
+     * We use this option to decide when to enable workers for busybox tools. This flag is also a
+     * guard against enabling workers using nothing but --persistent_android_resource_processor.
+     *
+     * <p>Consequently, we use this option to decide between param files or regular command line
+     * parameters. If we're not using workers or on Windows, there's no need to always use param
+     * files for I/O performance reasons.
+     */
+    @Option(
+        name = "internal_persistent_android_dex_desugar",
+        documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
+        effectTags = {
+          OptionEffectTag.HOST_MACHINE_RESOURCE_OPTIMIZATIONS,
+          OptionEffectTag.EXECUTION,
+        },
+        defaultValue = "false",
+        help = "Tracking flag for when dexing and desugaring workers are enabled.")
+    public boolean persistentDexDesugar;
+
+    @Option(
+        name = "internal_persistent_multiplex_android_dex_desugar",
+        documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
+        effectTags = {
+          OptionEffectTag.HOST_MACHINE_RESOURCE_OPTIMIZATIONS,
+          OptionEffectTag.EXECUTION,
+        },
+        defaultValue = "false",
+        help = "Tracking flag for when multiplexed dexing and desugaring workers are enabled.")
+    public boolean persistentMultiplexDexDesugar;
+
     @Option(
         name = "experimental_remove_r_classes_from_instrumentation_test_jar",
         defaultValue = "true",
@@ -1155,6 +1185,8 @@ public FragmentOptions getHost() {
   private final boolean dataBindingAndroidX;
   private final boolean persistentBusyboxTools;
   private final boolean persistentMultiplexBusyboxTools;
+  private final boolean persistentDexDesugar;
+  private final boolean persistentMultiplexDexDesugar;
   private final boolean filterRJarsFromAndroidTest;
   private final boolean removeRClassesFromInstrumentationTestJar;
   private final boolean alwaysFilterDuplicateClassesFromAndroidTest;
@@ -1216,6 +1248,8 @@ public AndroidConfiguration(BuildOptions buildOptions) throws InvalidConfigurati
     this.dataBindingAndroidX = options.dataBindingAndroidX;
     this.persistentBusyboxTools = options.persistentBusyboxTools;
     this.persistentMultiplexBusyboxTools = options.persistentMultiplexBusyboxTools;
+    this.persistentDexDesugar = options.persistentDexDesugar;
+    this.persistentMultiplexDexDesugar = options.persistentMultiplexDexDesugar;
     this.filterRJarsFromAndroidTest = options.filterRJarsFromAndroidTest;
     this.removeRClassesFromInstrumentationTestJar =
         options.removeRClassesFromInstrumentationTestJar;
@@ -1319,12 +1353,6 @@ public ImmutableList<String> getTargetDexoptsThatPreventIncrementalDexing() {
     return targetDexoptsThatPreventIncrementalDexing;
   }
 
-  /** Whether to assume the dexbuilder tool supports local worker mode. */
-  @Override
-  public boolean useWorkersWithDexbuilder() {
-    return useWorkersWithDexbuilder;
-  }
-
   @Override
   public boolean desugarJava8() {
     return desugarJava8;
@@ -1473,6 +1501,16 @@ public boolean persistentMultiplexBusyboxTools() {
     return persistentMultiplexBusyboxTools;
   }
 
+  @Override
+  public boolean persistentDexDesugar() {
+    return persistentDexDesugar;
+  }
+
+  @Override
+  public boolean persistentMultiplexDexDesugar() {
+    return persistentMultiplexDexDesugar;
+  }
+
   @Override
   public boolean incompatibleUseToolchainResolution() {
     return incompatibleUseToolchainResolution;
@@ -1500,6 +1538,7 @@ public boolean alwaysFilterDuplicateClassesFromAndroidTest() {
     return alwaysFilterDuplicateClassesFromAndroidTest;
   }
 
+  @Override
   public boolean filterLibraryJarWithProgramJar() {
     return filterLibraryJarWithProgramJar;
   }

src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java

@@ -27,6 +27,7 @@
 import com.google.common.base.Joiner;
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
@@ -526,7 +527,9 @@ private static Artifact createDesugarAction(
             .addOutput(result)
             .setMnemonic("Desugar")
             .setProgressMessage("Desugaring %s for Android", jar.prettyPrint())
-            .setExecutionInfo(ExecutionRequirements.WORKER_MODE_ENABLED);
+            .setExecutionInfo(createDexingDesugaringExecRequirements(ruleContext)
+                .putAll(ExecutionRequirements.WORKER_MODE_ENABLED)
+                .buildKeepingLast());
 
     // SpawnAction.Builder.build() is documented as being safe for re-use. So we can call build here
     // to get the action's inputs for vetting path stripping safety, then call it again later to
@@ -617,9 +620,9 @@ static Artifact createDexArchiveAction(
         new SpawnAction.Builder()
             .useDefaultShellEnvironment()
             .setExecutable(ruleContext.getExecutablePrerequisite(dexbuilderPrereq))
-            .setExecutionInfo(
-                TargetUtils.getExecutionInfo(
-                    ruleContext.getRule(), ruleContext.isAllowTagsPropagation()))
+            .setExecutionInfo(createDexingDesugaringExecRequirements(ruleContext)
+                .putAll(TargetUtils.getExecutionInfo(ruleContext.getRule(), ruleContext.isAllowTagsPropagation()))
+                .buildKeepingLast())
             // WorkerSpawnStrategy expects the last argument to be @paramfile
             .addInput(jar)
             .addOutput(dexArchive)
@@ -648,9 +651,6 @@ static Artifact createDexArchiveAction(
     dexbuilder
         .addCommandLine(args.build(), ParamFileInfo.builder(UNQUOTED).setUseAlways(true).build())
         .stripOutputPaths(stripOutputPaths);
-    if (getAndroidConfig(ruleContext).useWorkersWithDexbuilder()) {
-      dexbuilder.setExecutionInfo(ExecutionRequirements.WORKER_MODE_ENABLED);
-    }
     ruleContext.registerAction(dexbuilder.build(ruleContext));
     return dexArchive;
   }
@@ -660,6 +660,22 @@ private static Set<Set<String>> aspectDexopts(RuleContext ruleContext) {
         normalizeDexopts(getAndroidConfig(ruleContext).getDexoptsSupportedInIncrementalDexing()));
   }
 
+  /**
+   * Creates the execution requires for the DexBuilder and Desugar actions
+   */
+  private static ImmutableMap.Builder<String, String> createDexingDesugaringExecRequirements(RuleContext ruleContext) {
+    final ImmutableMap.Builder<String, String> executionInfo = ImmutableMap.builder();
+    AndroidConfiguration androidConfiguration = getAndroidConfig(ruleContext);
+    if (androidConfiguration.persistentDexDesugar()) {
+      executionInfo.putAll(ExecutionRequirements.WORKER_MODE_ENABLED);
+      if (androidConfiguration.persistentMultiplexDexDesugar()) {
+        executionInfo.putAll(ExecutionRequirements.WORKER_MULTIPLEX_MODE_ENABLED);
+      }
+    }
+
+    return executionInfo;
+  }
+
   /**
    * Derives options to use in incremental dexing actions from the given context and dx flags, where
    * the latter typically come from a {@code dexopts} attribute on a top-level target. This method

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/android/AndroidConfigurationApi.java

@@ -94,13 +94,6 @@ public interface AndroidConfigurationApi extends StarlarkValue {
       documented = false)
   ImmutableList<String> getTargetDexoptsThatPreventIncrementalDexing();
 
-  @StarlarkMethod(
-      name = "use_workers_with_dexbuilder",
-      structField = true,
-      doc = "",
-      documented = false)
-  boolean useWorkersWithDexbuilder();
-
   @StarlarkMethod(name = "desugar_java8", structField = true, doc = "", documented = false)
   boolean desugarJava8();
 
@@ -238,6 +231,20 @@ public interface AndroidConfigurationApi extends StarlarkValue {
       documented = false)
   boolean persistentMultiplexBusyboxTools();
 
+  @StarlarkMethod(
+      name = "persistent_android_dex_desugar",
+      structField = true,
+      doc = "",
+      documented = false)
+  boolean persistentDexDesugar();
+
+  @StarlarkMethod(
+      name = "persistent_multiplex_android_dex_desugar",
+      structField = true,
+      doc = "",
+      documented = false)
+  boolean persistentMultiplexDexDesugar();
+
   @StarlarkMethod(
       name = "get_output_directory_name",
       structField = true,

src/test/shell/bazel/android/desugarer_integration_test.sh

@@ -124,10 +124,24 @@ function test_java_8_android_binary_worker_strategy() {
   setup_android_sdk_support
   create_java_8_android_binary
 
-  bazel build \
-   --strategy=Desugar=worker \
-   --desugar_for_android //java/bazel:bin \
-      || fail "build failed"
+  assert_build //java/bazel:bin \
+    --persistent_android_dex_desugar \
+    --worker_verbose &> $TEST_log
+  expect_log "Created new non-sandboxed Desugar worker (id [0-9]\+)"
+  expect_log "Created new non-sandboxed DexBuilder worker (id [0-9]\+)"
+}
+
+function test_java_8_android_binary_multiplex_worker_strategy() {
+  create_new_workspace
+  setup_android_sdk_support
+  create_java_8_android_binary
+
+  assert_build //java/bazel:bin \
+    --experimental_worker_multiplex \
+    --persistent_multiplex_android_dex_desugar \
+    --worker_verbose &> $TEST_log
+  expect_log "Created new non-sandboxed Desugar multiplex-worker (id [0-9]\+)"
+  expect_log "Created new non-sandboxed DexBuilder multiplex-worker (id [0-9]\+)"
 }
 
 run_suite "Android desugarer integration tests"

tools/android/BUILD.tools

@@ -47,8 +47,6 @@ java_binary(
     runtime_deps = ["//src/tools/android/java/com/google/devtools/build/android/r8"],
 )
 
-# NOTE: d8 dex builder doesn't support the persistent worker mode at the moment. To use this config,
-# without a build error, --nouse_workers_with_dexbuilder flag must also be specified.
 config_setting(
     name = "d8_incremental_dexing",
     values = {