Cache MerkleTree creation in RemoteExecutionService.java

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

@@ -52,8 +52,11 @@
 import build.bazel.remote.execution.v2.RequestMetadata;
 import build.bazel.remote.execution.v2.SymlinkNode;
 import build.bazel.remote.execution.v2.Tree;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -69,6 +72,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 +120,7 @@
 import java.util.Objects;
 import java.util.SortedMap;
 import java.util.TreeSet;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import javax.annotation.Nullable;
 
 /**
@@ -133,6 +138,7 @@ public class RemoteExecutionService {
   @Nullable private final RemoteExecutionClient remoteExecutor;
   private final ImmutableSet<PathFragment> filesToDownload;
   @Nullable private final Path captureCorruptedOutputsDir;
+  private final Cache<Object, MerkleTree> merkleTreeCache;
 
   public RemoteExecutionService(
       Path execRoot,
@@ -154,6 +160,13 @@ public RemoteExecutionService(
     this.remoteCache = remoteCache;
     this.remoteExecutor = remoteExecutor;
 
+    CacheBuilder merkleTreeCacheBuilder = CacheBuilder.newBuilder()
+        .softValues();
+    if (remoteOptions.remoteMerkleTreesCacheSize != -1) {
+      merkleTreeCacheBuilder.maximumSize(remoteOptions.remoteMerkleTreesCacheSize);
+    }
+    merkleTreeCache = merkleTreeCacheBuilder.build();
+
     ImmutableSet.Builder<PathFragment> filesToDownloadBuilder = ImmutableSet.builder();
     for (ActionInput actionInput : filesToDownload) {
       filesToDownloadBuilder.add(actionInput.getExecPath());
@@ -331,12 +344,53 @@ public boolean mayBeExecutedRemotely(Spawn spawn) {
         && Spawns.mayBeExecutedRemotely(spawn);
   }
 
+  private MerkleTree buildInputMerkleTree(SpawnExecutionContext context)
+      throws IOException, ForbiddenActionInputException {
+    if (remoteOptions.remoteMerkleTreesCacheSize != 0) {
+      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.getIfPresent(nodeKey);
+    if (result == null) {
+      result = uncachedBuildMerkleTreeVisitor(walker, metadataProvider);
+      merkleTreeCache.put(nodeKey, result);
+    }
+    return result;
+  }
+
+  @VisibleForTesting
+  public MerkleTree uncachedBuildMerkleTreeVisitor(
+      SpawnExecutionContext.InputWalker walker, MetadataProvider metadataProvider)
+      throws IOException, ForbiddenActionInputException {
+    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));
+        });
+    return MerkleTree.merge(subMerkleTrees, digestUtil);
+  }
+
   /** 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,19 @@ public RemoteOutputsStrategyConverter() {
               + " discard the remotely cached values if they don't match the expected value.")
   public boolean remoteVerifyDownloads;
 
+  @Option(
+      name = "experimental_remote_merkle_tree_cache_size",
+      defaultValue = "0",
+      documentationCategory = OptionDocumentationCategory.REMOTE,
+      effectTags = {OptionEffectTag.UNKNOWN},
+      help =
+          "The number of Merkle trees to memoize to improve the remote cache hit checking speed. "
+              + "Enabling this option will increase the memory foot print, although the cache is "
+              + "automatically pruned according to Java's handling of soft references. If set to "
+              + "0 (default), the memoization is disabled. If set to -1, the cache size is "
+              + "unlimited.")
+  public long remoteMerkleTreesCacheSize;
+
   @Option(
       name = "remote_download_symlink_template",
       defaultValue = "",

src/test/java/com/google/devtools/build/lib/remote/RemoteExecutionServiceTest.java

@@ -46,6 +46,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.actions.ActionInput;
+import com.google.devtools.build.lib.actions.ActionInputHelper;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
@@ -58,6 +59,7 @@
 import com.google.devtools.build.lib.actions.cache.MetadataInjector;
 import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
 import com.google.devtools.build.lib.clock.JavaClock;
+import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.exec.util.FakeOwner;
@@ -1044,6 +1046,92 @@ public void downloadOutputs_missingInMemoryOutputWithMinimal_returnsNull() throw
     verify(injector, never()).injectFile(eq(a1), remoteFileMatchingDigest(d1));
   }
 
+  @Test
+  public void buildMerkleTree_withMemoization_works() throws Exception {
+    // Test that Merkle tree building can be memoized.
+
+    // TODO: Would like to check that NestedSet.getNonLeaves() is only called once per node, but
+    //       cannot Mockito.spy on NestedSet as it is final.
+
+    // arrange
+    /*
+     * First:
+     *   /bar/file
+     *   /foo1/file
+     * Second:
+     *   /bar/file
+     *   /foo2/file
+     */
+
+    // arrange
+    // Single node NestedSets are folded, so always add a dummy file everywhere.
+    ActionInput dummyFile = ActionInputHelper.fromPath("dummy");
+    fakeFileCache.createScratchInput(dummyFile, "dummy");
+
+    ActionInput barFile = ActionInputHelper.fromPath("bar/file");
+    NestedSet<ActionInput> nodeBar = NestedSetBuilder.create(
+        Order.STABLE_ORDER, dummyFile, barFile);
+    fakeFileCache.createScratchInput(barFile, "bar");
+
+    ActionInput foo1File = ActionInputHelper.fromPath("foo1/file");
+    NestedSet<ActionInput> nodeFoo1 = NestedSetBuilder.create(
+        Order.STABLE_ORDER, dummyFile, foo1File);
+    fakeFileCache.createScratchInput(foo1File, "foo1");
+
+    ActionInput foo2File = ActionInputHelper.fromPath("foo2/file");
+    NestedSet<ActionInput> nodeFoo2 = NestedSetBuilder.create(
+        Order.STABLE_ORDER, dummyFile, foo2File);
+    fakeFileCache.createScratchInput(foo2File, "foo2");
+
+    NestedSet<ActionInput> nodeRoot1 = new NestedSetBuilder(Order.STABLE_ORDER)
+        .add(dummyFile)
+        .addTransitive(nodeBar)
+        .addTransitive(nodeFoo1)
+        .build();
+    NestedSet<ActionInput> nodeRoot2 = new NestedSetBuilder(Order.STABLE_ORDER)
+        .add(dummyFile)
+        .addTransitive(nodeBar)
+        .addTransitive(nodeFoo2)
+        .build();
+
+    Spawn spawn1 = new SimpleSpawn(
+        new FakeOwner("foo", "bar", "//dummy:label"),
+        /*arguments=*/ ImmutableList.of(),
+        /*environment=*/ ImmutableMap.of(),
+        /*executionInfo=*/ ImmutableMap.of(),
+        /*inputs=*/ nodeRoot1,
+        /*outputs=*/ ImmutableSet.of(),
+        ResourceSet.ZERO);
+    Spawn spawn2 = new SimpleSpawn(
+        new FakeOwner("foo", "bar", "//dummy:label"),
+        /*arguments=*/ ImmutableList.of(),
+        /*environment=*/ ImmutableMap.of(),
+        /*executionInfo=*/ ImmutableMap.of(),
+        /*inputs=*/ nodeRoot2,
+        /*outputs=*/ ImmutableSet.of(),
+        ResourceSet.ZERO);
+
+    FakeSpawnExecutionContext context1 = newSpawnExecutionContext(spawn1);
+    FakeSpawnExecutionContext context2 = newSpawnExecutionContext(spawn2);
+    RemoteOptions remoteOptions = Options.getDefaults(RemoteOptions.class);
+    remoteOptions.remoteMerkleTreesCacheSize = -1;
+    RemoteExecutionService service = spy(newRemoteExecutionService(remoteOptions));
+
+    // act first time
+    service.buildRemoteAction(spawn1, context1);
+
+    // assert first time
+    // Called for: manifests, runfiles, nodeRoot1, nodeFoo1 and nodeBar.
+    verify(service, times(5)).uncachedBuildMerkleTreeVisitor(any(), any());
+
+    // act second time
+    service.buildRemoteAction(spawn2, context2);
+
+    // assert second time
+    // Called again for: manifests, runfiles, nodeRoot2 and nodeFoo2 but not nodeBar (cached).
+    verify(service, times(5 + 4)).uncachedBuildMerkleTreeVisitor(any(), any());
+  }
+
   private Spawn newSpawnFromResult(RemoteActionResult result) {
     return newSpawnFromResult(ImmutableMap.of(), result);
   }