Add support for async worker processes (facebook#2547)

* Allow WorkerProcess to be used for multiple jobs at the same time. No users yet, but next up is the process pool. * Introduce a new `async` option to worker tools. This makes buck call the tool concurrently, without waiting for responses. It's up to the tool to provide responses with the right ids, in any order. * Write some tests, fix some bugs * Add back license * Update WorkerProcessPool.java * Fix windows tests Don't crash all of buck when a worker process fails. * Code review feedback and clarifications, I hope
ifood · Lcsmarcal · Oct 20, 2020 · Jun 25, 2020 · Jun 29, 2020 · Jun 29, 2020
commit 32440f9bb50e802a7671eef5f32a19a76972c319
diff --git a/src/com/facebook/buck/features/js/JsUtil.java b/src/com/facebook/buck/features/js/JsUtil.java
@@ -112,6 +112,7 @@ static WorkerShellStep jsonWorkerShellStepAddingFlavors(
                 tool.getCommandPrefix(pathResolver),
                 tool.getEnvironment(pathResolver),
                 worker.getMaxWorkers(),
+                worker.isAsync(),
                 worker.isPersistent()
                     ? Optional.of(
                         WorkerProcessIdentity.of(buildTarget.toString(), worker.getInstanceKey()))

diff --git a/src/com/facebook/buck/rules/macros/WorkerMacroArg.java b/src/com/facebook/buck/rules/macros/WorkerMacroArg.java
@@ -133,6 +133,10 @@ public int getMaxWorkers() {
     return workerTool.getMaxWorkers();
   }
 
+  public boolean isAsync() {
+    return workerTool.isAsync();
+  }
+
   public String getJobArgs(SourcePathResolverAdapter pathResolver) {
     return Arg.stringify(arg, pathResolver).trim();
   }

diff --git a/src/com/facebook/buck/shell/DefaultWorkerToolRule.java b/src/com/facebook/buck/shell/DefaultWorkerToolRule.java
@@ -59,6 +59,7 @@ protected DefaultWorkerToolRule(
       SourcePathRuleFinder ruleFinder,
       Tool tool,
       int maxWorkers,
+      boolean isAsync,
       boolean isPersistent) {
     super(
         buildTarget,
@@ -71,6 +72,7 @@ protected DefaultWorkerToolRule(
         new DefaultWorkerTool(
             new DefaultWorkerToolDelegatingTool(tool, getSourcePathToOutput()),
             maxWorkers,
+            isAsync,
             isPersistent,
             buildTarget,
             generateNewUUID());
@@ -141,6 +143,13 @@ static class DefaultWorkerTool implements WorkerTool {
     @CustomFieldBehavior(DefaultFieldSerialization.class)
     private final Integer maxWorkers;
 
+    /**
+     * Important : Do not add this field into RuleKey. Rule key should not change in case of async
+     * variable modification.
+     */
+    @CustomFieldBehavior(DefaultFieldSerialization.class)
+    private final boolean isAsync;
+
     /**
      * Important : Do not add this field into RuleKey. Rule key should not change in case of
      * instance key modification (that is calculated during creation as random UUID).
@@ -149,11 +158,17 @@ static class DefaultWorkerTool implements WorkerTool {
     private HashCode instanceKey;
 
     DefaultWorkerTool(
-        Tool tool, int maxWorkers, boolean isPersistent, BuildTarget buildTarget, UUID uuid) {
+        Tool tool,
+        int maxWorkers,
+        boolean isAsync,
+        boolean isPersistent,
+        BuildTarget buildTarget,
+        UUID uuid) {
       this.tool = tool;
       this.maxWorkers = maxWorkers;
       this.isPersistent = isPersistent;
       this.buildTarget = buildTarget;
+      this.isAsync = isAsync;
       this.instanceKey = calculateInstanceKey(uuid);
     }
 
@@ -180,6 +195,11 @@ public int getMaxWorkers() {
       return maxWorkers;
     }
 
+    @Override
+    public boolean isAsync() {
+      return isAsync;
+    }
+
     @Override
     public boolean isPersistent() {
       return isPersistent;

diff --git a/src/com/facebook/buck/shell/GenruleBuildable.java b/src/com/facebook/buck/shell/GenruleBuildable.java
@@ -705,7 +705,8 @@ public void addEnvironmentVariables(
     androidTools.ifPresent(
         tools -> {
           environmentVariablesBuilder.put("ANDROID_HOME", tools.getAndroidSdkLocation().toString());
-          environmentVariablesBuilder.put("ANDROID_SDK_ROOT", tools.getAndroidSdkLocation().toString());
+          environmentVariablesBuilder.put(
+              "ANDROID_SDK_ROOT", tools.getAndroidSdkLocation().toString());
           environmentVariablesBuilder.put("DX", tools.getAndroidPathToDx().toString());
           environmentVariablesBuilder.put("ZIPALIGN", tools.getAndroidPathToZipalign().toString());
           environmentVariablesBuilder.put(
@@ -760,6 +761,7 @@ private static Optional<WorkerJobParams> convertToWorkerJobParams(
                   workerMacroArg.getStartupCommand(),
                   workerMacroArg.getEnvironment(),
                   workerMacroArg.getMaxWorkers(),
+                  workerMacroArg.isAsync(),
                   workerMacroArg.getPersistentWorkerKey().isPresent()
                       ? Optional.of(
                           WorkerProcessIdentity.of(

diff --git a/src/com/facebook/buck/shell/WorkerShellStep.java b/src/com/facebook/buck/shell/WorkerShellStep.java
@@ -28,13 +28,14 @@
 import com.facebook.buck.worker.WorkerJobParams;
 import com.facebook.buck.worker.WorkerJobResult;
 import com.facebook.buck.worker.WorkerProcessPool;
-import com.facebook.buck.worker.WorkerProcessPool.BorrowedWorkerProcess;
 import com.facebook.buck.worker.WorkerProcessPoolFactory;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.ExecutionException;
 import java.util.stream.Collectors;
 
 public class WorkerShellStep implements Step {
@@ -78,9 +79,14 @@ public StepExecutionResult execute(ExecutionContext context)
     WorkerJobParams paramsToUse = getWorkerJobParamsToUse(context.getPlatform());
     WorkerProcessPool pool =
         factory.getWorkerProcessPool(context, paramsToUse.getWorkerProcessParams());
-    WorkerJobResult result;
-    try (BorrowedWorkerProcess process = pool.borrowWorkerProcess()) {
-      result = process.submitAndWaitForJob(getExpandedJobArgs(context));
+    WorkerJobResult result = null;
+    try {
+      result = pool.submitJob(getExpandedJobArgs(context)).get();
+    } catch (ExecutionException e) {
+      if (e.getCause() != null) {
+        Throwables.throwIfUnchecked(e.getCause());
+      }
+      throw new RuntimeException(e);
     }
 
     Verbosity verbosity = context.getVerbosity();

diff --git a/src/com/facebook/buck/shell/WorkerTool.java b/src/com/facebook/buck/shell/WorkerTool.java
@@ -34,4 +34,6 @@ public interface WorkerTool extends AddsToRuleKey {
   boolean isPersistent();
 
   HashCode getInstanceKey();
+
+  boolean isAsync();
 }
diff --git a/src/com/facebook/buck/shell/WorkerToolDescription.java b/src/com/facebook/buck/shell/WorkerToolDescription.java
@@ -124,6 +124,8 @@ public BuildRule createBuildRule(
       builder.addEnv(e.getKey(), macrosConverter.convert(e.getValue()));
     }
 
+    boolean async = args.getSoloAsync().orElse(false);
+
     Preconditions.checkArgument(
         !(args.getMaxWorkers().isPresent() && args.getMaxWorkersPerThreadPercent().isPresent()),
         "max_workers and max_workers_per_thread_percent must not be used together.");
@@ -154,6 +156,7 @@ public BuildRule createBuildRule(
         graphBuilder,
         tool,
         maxWorkers,
+        async,
         args.getPersistent()
             .orElse(buckConfig.getBooleanValue(CONFIG_SECTION, CONFIG_PERSISTENT_KEY, false)));
   }
@@ -197,5 +200,7 @@ default Either<StringWithMacros, ImmutableList<StringWithMacros>> getArgs() {
     Optional<Integer> getMaxWorkersPerThreadPercent();
 
     Optional<Boolean> getPersistent();
+
+    Optional<Boolean> getSoloAsync();
   }
 }
diff --git a/src/com/facebook/buck/worker/BUCK b/src/com/facebook/buck/worker/BUCK
@@ -39,6 +39,8 @@ java_immutables_library(
         "WorkerProcess.java",
         "WorkerProcessCommand.java",
         "WorkerProcessPool.java",
+        "WorkerProcessPoolAsync.java",
+        "WorkerProcessPoolSync.java",
         "WorkerProcessProtocol.java",
         "WorkerProcessProtocolZero.java",
     ],