From ac96cd37eb46f8f7a4ddd9680ffa24c3e35b6292 Mon Sep 17 00:00:00 2001
From: Benjamin Lee <ben@ben.cm>
Date: Fri, 22 Apr 2022 12:42:26 -0700
Subject: [PATCH] Add multiplex worker support to dexbuilder

---
 .../lib/actions/ExecutionRequirements.java      |  3 +++
 .../lib/rules/android/AndroidConfiguration.java | 17 +++++++++++++++++
 .../lib/rules/android/DexArchiveAspect.java     |  3 +++
 .../android/AndroidConfigurationApi.java        |  7 +++++++
 4 files changed, 30 insertions(+)

diff --git a/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java b/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java
index ffcc82d090ece7..db9f08e201c245 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java
@@ -176,6 +176,9 @@ public enum WorkerProtocolFormat {
   public static final ImmutableMap<String, String> WORKER_MODE_ENABLED =
       ImmutableMap.of(SUPPORTS_WORKERS, "1");
 
+  public static final ImmutableMap<String, String> WORKER_MULTIPLEX_MODE_ENABLED =
+      ImmutableMap.of(SUPPORTS_MULTIPLEX_WORKERS, "1");
+
   /**
    * Requires local execution without sandboxing for a spawn.
    *
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java
index ec9cc583e19b39..fb9861c47ea408 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidConfiguration.java
@@ -552,6 +552,14 @@ public static class Options extends FragmentOptions {
         help = "Whether dexbuilder supports being run in local worker mode.")
     public boolean useWorkersWithDexbuilder;
 
+    @Option(
+        name = "use_multiplex_workers_with_dexbuilder",
+        defaultValue = "false",
+        documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
+        effectTags = {OptionEffectTag.EXECUTION},
+        help = "Whether dexbuilder supports being run in local multiplex worker mode.")
+    public boolean useMultiplexWorkersWithDexbuilder;
+
     @Option(
         name = "experimental_android_rewrite_dexes_with_rex",
         defaultValue = "false",
@@ -1016,6 +1024,7 @@ public FragmentOptions getHost() {
       host.dexoptsSupportedInDexMerger = dexoptsSupportedInDexMerger;
       host.dexoptsSupportedInDexSharder = dexoptsSupportedInDexSharder;
       host.useWorkersWithDexbuilder = useWorkersWithDexbuilder;
+      host.useMultiplexWorkersWithDexbuilder = useMultiplexWorkersWithDexbuilder;
       host.manifestMerger = manifestMerger;
       host.manifestMergerOrder = manifestMergerOrder;
       host.allowAndroidLibraryDepsWithoutSrcs = allowAndroidLibraryDepsWithoutSrcs;
@@ -1042,6 +1051,7 @@ public FragmentOptions getHost() {
   private final ImmutableList<String> dexoptsSupportedInDexMerger;
   private final ImmutableList<String> dexoptsSupportedInDexSharder;
   private final boolean useWorkersWithDexbuilder;
+  private final boolean useMultiplexWorkersWithDexbuilder;
   private final boolean desugarJava8;
   private final boolean desugarJava8Libs;
   private final boolean checkDesugarDeps;
@@ -1096,6 +1106,7 @@ public AndroidConfiguration(BuildOptions buildOptions) throws InvalidConfigurati
     this.dexoptsSupportedInDexMerger = ImmutableList.copyOf(options.dexoptsSupportedInDexMerger);
     this.dexoptsSupportedInDexSharder = ImmutableList.copyOf(options.dexoptsSupportedInDexSharder);
     this.useWorkersWithDexbuilder = options.useWorkersWithDexbuilder;
+    this.useMultiplexWorkersWithDexbuilder = options.useMultiplexWorkersWithDexbuilder;
     this.desugarJava8 = options.desugarJava8;
     this.desugarJava8Libs = options.desugarJava8Libs;
     this.checkDesugarDeps = options.checkDesugarDeps;
@@ -1234,6 +1245,12 @@ public boolean useWorkersWithDexbuilder() {
     return useWorkersWithDexbuilder;
   }
 
+  /** Whether to assume the dexbuilder tool supports local worker mode. */
+  @Override
+  public boolean useMultiplexWorkersWithDexbuilder() {
+    return useMultiplexWorkersWithDexbuilder;
+  }
+
   @Override
   public boolean desugarJava8() {
     return desugarJava8;
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java b/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java
index f88fcd12a9f480..9f6f0b43ee6b86 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/DexArchiveAspect.java
@@ -642,6 +642,9 @@ static Artifact createDexArchiveAction(
     if (getAndroidConfig(ruleContext).useWorkersWithDexbuilder()) {
       dexbuilder.setExecutionInfo(ExecutionRequirements.WORKER_MODE_ENABLED);
     }
+    if (getAndroidConfig(ruleContext).useMultiplexWorkersWithDexbuilder()) {
+      dexbuilder.setExecutionInfo(ExecutionRequirements.WORKER_MULTIPLEX_MODE_ENABLED);
+    }
     ruleContext.registerAction(dexbuilder.build(ruleContext));
     return dexArchive;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/android/AndroidConfigurationApi.java b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/android/AndroidConfigurationApi.java
index 197d3c16a4e5a9..331a3404ccb1a2 100644
--- a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/android/AndroidConfigurationApi.java
+++ b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/android/AndroidConfigurationApi.java
@@ -101,6 +101,13 @@ public interface AndroidConfigurationApi extends StarlarkValue {
       documented = false)
   boolean useWorkersWithDexbuilder();
 
+  @StarlarkMethod(
+      name = "use_multiplex_workers_with_dexbuilder",
+      structField = true,
+      doc = "",
+      documented = false)
+  boolean useMultiplexWorkersWithDexbuilder();
+
   @StarlarkMethod(name = "desugar_java8", structField = true, doc = "", documented = false)
   boolean desugarJava8();