add flag

Author: SupSaiYaJin

src/java_tools/buildjar/java/com/google/devtools/build/java/turbine/TurbineWorkerWrapper.java

@@ -5,6 +5,7 @@
 import com.google.devtools.build.lib.worker.WorkRequestHandler;
 
 import java.io.IOException;
+import java.io.PrintStream;
 import java.io.PrintWriter;
 import java.time.Duration;
 import java.util.List;
@@ -16,10 +17,13 @@ public class TurbineWorkerWrapper {
 
   public static void main(String[] args) throws IOException {
     if (args.length == 1 && args[0].equals("--persistent_worker")) {
+      PrintStream realStdErr = System.err;
       WorkRequestHandler workerHandler =
           new WorkRequestHandler.WorkRequestHandlerBuilder(
-              TurbineWorkerWrapper::turbine,
-              System.err,
+              new WorkRequestHandler.WorkRequestCallback(
+                  (request, pw) ->
+                      turbine(request.getArgumentsList(), pw)),
+              realStdErr,
               new ProtoWorkerMessageProcessor(System.in, System.out))
               .setCpuUsageBeforeGc(Duration.ofSeconds(10))
               .build();
@@ -28,7 +32,7 @@ public static void main(String[] args) throws IOException {
         workerHandler.processRequests();
         exitCode = 0;
       } catch (IOException e) {
-        System.err.println(e.getMessage());
+        realStdErr.println(e.getMessage());
       } finally {
         // Prevent hanging threads from keeping the worker alive.
         System.exit(exitCode);
@@ -42,7 +46,7 @@ private static int turbine(List<String> args, PrintWriter pw) {
       try {
           Main.compile(args.toArray(new String[0]));
       } catch (Throwable e) {
-          System.err.println(e.getMessage());
+          pw.println(e.getMessage());
           return 1;
       }
       return 0;

src/main/java/com/google/devtools/build/lib/rules/java/JavaHeaderCompileAction.java

@@ -449,19 +449,21 @@ public void build(JavaToolchainProvider javaToolchain) {
         }
       }
 
-      ImmutableMap<String, String> executionInfo =
-          TargetUtils.getExecutionInfo(ruleContext.getRule(), ruleContext.isAllowTagsPropagation());
+      ImmutableMap.Builder<String, String> executionInfoBuilder = ImmutableMap.builder();
+      executionInfoBuilder.putAll(
+          TargetUtils.getExecutionInfo(ruleContext.getRule(), ruleContext.isAllowTagsPropagation()));
       if (javaConfiguration.inmemoryJdepsFiles()) {
-        executionInfo =
-            ImmutableMap.of(
-                ExecutionRequirements.REMOTE_EXECUTION_INLINE_OUTPUTS,
-                outputDepsProto.getExecPathString());
+        executionInfoBuilder.put(
+            ExecutionRequirements.REMOTE_EXECUTION_INLINE_OUTPUTS,
+            outputDepsProto.getExecPathString());
       }
-      executionInfo = ImmutableMap.<String, String>builder()
-          .putAll(executionInfo)
-          .putAll(ExecutionRequirements.WORKER_MODE_ENABLED)
-          .putAll(ExecutionRequirements.WORKER_MULTIPLEX_MODE_ENABLED)
-          .build();
+      if (javaToolchain.getHeaderCompilerSupportsWorkers()) {
+        executionInfoBuilder.putAll(ExecutionRequirements.WORKER_MODE_ENABLED);
+      }
+      if (javaToolchain.getHeaderCompilerSupportsMultiplexWorkers()) {
+        executionInfoBuilder.putAll(ExecutionRequirements.WORKER_MULTIPLEX_MODE_ENABLED);
+      }
+      ImmutableMap<String, String> executionInfo = executionInfoBuilder.build();
       if (useDirectClasspath) {
         NestedSet<Artifact> classpath;
         if (!directJars.isEmpty() || classpathEntries.isEmpty()) {

src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchain.java

@@ -68,6 +68,10 @@ public ConfiguredTarget create(RuleContext ruleContext)
         ruleContext.attributes().get("javac_supports_multiplex_workers", Type.BOOLEAN);
     boolean javacSupportsWorkerCancellation =
         ruleContext.attributes().get("javac_supports_worker_cancellation", Type.BOOLEAN);
+    boolean headerCompilerSupportsWorkers =
+        ruleContext.attributes().get("header_compiler_supports_workers", Type.BOOLEAN);
+    boolean headerCompilerSupportsMultiplexWorkers =
+        ruleContext.attributes().get("header_compiler_supports_multiplex_workers", Type.BOOLEAN);
     ImmutableSet<String> headerCompilerBuiltinProcessors =
         ImmutableSet.copyOf(
             ruleContext.attributes().get("header_compiler_builtin_processors", Type.STRING_LIST));
@@ -168,6 +172,8 @@ public ConfiguredTarget create(RuleContext ruleContext)
             javacSupportsWorkers,
             javacSupportsMultiplexWorkers,
             javacSupportsWorkerCancellation,
+            headerCompilerSupportsWorkers,
+            headerCompilerSupportsMultiplexWorkers,
             bootclasspath,
             tools,
             javabuilder,

src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainProvider.java

@@ -97,6 +97,8 @@ public static JavaToolchainProvider create(
       boolean javacSupportsWorkers,
       boolean javacSupportsMultiplexWorkers,
       boolean javacSupportsWorkerCancellation,
+      boolean headerCompilerSupportsWorkers,
+      boolean headerCompilerSupportsMultiplexWorkers,
       BootClassPathInfo bootclasspath,
       NestedSet<Artifact> tools,
       JavaToolchainTool javaBuilder,
@@ -149,6 +151,8 @@ public static JavaToolchainProvider create(
         javacSupportsWorkers,
         javacSupportsMultiplexWorkers,
         javacSupportsWorkerCancellation,
+        headerCompilerSupportsWorkers,
+        headerCompilerSupportsMultiplexWorkers,
         packageConfiguration,
         jacocoRunner,
         proguardAllowlister,
@@ -182,6 +186,8 @@ public static JavaToolchainProvider create(
   private final boolean javacSupportsWorkers;
   private final boolean javacSupportsMultiplexWorkers;
   private final boolean javacSupportsWorkerCancellation;
+  private final boolean headerCompilerSupportsWorkers;
+  private final boolean headerCompilerSupportsMultiplexWorkers;
   private final ImmutableList<JavaPackageConfigurationProvider> packageConfiguration;
   private final FilesToRunProvider jacocoRunner;
   private final FilesToRunProvider proguardAllowlister;
@@ -215,6 +221,8 @@ private JavaToolchainProvider(
       boolean javacSupportsWorkers,
       boolean javacSupportsMultiplexWorkers,
       boolean javacSupportsWorkerCancellation,
+      boolean headerCompilerSupportsWorkers,
+      boolean headerCompilerSupportsMultiplexWorkers,
       ImmutableList<JavaPackageConfigurationProvider> packageConfiguration,
       FilesToRunProvider jacocoRunner,
       FilesToRunProvider proguardAllowlister,
@@ -247,6 +255,8 @@ private JavaToolchainProvider(
     this.javacSupportsWorkers = javacSupportsWorkers;
     this.javacSupportsMultiplexWorkers = javacSupportsMultiplexWorkers;
     this.javacSupportsWorkerCancellation = javacSupportsWorkerCancellation;
+    this.headerCompilerSupportsWorkers = headerCompilerSupportsWorkers;
+    this.headerCompilerSupportsMultiplexWorkers = headerCompilerSupportsMultiplexWorkers;
     this.packageConfiguration = packageConfiguration;
     this.jacocoRunner = jacocoRunner;
     this.proguardAllowlister = proguardAllowlister;
@@ -419,6 +429,16 @@ public boolean getJavacSupportsMultiplexWorkers() {
     return javacSupportsMultiplexWorkers;
   }
 
+  /** Returns whether JavaHeaderCompiler supports running as a persistent worker or not. */
+  public boolean getHeaderCompilerSupportsWorkers() {
+    return headerCompilerSupportsWorkers;
+  }
+
+  /** Returns whether JavaHeaderCompiler supports running persistent workers in multiplex mode */
+  public boolean getHeaderCompilerSupportsMultiplexWorkers() {
+    return headerCompilerSupportsMultiplexWorkers;
+  }
+
   /** Returns whether JavaBuilders supports running persistent workers with cancellation */
   public boolean getJavacSupportsWorkerCancellation() {
     return javacSupportsWorkerCancellation;

src/main/java/com/google/devtools/build/lib/rules/java/JavaToolchainRule.java

@@ -129,6 +129,14 @@ The Java target version (e.g., '6' or '7'). It specifies for which Java runtime
         True if JavaBuilder supports cancellation of persistent workers, false if it doesn't.
         <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
         .add(attr("javac_supports_worker_cancellation", BOOLEAN).value(true))
+        /* <!-- #BLAZE_RULE(java_toolchain).ATTRIBUTE(header_compiler_supports_workers) -->
+        True if JavaHeaderCompiler supports running as a persistent worker, false if it doesn't.
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(attr("header_compiler_supports_workers", BOOLEAN).value(false))
+        /* <!-- #BLAZE_RULE(java_toolchain).ATTRIBUTE(header_compiler_supports_multiplex_workers) -->
+        True if JavaHeaderCompiler supports running as a multiplex persistent worker, false if it doesn't.
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(attr("header_compiler_supports_multiplex_workers", BOOLEAN).value(false))
         /* <!-- #BLAZE_RULE(java_toolchain).ATTRIBUTE(tools) -->
         Labels of tools available for label-expansion in jvm_opts.
         <!-- #END_BLAZE_RULE.ATTRIBUTE --> */