Support multiplex workers in ResourceProcessorBusyBox

Adding support for multiplex workers inside of `ResourceProcessorBusyBox` and moving it's worker implementation over to the generic work request handler.

These PRs need to land for this to work:
- #14424
- #14425
- #14427

For those not on rolling releases, the other required PRs that have already merged are:
- #14144
- #14145
- #14146

Closes #14428.

PiperOrigin-RevId: 456561596
Change-Id: I098d5a323ac6558ad0f5f8190e29f45a7a37b4d4

Author: Bencodes

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.
    *

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

@@ -882,6 +882,17 @@ public static class Options extends FragmentOptions {
         help = "Tracking flag for when busybox workers are enabled.")
     public boolean persistentBusyboxTools;
 
+    @Option(
+        name = "experimental_persistent_multiplex_busybox_tools",
+        documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
+        effectTags = {
+          OptionEffectTag.HOST_MACHINE_RESOURCE_OPTIMIZATIONS,
+          OptionEffectTag.EXECUTION,
+        },
+        defaultValue = "false",
+        help = "Tracking flag for when multiplex busybox workers are enabled.")
+    public boolean experimentalPersistentMultiplexBusyboxTools;
+
     @Option(
         name = "experimental_remove_r_classes_from_instrumentation_test_jar",
         defaultValue = "true",
@@ -1007,6 +1018,8 @@ public FragmentOptions getHost() {
       host.oneVersionEnforcementUseTransitiveJarsForBinaryUnderTest =
           oneVersionEnforcementUseTransitiveJarsForBinaryUnderTest;
       host.persistentBusyboxTools = persistentBusyboxTools;
+      host.experimentalPersistentMultiplexBusyboxTools =
+          experimentalPersistentMultiplexBusyboxTools;
 
       // Unless the build was started from an Android device, host means MAIN.
       host.configurationDistinguisher = ConfigurationDistinguisher.MAIN;
@@ -1052,6 +1065,7 @@ public FragmentOptions getHost() {
   private final boolean dataBindingUpdatedArgs;
   private final boolean dataBindingAndroidX;
   private final boolean persistentBusyboxTools;
+  private final boolean experimentalPersistentMultiplexBusyboxTools;
   private final boolean filterRJarsFromAndroidTest;
   private final boolean removeRClassesFromInstrumentationTestJar;
   private final boolean alwaysFilterDuplicateClassesFromAndroidTest;
@@ -1111,6 +1125,8 @@ public AndroidConfiguration(BuildOptions buildOptions) throws InvalidConfigurati
     this.dataBindingUpdatedArgs = options.dataBindingUpdatedArgs;
     this.dataBindingAndroidX = options.dataBindingAndroidX;
     this.persistentBusyboxTools = options.persistentBusyboxTools;
+    this.experimentalPersistentMultiplexBusyboxTools =
+        options.experimentalPersistentMultiplexBusyboxTools;
     this.filterRJarsFromAndroidTest = options.filterRJarsFromAndroidTest;
     this.removeRClassesFromInstrumentationTestJar =
         options.removeRClassesFromInstrumentationTestJar;
@@ -1362,6 +1378,11 @@ public boolean persistentBusyboxTools() {
     return persistentBusyboxTools;
   }
 
+  @Override
+  public boolean persistentMultiplexBusyboxTools() {
+    return experimentalPersistentMultiplexBusyboxTools;
+  }
+
   @Override
   public boolean incompatibleUseToolchainResolution() {
     return incompatibleUseToolchainResolution;

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

@@ -64,6 +64,7 @@ public class AndroidDataContext implements AndroidDataContextApi {
   private final FilesToRunProvider busybox;
   private final AndroidSdkProvider sdk;
   private final boolean persistentBusyboxToolsEnabled;
+  private final boolean persistentMultiplexBusyboxTools;
   private final boolean optOutOfResourcePathShortening;
   private final boolean optOutOfResourceNameObfuscation;
   private final boolean throwOnShrinkResources;
@@ -90,6 +91,7 @@ public static AndroidDataContext makeContext(RuleContext ruleContext) {
         ruleContext,
         ruleContext.getExecutablePrerequisite("$android_resources_busybox"),
         androidConfig.persistentBusyboxTools(),
+        androidConfig.persistentMultiplexBusyboxTools(),
         AndroidSdkProvider.fromRuleContext(ruleContext),
         hasExemption(ruleContext, "allow_raw_access_to_resource_paths", false),
         hasExemption(ruleContext, "allow_resource_name_obfuscation_opt_out", false),
@@ -114,6 +116,7 @@ protected AndroidDataContext(
       RuleContext ruleContext,
       FilesToRunProvider busybox,
       boolean persistentBusyboxToolsEnabled,
+      boolean persistentMultiplexBusyboxTools,
       AndroidSdkProvider sdk,
       boolean optOutOfResourcePathShortening,
       boolean optOutOfResourceNameObfuscation,
@@ -126,6 +129,7 @@ protected AndroidDataContext(
       boolean includeProguardLocationReferences,
       ImmutableMap<String, String> executionInfo) {
     this.persistentBusyboxToolsEnabled = persistentBusyboxToolsEnabled;
+    this.persistentMultiplexBusyboxTools = persistentMultiplexBusyboxTools;
     this.ruleContext = ruleContext;
     this.busybox = busybox;
     this.sdk = sdk;
@@ -222,6 +226,10 @@ public boolean isPersistentBusyboxToolsEnabled() {
     return persistentBusyboxToolsEnabled;
   }
 
+  public boolean isPersistentMultiplexBusyboxTools() {
+    return persistentMultiplexBusyboxTools;
+  }
+
   public boolean optOutOfResourcePathShortening() {
     return optOutOfResourcePathShortening;
   }

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

@@ -344,10 +344,12 @@ public void buildAndRegister(String message, String mnemonic) {
 
     if (dataContext.isPersistentBusyboxToolsEnabled()) {
       commandLine.add("--logWarnings=false");
-      spawnActionBuilder
-          .addCommandLine(commandLine.build(), WORKERS_FORCED_PARAM_FILE_INFO);
-
+      spawnActionBuilder.addCommandLine(commandLine.build(), WORKERS_FORCED_PARAM_FILE_INFO);
       executionInfo.putAll(ExecutionRequirements.WORKER_MODE_ENABLED);
+
+      if (dataContext.isPersistentMultiplexBusyboxTools()) {
+        executionInfo.putAll(ExecutionRequirements.WORKER_MULTIPLEX_MODE_ENABLED);
+      }
     } else {
       spawnActionBuilder.addCommandLine(commandLine.build(), FORCED_PARAM_FILE_INFO);
     }

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

@@ -231,6 +231,13 @@ public interface AndroidConfigurationApi extends StarlarkValue {
       documented = false)
   boolean persistentBusyboxTools();
 
+  @StarlarkMethod(
+      name = "experimental_persistent_multiplex_busybox_tools",
+      structField = true,
+      doc = "",
+      documented = false)
+  boolean persistentMultiplexBusyboxTools();
+
   @StarlarkMethod(
       name = "get_output_directory_name",
       structField = true,

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

@@ -76,6 +76,8 @@ java_library(
         ":dependency_info",
         "//src/java_tools/singlejar/java/com/google/devtools/build/singlejar:libSingleJar",
         "//src/java_tools/singlejar/java/com/google/devtools/build/zip",
+        "//src/main/java/com/google/devtools/build/lib/worker",
+        "//src/main/java/com/google/devtools/build/lib/worker:work_request_handlers",
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/protobuf:worker_protocol_java_proto",
         "//src/tools/android/java/com/google/devtools/build/android/junctions",

src/tools/android/java/com/google/devtools/build/android/ResourceProcessorBusyBox.java

@@ -14,10 +14,12 @@
 
 package com.google.devtools.build.android;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
+
 import com.google.devtools.build.android.aapt2.Aapt2Exception;
 import com.google.devtools.build.android.resources.JavaIdentifierValidator.InvalidJavaIdentifier;
-import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
-import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
+import com.google.devtools.build.lib.worker.ProtoWorkerMessageProcessor;
+import com.google.devtools.build.lib.worker.WorkRequestHandler;
 import com.google.devtools.common.options.EnumConverter;
 import com.google.devtools.common.options.Option;
 import com.google.devtools.common.options.OptionDocumentationCategory;
@@ -30,7 +32,9 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.PrintStream;
+import java.io.PrintWriter;
 import java.nio.file.FileSystems;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Properties;
@@ -183,43 +187,56 @@ private static int runPersistentWorker() throws Exception {
     PrintStream ps = new PrintStream(buf, true);
     PrintStream realStdOut = System.out;
     PrintStream realStdErr = System.err;
-    try {
-      // Redirect all stdout and stderr output for logging.
-      System.setOut(ps);
-      System.setErr(ps);
-
-      while (true) {
-        try {
-          WorkRequest request = WorkRequest.parseDelimitedFrom(System.in);
-          if (request == null) {
-            break;
-          }
-
-          int exitCode = processRequest(request.getArgumentsList());
-          ps.flush();
-
-          WorkResponse.newBuilder()
-              .setExitCode(exitCode)
-              .setRequestId(request.getRequestId())
-              .setOutput(buf.toString())
-              .build()
-              .writeDelimitedTo(realStdOut);
 
-          realStdOut.flush();
-          buf.reset();
-        } catch (IOException e) {
-          logger.severe(e.getMessage());
-          e.printStackTrace(realStdErr);
-          return 1;
-        }
-      }
+    // Redirect all stdout and stderr output for logging.
+    System.setOut(ps);
+    System.setErr(ps);
+    try {
+      WorkRequestHandler workerHandler =
+          new WorkRequestHandler.WorkRequestHandlerBuilder(
+                  new WorkRequestHandler.WorkRequestCallback(
+                      (request, pw) -> processRequest(request.getArgumentsList(), pw, buf)),
+                  realStdErr,
+                  new ProtoWorkerMessageProcessor(System.in, realStdOut))
+              .setCpuUsageBeforeGc(Duration.ofSeconds(10))
+              .build();
+      workerHandler.processRequests();
+    } catch (IOException e) {
+      logger.severe(e.getMessage());
+      e.printStackTrace(realStdErr);
+      return 1;
     } finally {
       System.setOut(realStdOut);
       System.setErr(realStdErr);
     }
     return 0;
   }
 
+  /**
+   * Processes the request for the given args and writes the captured byte array buffer to the
+   * WorkRequestHandler print writer.
+   */
+  private static int processRequest(List<String> args, PrintWriter pw, ByteArrayOutputStream buf) {
+    int exitCode;
+    try {
+      // Process the actual request and grab the exit code
+      exitCode = processRequest(args);
+    } catch (Exception e) {
+      e.printStackTrace(pw);
+      exitCode = 1;
+    } finally {
+      // Write the captured buffer to the work response. We synchronize to avoid race conditions
+      // while reading from and calling reset on the shared ByteArrayOutputStream.
+      synchronized (buf) {
+        String captured = buf.toString(UTF_8).trim();
+        buf.reset();
+        pw.print(captured);
+      }
+    }
+
+    return exitCode;
+  }
+
   private static int processRequest(List<String> args) throws Exception {
     OptionsParser optionsParser =
         OptionsParser.builder()