Cache MerkleTree creation in RemoteExecutionService.java

MerkleTree calculations are now cached for each node in the input
NestedSets (depsets). This drastically improves the speed when checking
for remote cache hits. One example reduced the Merkle tree calculation
time from 78 ms to 3 ms for 3000 inputs.

This caching can be disabled using --remote_cache_merkle_trees=false
which will reduce the memory footprint. The caching is discarded after
each build to free up memory, the cache setup time is negligible.

Fixes bazelbuild#10875.

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

@@ -69,6 +69,7 @@
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue;
 import com.google.devtools.build.lib.actions.ForbiddenActionInputException;
+import com.google.devtools.build.lib.actions.MetadataProvider;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.Spawns;
 import com.google.devtools.build.lib.actions.UserExecException;
@@ -116,6 +117,8 @@
 import java.util.Objects;
 import java.util.SortedMap;
 import java.util.TreeSet;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import javax.annotation.Nullable;
 
 /**
@@ -133,6 +136,7 @@ public class RemoteExecutionService {
   @Nullable private final RemoteExecutionClient remoteExecutor;
   private final ImmutableSet<PathFragment> filesToDownload;
   @Nullable private final Path captureCorruptedOutputsDir;
+  private final ConcurrentHashMap<Object, MerkleTree> merkleTreeCache = new ConcurrentHashMap<>();
 
   public RemoteExecutionService(
       Path execRoot,
@@ -331,12 +335,49 @@ public boolean mayBeExecutedRemotely(Spawn spawn) {
         && Spawns.mayBeExecutedRemotely(spawn);
   }
 
+  private MerkleTree buildInputMerkleTree(SpawnExecutionContext context)
+      throws IOException, ForbiddenActionInputException {
+    if (remoteOptions.remoteCacheMerkleTrees) {
+      MetadataProvider metadataProvider = context.getMetadataProvider();
+      ConcurrentLinkedQueue<MerkleTree> subMerkleTrees = new ConcurrentLinkedQueue();
+      remotePathResolver.walkInputs(
+          context,
+          (Object nodeKey, SpawnExecutionContext.InputWalker walker) -> {
+            subMerkleTrees.add(buildMerkleTreeVisitor(nodeKey, walker, metadataProvider));
+          });
+      return MerkleTree.merge(subMerkleTrees, digestUtil);
+    } else {
+      SortedMap<PathFragment, ActionInput> inputMap = remotePathResolver.getInputMapping(context);
+      return MerkleTree.build(inputMap, context.getMetadataProvider(), execRoot, digestUtil);
+    }
+  }
+
+  private MerkleTree buildMerkleTreeVisitor(
+      Object nodeKey, SpawnExecutionContext.InputWalker walker, MetadataProvider metadataProvider)
+      throws IOException, ForbiddenActionInputException {
+    MerkleTree result = merkleTreeCache.get(nodeKey);
+    if (result == null) {
+      ConcurrentLinkedQueue<MerkleTree> subMerkleTrees = new ConcurrentLinkedQueue();
+      subMerkleTrees.add(MerkleTree.build(
+          walker.getLeavesInputMapping(), metadataProvider, execRoot, digestUtil));
+      walker.visitNonLeaves(
+          (Object subNodeKey, SpawnExecutionContext.InputWalker subWalker) -> {
+            subMerkleTrees.add(buildMerkleTreeVisitor(
+                subNodeKey, subWalker, metadataProvider));
+          });
+      result = MerkleTree.merge(subMerkleTrees, digestUtil);
+      MerkleTree existingResult = merkleTreeCache.putIfAbsent(nodeKey, result);
+      if (existingResult != null) {
+        result = existingResult;
+      }
+    }
+    return result;
+  }
+
   /** Creates a new {@link RemoteAction} instance from spawn. */
   public RemoteAction buildRemoteAction(Spawn spawn, SpawnExecutionContext context)
       throws IOException, UserExecException, ForbiddenActionInputException {
-    SortedMap<PathFragment, ActionInput> inputMap = remotePathResolver.getInputMapping(context);
-    final MerkleTree merkleTree =
-        MerkleTree.build(inputMap, context.getMetadataProvider(), execRoot, digestUtil);
+    final MerkleTree merkleTree = buildInputMerkleTree(context);
 
     // Get the remote platform properties.
     Platform platform = PlatformUtils.getPlatformProto(spawn, remoteOptions);

src/main/java/com/google/devtools/build/lib/remote/options/RemoteOptions.java

@@ -504,6 +504,16 @@ public RemoteOutputsStrategyConverter() {
               + " discard the remotely cached values if they don't match the expected value.")
   public boolean remoteVerifyDownloads;
 
+  @Option(
+      name = "remote_cache_merkle_trees",
+      defaultValue = "true",
+      documentationCategory = OptionDocumentationCategory.REMOTE,
+      effectTags = {OptionEffectTag.UNKNOWN},
+      help =
+          "Let the Merkle tree calculation be cached to improve speed for remote cache hit "
+              + "checking. Disabling this option will decrease the memory foot print.")
+  public boolean remoteCacheMerkleTrees;
+
   @Option(
       name = "remote_download_symlink_template",
       defaultValue = "",