Fix worker and multiplex workers for DexBuilder and Desugar actions

Fixing up the DexBuilder and Desugar actions so that they correctly spawn worker or multiplexed worker actions. `--modify_execution_info` doesn't work as expected and is not additive, which results in the previous execution infos being removed.

Closes #17351.

PiperOrigin-RevId: 523696356
Change-Id: Iada7fb75df5b4d2e3ba1308110977899567f2bc2

Author: Bencodes

src/main/java/com/google/devtools/build/lib/bazel/rules/BazelRulesModule.java

@@ -307,6 +307,15 @@ public static class BuildGraveyardOptions extends OptionsBase {
         help = "No-op, will be removed soon.",
         allowMultiple = true)
     public List<String> highPriorityWorkers;
+
+    @Option(
+        name = "use_workers_with_dexbuilder",
+        defaultValue = "true",
+        documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
+        effectTags = {OptionEffectTag.EXECUTION},
+        help = "This option is deprecated and has no effect.")
+    @Deprecated
+    public boolean useWorkersWithDexbuilder;
   }
 
   /** This is where deprecated Bazel-specific options only used by the build command go to die. */

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

@@ -39,6 +39,7 @@
 import com.google.devtools.common.options.OptionEffectTag;
 import com.google.devtools.common.options.OptionMetadataTag;
 import java.util.List;
+import java.util.Locale;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.StarlarkValue;
 
@@ -163,11 +164,11 @@ public enum AndroidManifestMerger {
     ANDROID,
     FORCE_ANDROID;
 
-    public static List<String> getAttributeValues() {
+    public static ImmutableList<String> getAttributeValues() {
       return ImmutableList.of(
-          LEGACY.name().toLowerCase(),
-          ANDROID.name().toLowerCase(),
-          FORCE_ANDROID.name().toLowerCase(),
+          LEGACY.name().toLowerCase(Locale.ROOT),
+          ANDROID.name().toLowerCase(Locale.ROOT),
+          FORCE_ANDROID.name().toLowerCase(Locale.ROOT),
           getRuleAttributeDefault());
     }
 
@@ -533,14 +534,6 @@ public static class Options extends FragmentOptions {
         help = "dx flags supported in tool that groups classes for inclusion in final .dex files.")
     public List<String> dexoptsSupportedInDexSharder;
 
-    @Option(
-        name = "use_workers_with_dexbuilder",
-        defaultValue = "true",
-        documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
-        effectTags = {OptionEffectTag.EXECUTION},
-        help = "Whether dexbuilder supports being run in local worker mode.")
-    public boolean useWorkersWithDexbuilder;
-
     @Option(
         name = "experimental_android_rewrite_dexes_with_rex",
         defaultValue = "false",
@@ -905,10 +898,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 +915,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 +967,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",
@@ -1100,7 +1123,6 @@ public FragmentOptions getExec() {
       exec.dexoptsSupportedInIncrementalDexing = dexoptsSupportedInIncrementalDexing;
       exec.dexoptsSupportedInDexMerger = dexoptsSupportedInDexMerger;
       exec.dexoptsSupportedInDexSharder = dexoptsSupportedInDexSharder;
-      exec.useWorkersWithDexbuilder = useWorkersWithDexbuilder;
       exec.manifestMerger = manifestMerger;
       exec.manifestMergerOrder = manifestMergerOrder;
       exec.allowAndroidLibraryDepsWithoutSrcs = allowAndroidLibraryDepsWithoutSrcs;
@@ -1128,7 +1150,6 @@ public FragmentOptions getExec() {
   private final ImmutableList<String> targetDexoptsThatPreventIncrementalDexing;
   private final ImmutableList<String> dexoptsSupportedInDexMerger;
   private final ImmutableList<String> dexoptsSupportedInDexSharder;
-  private final boolean useWorkersWithDexbuilder;
   private final boolean desugarJava8;
   private final boolean desugarJava8Libs;
   private final boolean checkDesugarDeps;
@@ -1155,6 +1176,8 @@ public FragmentOptions getExec() {
   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;
@@ -1184,7 +1207,6 @@ public AndroidConfiguration(BuildOptions buildOptions) throws InvalidConfigurati
         ImmutableList.copyOf(options.nonIncrementalPerTargetDexopts);
     this.dexoptsSupportedInDexMerger = ImmutableList.copyOf(options.dexoptsSupportedInDexMerger);
     this.dexoptsSupportedInDexSharder = ImmutableList.copyOf(options.dexoptsSupportedInDexSharder);
-    this.useWorkersWithDexbuilder = options.useWorkersWithDexbuilder;
     this.desugarJava8 = options.desugarJava8;
     this.desugarJava8Libs = options.desugarJava8Libs;
     this.checkDesugarDeps = options.checkDesugarDeps;
@@ -1216,6 +1238,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 +1343,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 +1491,16 @@ public boolean persistentMultiplexBusyboxTools() {
     return persistentMultiplexBusyboxTools;
   }
 
+  @Override
+  public boolean persistentDexDesugar() {
+    return persistentDexDesugar;
+  }
+
+  @Override
+  public boolean persistentMultiplexDexDesugar() {
+    return persistentMultiplexDexDesugar;
+  }
+
   @Override
   public boolean incompatibleUseToolchainResolution() {
     return incompatibleUseToolchainResolution;

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

@@ -28,6 +28,7 @@
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
@@ -536,7 +537,10 @@ 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
@@ -628,8 +632,11 @@ static Artifact createDexArchiveAction(
             .useDefaultShellEnvironment()
             .setExecutable(ruleContext.getExecutablePrerequisite(dexbuilderPrereq))
             .setExecutionInfo(
-                TargetUtils.getExecutionInfo(
-                    ruleContext.getRule(), ruleContext.isAllowTagsPropagation()))
+                createDexingDesugaringExecRequirements(ruleContext)
+                    .putAll(
+                        TargetUtils.getExecutionInfo(
+                            ruleContext.getRule(), ruleContext.isAllowTagsPropagation()))
+                    .buildKeepingLast())
             // WorkerSpawnStrategy expects the last argument to be @paramfile
             .addInput(jar)
             .addOutput(dexArchive)
@@ -658,9 +665,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;
   }
@@ -670,6 +674,21 @@ 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 = {