Experimental support for remote persistent workers

Add a new --experimental_remote_mark_tool_inputs flag, which makes Bazel tag
tool inputs when executing actions remotely, and also adds a tools input key
to the platform proto sent as part of the remote execution request.

This allows a remote execution system to implement persistent workers, i.e.,
to keep worker processes around and reuse them for subsequent actions. In a
trivial example, this improves build performance by ~3x.

We use "persistentWorkerKey" for the platform property, with the value being
a hash of the tool inputs, and "bazel_tool_input" as the node property name,
with an empty string as value (this is just a boolean tag).

Implements bazelbuild#10091.

Change-Id: Iccb36081fee399855be7c487c2d4091cb36f8df3

src/main/java/com/google/devtools/build/lib/analysis/platform/PlatformUtils.java

@@ -17,6 +17,7 @@
 import build.bazel.remote.execution.v2.Platform;
 import build.bazel.remote.execution.v2.Platform.Property;
 import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Ordering;
 import com.google.devtools.build.lib.actions.Spawn;
@@ -64,41 +65,49 @@ public static Platform buildPlatformProto(Map<String, String> executionPropertie
   @Nullable
   public static Platform getPlatformProto(Spawn spawn, @Nullable RemoteOptions remoteOptions)
       throws UserExecException {
+    return getPlatformProto(spawn, remoteOptions, ImmutableMap.of());
+  }
+
+  @Nullable
+  public static Platform getPlatformProto(
+      Spawn spawn, @Nullable RemoteOptions remoteOptions, Map<String, String> additionalProperties)
+      throws UserExecException {
     SortedMap<String, String> defaultExecProperties =
         remoteOptions != null
             ? remoteOptions.getRemoteDefaultExecProperties()
             : ImmutableSortedMap.of();
 
     if (spawn.getExecutionPlatform() == null
         && spawn.getCombinedExecProperties().isEmpty()
-        && defaultExecProperties.isEmpty()) {
+        && defaultExecProperties.isEmpty()
+        && additionalProperties.isEmpty()) {
       return null;
     }
 
-    Platform.Builder platformBuilder = Platform.newBuilder();
-
+    Map<String, String> properties;
     if (!spawn.getCombinedExecProperties().isEmpty()) {
-      Map<String, String> combinedExecProperties;
       // Apply default exec properties if the execution platform does not already set
       // exec_properties
       if (spawn.getExecutionPlatform() == null
           || spawn.getExecutionPlatform().execProperties().isEmpty()) {
-        combinedExecProperties = new HashMap<>();
-        combinedExecProperties.putAll(defaultExecProperties);
-        combinedExecProperties.putAll(spawn.getCombinedExecProperties());
+        properties = new HashMap<>();
+        properties.putAll(defaultExecProperties);
+        properties.putAll(spawn.getCombinedExecProperties());
       } else {
-        combinedExecProperties = spawn.getCombinedExecProperties();
-      }
-
-      for (Map.Entry<String, String> entry : combinedExecProperties.entrySet()) {
-        platformBuilder.addPropertiesBuilder().setName(entry.getKey()).setValue(entry.getValue());
+        properties = spawn.getCombinedExecProperties();
       }
     } else if (spawn.getExecutionPlatform() != null
         && !Strings.isNullOrEmpty(spawn.getExecutionPlatform().remoteExecutionProperties())) {
-      // Try and get the platform info from the execution properties.
+      properties = new HashMap<>();
+      // Try and get the platform info from the execution properties. This is pretty inefficient; it
+      // would be better to store the parsed properties instead of the String text proto.
       try {
+        Platform.Builder platformBuilder = Platform.newBuilder();
         TextFormat.getParser()
             .merge(spawn.getExecutionPlatform().remoteExecutionProperties(), platformBuilder);
+        for (Property property : platformBuilder.getPropertiesList()) {
+          properties.put(property.getName(), property.getValue());
+        }
       } catch (ParseException e) {
         String message =
             String.format(
@@ -108,12 +117,23 @@ public static Platform getPlatformProto(Spawn spawn, @Nullable RemoteOptions rem
             e, createFailureDetail(message, Code.INVALID_REMOTE_EXECUTION_PROPERTIES));
       }
     } else {
-      for (Map.Entry<String, String> property : defaultExecProperties.entrySet()) {
-        platformBuilder.addProperties(
-            Property.newBuilder().setName(property.getKey()).setValue(property.getValue()).build());
+      properties = defaultExecProperties;
+    }
+
+    if (!additionalProperties.isEmpty()) {
+      if (properties.isEmpty()) {
+        properties = additionalProperties;
+      } else {
+        // Merge the two maps.
+        properties = new HashMap<>(properties);
+        properties.putAll(additionalProperties);
       }
     }
 
+    Platform.Builder platformBuilder = Platform.newBuilder();
+    for (Map.Entry<String, String> entry : properties.entrySet()) {
+      platformBuilder.addPropertiesBuilder().setName(entry.getKey()).setValue(entry.getValue());
+    }
     sortPlatformProperties(platformBuilder);
     return platformBuilder.build();
   }

src/main/java/com/google/devtools/build/lib/exec/local/LocalEnvProvider.java

@@ -24,6 +24,7 @@
  * probably should not exist, but is currently necessary for our local MacOS support.
  */
 public interface LocalEnvProvider {
+  LocalEnvProvider NOOP = (env, binTools, fallbackTmpDir) -> ImmutableMap.copyOf(env);
 
   /**
    * Creates a local environment provider for the current OS.

src/main/java/com/google/devtools/build/lib/remote/BUILD

@@ -74,6 +74,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/exec:spawn_input_expander",
         "//src/main/java/com/google/devtools/build/lib/exec:spawn_runner",
         "//src/main/java/com/google/devtools/build/lib/exec:spawn_strategy_registry",
+        "//src/main/java/com/google/devtools/build/lib/exec/local",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/packages/semantics",
         "//src/main/java/com/google/devtools/build/lib/profiler",
@@ -90,6 +91,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/sandbox",
         "//src/main/java/com/google/devtools/build/lib/skyframe:mutable_supplier",
         "//src/main/java/com/google/devtools/build/lib/skyframe:tree_artifact_value",
+        "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:abrupt_exit_exception",
         "//src/main/java/com/google/devtools/build/lib/util:detailed_exit_code",
         "//src/main/java/com/google/devtools/build/lib/util:exit_code",
@@ -98,6 +100,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:output_service",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
+        "//src/main/java/com/google/devtools/build/lib/worker",
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/protobuf:failure_details_java_proto",
         "//third_party:auth",

src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java

@@ -75,14 +75,14 @@
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnResult;
 import com.google.devtools.build.lib.actions.Spawns;
-import com.google.devtools.build.lib.actions.UserExecException;
 import com.google.devtools.build.lib.actions.cache.MetadataInjector;
 import com.google.devtools.build.lib.analysis.platform.PlatformUtils;
 import com.google.devtools.build.lib.buildtool.buildevent.BuildInterruptedEvent;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.Reporter;
 import com.google.devtools.build.lib.exec.SpawnInputExpander.InputWalker;
 import com.google.devtools.build.lib.exec.SpawnRunner.SpawnExecutionContext;
+import com.google.devtools.build.lib.exec.local.LocalEnvProvider;
 import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.ProfilerTask;
 import com.google.devtools.build.lib.profiler.SilentCloseable;
@@ -108,10 +108,15 @@
 import com.google.devtools.build.lib.remote.util.Utils.InMemoryOutput;
 import com.google.devtools.build.lib.server.FailureDetails.RemoteExecution;
 import com.google.devtools.build.lib.skyframe.TreeArtifactValue;
+import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.devtools.build.lib.util.io.FileOutErr;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.devtools.build.lib.worker.WorkerKey;
+import com.google.devtools.build.lib.worker.WorkerOptions;
+import com.google.devtools.build.lib.worker.WorkerParser;
+import com.google.devtools.common.options.Options;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.ExtensionRegistry;
 import com.google.protobuf.Message;
@@ -365,13 +370,14 @@ private SortedMap<PathFragment, ActionInput> buildOutputDirMap(Spawn spawn) {
     return outputDirMap;
   }
 
-  private MerkleTree buildInputMerkleTree(Spawn spawn, SpawnExecutionContext context)
+  private MerkleTree buildInputMerkleTree(
+      Spawn spawn, SpawnExecutionContext context, ToolSignature toolSignature)
       throws IOException, ForbiddenActionInputException {
     // Add output directories to inputs so that they are created as empty directories by the
     // executor. The spec only requires the executor to create the parent directory of an output
     // directory, which differs from the behavior of both local and sandboxed execution.
     SortedMap<PathFragment, ActionInput> outputDirMap = buildOutputDirMap(spawn);
-    if (remoteOptions.remoteMerkleTreeCache) {
+    if (remoteOptions.remoteMerkleTreeCache && toolSignature == null) {
       MetadataProvider metadataProvider = context.getMetadataProvider();
       ConcurrentLinkedQueue<MerkleTree> subMerkleTrees = new ConcurrentLinkedQueue<>();
       remotePathResolver.walkInputs(
@@ -394,7 +400,12 @@ private MerkleTree buildInputMerkleTree(Spawn spawn, SpawnExecutionContext conte
         newInputMap.putAll(outputDirMap);
         inputMap = newInputMap;
       }
-      return MerkleTree.build(inputMap, context.getMetadataProvider(), execRoot, digestUtil);
+      return MerkleTree.build(
+          inputMap,
+          toolSignature == null ? ImmutableSet.of() : toolSignature.toolInputs,
+          context.getMetadataProvider(),
+          execRoot,
+          digestUtil);
     }
   }
 
@@ -441,11 +452,24 @@ private static ByteString buildSalt(Spawn spawn) {
 
   /** Creates a new {@link RemoteAction} instance from spawn. */
   public RemoteAction buildRemoteAction(Spawn spawn, SpawnExecutionContext context)
-      throws IOException, UserExecException, ForbiddenActionInputException {
-    final MerkleTree merkleTree = buildInputMerkleTree(spawn, context);
+      throws IOException, ExecException, ForbiddenActionInputException, InterruptedException {
+    ToolSignature toolSignature =
+        remoteOptions.markToolInputs
+                && Spawns.supportsWorkers(spawn)
+                && !spawn.getToolFiles().isEmpty()
+            ? computePersistentWorkerSignature(spawn, context)
+            : null;
+    final MerkleTree merkleTree = buildInputMerkleTree(spawn, context, toolSignature);
 
     // Get the remote platform properties.
     Platform platform = PlatformUtils.getPlatformProto(spawn, remoteOptions);
+    if (toolSignature != null) {
+      platform =
+          PlatformUtils.getPlatformProto(
+              spawn, remoteOptions, ImmutableMap.of("persistentWorkerKey", toolSignature.key));
+    } else {
+      platform = PlatformUtils.getPlatformProto(spawn, remoteOptions);
+    }
 
     Command command =
         buildCommand(
@@ -485,6 +509,21 @@ public RemoteAction buildRemoteAction(Spawn spawn, SpawnExecutionContext context
         actionKey);
   }
 
+  @Nullable
+  private ToolSignature computePersistentWorkerSignature(Spawn spawn, SpawnExecutionContext context)
+      throws IOException, ExecException, InterruptedException {
+    WorkerParser workerParser =
+        new WorkerParser(
+            execRoot, Options.getDefaults(WorkerOptions.class), LocalEnvProvider.NOOP, null);
+    WorkerKey workerKey = workerParser.compute(spawn, context).getWorkerKey();
+    Fingerprint fingerprint = new Fingerprint();
+    fingerprint.addBytes(workerKey.getWorkerFilesCombinedHash().asBytes());
+    fingerprint.addIterableStrings(workerKey.getArgs());
+    fingerprint.addStringMap(workerKey.getEnv());
+    return new ToolSignature(
+        fingerprint.hexDigestAndReset(), workerKey.getWorkerFilesWithHashes().keySet());
+  }
+
   /** A value class representing the result of remotely executed {@link RemoteAction}. */
   public static class RemoteActionResult {
     private final ActionResult actionResult;
@@ -1468,4 +1507,18 @@ void report(Event evt) {
       reporter.handle(evt);
     }
   }
+
+  /**
+   * A simple value class combining a hash of the tool inputs (and their digests) as well as a set
+   * of the relative paths of all tool inputs.
+   */
+  private static final class ToolSignature {
+    private final String key;
+    private final Set<PathFragment> toolInputs;
+
+    private ToolSignature(String key, Set<PathFragment> toolInputs) {
+      this.key = key;
+      this.toolInputs = toolInputs;
+    }
+  }
 }

src/main/java/com/google/devtools/build/lib/remote/merkletree/DirectoryTree.java

@@ -73,6 +73,7 @@ static class FileNode extends Node {
     private final ByteString data;
     private final Digest digest;
     private final boolean isExecutable;
+    private final boolean toolInput;
 
     /**
      * Create a FileNode with its executable bit set.
@@ -83,27 +84,41 @@ static class FileNode extends Node {
      * https://github.com/bazelbuild/bazel/issues/13262 for more details.
      */
     static FileNode createExecutable(String pathSegment, Path path, Digest digest) {
-      return new FileNode(pathSegment, path, digest, /* isExecutable= */ true);
+      return new FileNode(pathSegment, path, digest, /* isExecutable= */ true, false);
     }
 
-    static FileNode createExecutable(String pathSegment, ByteString data, Digest digest) {
-      return new FileNode(pathSegment, data, digest, /* isExecutable= */ true);
+    static FileNode createExecutable(
+        String pathSegment, Path path, Digest digest, boolean toolInput) {
+      return new FileNode(pathSegment, path, digest, /* isExecutable= */ true, toolInput);
     }
 
-    private FileNode(String pathSegment, Path path, Digest digest, boolean isExecutable) {
+    static FileNode createExecutable(
+        String pathSegment, ByteString data, Digest digest, boolean toolInput) {
+      return new FileNode(pathSegment, data, digest, /* isExecutable= */ true, toolInput);
+    }
+
+    private FileNode(
+        String pathSegment, Path path, Digest digest, boolean isExecutable, boolean toolInput) {
       super(pathSegment);
       this.path = Preconditions.checkNotNull(path, "path");
       this.data = null;
       this.digest = Preconditions.checkNotNull(digest, "digest");
       this.isExecutable = isExecutable;
+      this.toolInput = toolInput;
     }
 
-    private FileNode(String pathSegment, ByteString data, Digest digest, boolean isExecutable) {
+    private FileNode(
+        String pathSegment,
+        ByteString data,
+        Digest digest,
+        boolean isExecutable,
+        boolean toolInput) {
       super(pathSegment);
       this.path = null;
       this.data = Preconditions.checkNotNull(data, "data");
       this.digest = Preconditions.checkNotNull(digest, "digest");
       this.isExecutable = isExecutable;
+      this.toolInput = toolInput;
     }
 
     Digest getDigest() {
@@ -122,9 +137,13 @@ public boolean isExecutable() {
       return isExecutable;
     }
 
+    boolean isToolInput() {
+      return toolInput;
+    }
+
     @Override
     public int hashCode() {
-      return Objects.hash(super.hashCode(), path, data, digest, isExecutable);
+      return Objects.hash(super.hashCode(), path, data, digest, toolInput, isExecutable);
     }
 
     @Override
@@ -135,6 +154,7 @@ public boolean equals(Object o) {
             && Objects.equals(path, other.path)
             && Objects.equals(data, other.data)
             && Objects.equals(digest, other.digest)
+            && toolInput == other.toolInput
             && isExecutable == other.isExecutable;
       }
       return false;