From 21d0239d2cdd3b0b0bd8084592fb71b0deef5b66 Mon Sep 17 00:00:00 2001
From: Fredrik Medley <fredrik.medley@gmail.com>
Date: Thu, 19 Aug 2021 19:49:37 +0200
Subject: [PATCH] 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 #10875.
---
 .../lib/remote/RemoteExecutionService.java    | 47 +++++++++++++++++--
 .../lib/remote/options/RemoteOptions.java     | 10 ++++
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java
index 18d2ae50950ecf..14e564376c2252 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java
+++ b/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);
diff --git a/src/main/java/com/google/devtools/build/lib/remote/options/RemoteOptions.java b/src/main/java/com/google/devtools/build/lib/remote/options/RemoteOptions.java
index cb2f1639807279..2e19b7843dadfb 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/options/RemoteOptions.java
+++ b/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 = "",