Cache MerkleTree creation in RemoteExecutionService.java

src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java

@@ -269,6 +269,11 @@ public ArtifactPathResolver getPathResolver() {
       return actionExecutionContext.getPathResolver();
     }
 
+    @Override
+    public SpawnInputExpander getSpawnInputExpander() {
+      return spawnInputExpander;
+    }
+
     @Override
     public void lockOutputFiles() throws InterruptedException {
       if (stopConcurrentSpawns != null) {

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

@@ -275,6 +275,7 @@ java_library(
         "SpawnSchedulingEvent.java",
     ],
     deps = [
+        ":spawn_input_expander",
         ":tree_deleter",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",

src/main/java/com/google/devtools/build/lib/exec/SpawnRunner.java

@@ -161,6 +161,12 @@ interface SpawnExecutionContext {
     // directories? Or maybe we need a separate method to return the set of directories?
     ArtifactExpander getArtifactExpander();
 
+    /** A spawn input expander. */
+    // TODO(moroten): This is only used for the remote cache and remote execution to optimize
+    // Merkle tree generation. Having both this and the getInputMapping method seems a bit
+    // duplicated.
+    SpawnInputExpander getSpawnInputExpander();
+
     /** The {@link ArtifactPathResolver} to use when directly writing output files. */
     default ArtifactPathResolver getPathResolver() {
       return ArtifactPathResolver.IDENTITY;

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

@@ -66,6 +66,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/exec:module_action_context_registry",
         "//src/main/java/com/google/devtools/build/lib/exec:remote_local_fallback_registry",
         "//src/main/java/com/google/devtools/build/lib/exec:spawn_cache",
+        "//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/packages",

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

@@ -56,6 +56,8 @@
 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;
@@ -73,6 +75,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.SpawnResult;
 import com.google.devtools.build.lib.actions.Spawns;
@@ -82,6 +85,7 @@
 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.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.ProfilerTask;
@@ -133,6 +137,7 @@
 import java.util.Objects;
 import java.util.SortedMap;
 import java.util.TreeSet;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
 import javax.annotation.Nullable;
@@ -154,6 +159,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;
 
   private final Scheduler scheduler;
 
@@ -184,6 +190,15 @@ public RemoteExecutionService(
     this.remoteOptions = remoteOptions;
     this.remoteCache = remoteCache;
     this.remoteExecutor = remoteExecutor;
+
+    CacheBuilder merkleTreeCacheBuilder = CacheBuilder.newBuilder()
+        .softValues();
+    // remoteMerkleTreesCacheSize = 0 means limitless.
+    if (remoteOptions.remoteMerkleTreesCacheSize != 0) {
+      merkleTreeCacheBuilder.maximumSize(remoteOptions.remoteMerkleTreesCacheSize);
+    }
+    this.merkleTreeCache = merkleTreeCacheBuilder.build();
+
     ImmutableSet.Builder<PathFragment> filesToDownloadBuilder = ImmutableSet.builder();
     for (ActionInput actionInput : filesToDownload) {
       filesToDownloadBuilder.add(actionInput.getExecPath());
@@ -363,12 +378,54 @@ public boolean mayBeExecutedRemotely(Spawn spawn) {
         && Spawns.mayBeExecutedRemotely(spawn);
   }
 
+  private MerkleTree buildInputMerkleTree(Spawn spawn, SpawnExecutionContext context)
+      throws IOException, ForbiddenActionInputException {
+    if (remoteOptions.remoteMerkleTreesCacheEnabled) {
+      MetadataProvider metadataProvider = context.getMetadataProvider();
+      ConcurrentLinkedQueue<MerkleTree> subMerkleTrees = new ConcurrentLinkedQueue();
+      remotePathResolver.walkInputs(
+          spawn,
+          context,
+          (Object nodeKey, 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, 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(
+      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, 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(spawn, context);
 
     // Get the remote platform properties.
     Platform platform = PlatformUtils.getPlatformProto(spawn, remoteOptions);

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

@@ -19,6 +19,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
         "//src/main/java/com/google/devtools/build/lib/concurrent",
+        "//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/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",

src/main/java/com/google/devtools/build/lib/remote/common/RemotePathResolver.java

@@ -18,6 +18,8 @@
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputHelper;
 import com.google.devtools.build.lib.actions.ForbiddenActionInputException;
+import com.google.devtools.build.lib.actions.Spawn;
+import com.google.devtools.build.lib.exec.SpawnInputExpander;
 import com.google.devtools.build.lib.exec.SpawnRunner.SpawnExecutionContext;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -43,6 +45,12 @@ public interface RemotePathResolver {
   SortedMap<PathFragment, ActionInput> getInputMapping(SpawnExecutionContext context)
       throws IOException, ForbiddenActionInputException;
 
+  void walkInputs(
+      Spawn spawn,
+      SpawnExecutionContext context,
+      SpawnInputExpander.InputVisitor visitor)
+      throws IOException, ForbiddenActionInputException;
+
   /** Resolves the output path relative to input root for the given {@link Path}. */
   String localPathToOutputPath(Path path);
 
@@ -99,6 +107,18 @@ public SortedMap<PathFragment, ActionInput> getInputMapping(SpawnExecutionContex
       return context.getInputMapping(PathFragment.EMPTY_FRAGMENT);
     }
 
+    @Override
+    public void walkInputs(
+        Spawn spawn, SpawnExecutionContext context, SpawnInputExpander.InputVisitor visitor)
+        throws IOException, ForbiddenActionInputException {
+      context.getSpawnInputExpander().walkInputs(
+          spawn,
+          context.getArtifactExpander(),
+          PathFragment.EMPTY_FRAGMENT,
+          context.getMetadataProvider(),
+          visitor);
+    }
+
     @Override
     public String localPathToOutputPath(Path path) {
       return path.relativeTo(execRoot).getPathString();
@@ -159,6 +179,18 @@ public SortedMap<PathFragment, ActionInput> getInputMapping(SpawnExecutionContex
       return context.getInputMapping(PathFragment.create(checkNotNull(getWorkingDirectory())));
     }
 
+    @Override
+    public void walkInputs(
+        Spawn spawn, SpawnExecutionContext context, SpawnInputExpander.InputVisitor visitor)
+        throws IOException, ForbiddenActionInputException {
+      context.getSpawnInputExpander().walkInputs(
+          spawn,
+          context.getArtifactExpander(),
+          PathFragment.create(checkNotNull(getWorkingDirectory())),
+          context.getMetadataProvider(),
+          visitor);
+    }
+
     private Path getBase() {
       if (incompatibleRemoteOutputPathsRelativeToInputRoot) {
         return execRoot.getParentDirectory();

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

@@ -505,6 +505,29 @@ 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_enabled",
+      defaultValue = "false",
+      documentationCategory = OptionDocumentationCategory.REMOTE,
+      effectTags = {OptionEffectTag.UNKNOWN},
+      help =
+          "If set to true, Merkle tree calculations will be memoized to improve the remote cache "
+              + "hit checking speed. The memory foot print of the cache is controlled by "
+              + "--experimental_remote_merkle_tree_cache_size.")
+  public boolean remoteMerkleTreesCacheEnabled;
+
+  @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. "
+              + "Even though the cache is automatically pruned according to Java's handling of "
+              + "soft references, out-of-memory errors can occur if set too high. If set to 0 "
+              + "(default), the cache size is unlimited.")
+  public long remoteMerkleTreesCacheSize;
+
   @Option(
       name = "remote_download_symlink_template",
       defaultValue = "",

src/test/java/com/google/devtools/build/lib/exec/local/BUILD

@@ -28,6 +28,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/actions:localhost_capacity",
         "//src/main/java/com/google/devtools/build/lib/exec:bin_tools",
         "//src/main/java/com/google/devtools/build/lib/exec:runfiles_tree_updater",
+        "//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/local",
         "//src/main/java/com/google/devtools/build/lib/exec/local:options",

src/test/java/com/google/devtools/build/lib/exec/local/LocalSpawnRunnerTest.java

@@ -49,6 +49,7 @@
 import com.google.devtools.build.lib.exec.BinTools;
 import com.google.devtools.build.lib.exec.RunfilesTreeUpdater;
 import com.google.devtools.build.lib.exec.SpawnExecutingEvent;
+import com.google.devtools.build.lib.exec.SpawnInputExpander;
 import com.google.devtools.build.lib.exec.SpawnRunner.ProgressStatus;
 import com.google.devtools.build.lib.exec.SpawnRunner.SpawnExecutionContext;
 import com.google.devtools.build.lib.exec.SpawnSchedulingEvent;
@@ -260,6 +261,11 @@ public ArtifactExpander getArtifactExpander() {
       throw new UnsupportedOperationException();
     }
 
+    @Override
+    public SpawnInputExpander getSpawnInputExpander() {
+      throw new UnsupportedOperationException();
+    }
+
     @Override
     public Duration getTimeout() {
       return Duration.ofMillis(timeoutMillis);

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

@@ -1386,6 +1386,93 @@ public void uploadInputsIfNotPresent_deduplicateFindMissingBlobCalls() throws Ex
     }
   }
 
+  @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.remoteMerkleTreesCacheEnabled = true;
+    remoteOptions.remoteMerkleTreesCacheSize = 0;
+    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);
   }

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

@@ -146,6 +146,11 @@ public ArtifactExpander getArtifactExpander() {
           throw new UnsupportedOperationException();
         }
 
+        @Override
+        public SpawnInputExpander getSpawnInputExpander() {
+          return new SpawnInputExpander(execRoot, /*strict*/ false);
+        }
+
         @Override
         public Duration getTimeout() {
           return Duration.ZERO;
@@ -159,7 +164,7 @@ public FileOutErr getFileOutErr() {
         @Override
         public SortedMap<PathFragment, ActionInput> getInputMapping(PathFragment baseDirectory)
             throws IOException, ForbiddenActionInputException {
-          return new SpawnInputExpander(execRoot, /*strict*/ false)
+          return getSpawnInputExpander()
               .getInputMapping(simpleSpawn, SIMPLE_ARTIFACT_EXPANDER, baseDirectory, fakeFileCache);
         }