Remote: Cache merkle trees

When --experimental_remote_merkle_tree_cache is set, Merkle tree calculations are 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.

The memory foot print of the cache is controlled by --experimental_remote_merkle_tree_cache_size.

The caching is discarded after each build to free up memory, the cache setup time is negligible.

Fixes bazelbuild#10875.

Closes bazelbuild#13879.

PiperOrigin-RevId: 405793372

Author: moroten

src/main/java/com/google/devtools/build/lib/actions/CompositeRunfilesSupplier.java

@@ -65,6 +65,20 @@ public static RunfilesSupplier of(RunfilesSupplier supplier1, RunfilesSupplier s
     this.suppliers = suppliers;
   }
 
+  @Override
+  public boolean equals(Object other) {
+    if (!(other instanceof CompositeRunfilesSupplier)) {
+      return false;
+    }
+    CompositeRunfilesSupplier that = (CompositeRunfilesSupplier) other;
+    return suppliers.equals(that.suppliers);
+  }
+
+  @Override
+  public int hashCode() {
+    return suppliers.hashCode();
+  }
+
   @Override
   public NestedSet<Artifact> getArtifacts() {
     NestedSetBuilder<Artifact> result = NestedSetBuilder.stableOrder();

src/main/java/com/google/devtools/build/lib/actions/EmptyRunfilesSupplier.java

@@ -24,12 +24,22 @@
 import java.util.Map;
 
 /** Empty implementation of RunfilesSupplier */
-public class EmptyRunfilesSupplier implements RunfilesSupplier {
+public final class EmptyRunfilesSupplier implements RunfilesSupplier {
 
   @AutoCodec public static final EmptyRunfilesSupplier INSTANCE = new EmptyRunfilesSupplier();
 
   private EmptyRunfilesSupplier() {}
 
+  @Override
+  public boolean equals(Object other) {
+    return (other instanceof EmptyRunfilesSupplier);
+  }
+
+  @Override
+  public int hashCode() {
+    return 0;
+  }
+
   @Override
   public NestedSet<Artifact> getArtifacts() {
     return NestedSetBuilder.<Artifact>stableOrder().build();

src/main/java/com/google/devtools/build/lib/analysis/SingleRunfilesSupplier.java

@@ -27,6 +27,7 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.lang.ref.SoftReference;
 import java.util.Map;
+import java.util.Objects;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
@@ -47,6 +48,7 @@ public static SingleRunfilesSupplier create(RunfilesSupport runfilesSupport) {
     return new SingleRunfilesSupplier(
         runfilesSupport.getRunfilesDirectoryExecPath(),
         runfilesSupport.getRunfiles(),
+        /*runfilesCachingEnabled=*/ false,
         /*manifest=*/ null,
         runfilesSupport.isBuildRunfileLinks(),
         runfilesSupport.isRunfilesEnabled());
@@ -68,7 +70,7 @@ public static SingleRunfilesSupplier createCaching(
     return new SingleRunfilesSupplier(
         runfilesDir,
         runfiles,
-        new RunfilesCacher(runfiles),
+        /*runfilesCachingEnabled=*/ true,
         /*manifest=*/ null,
         buildRunfileLinks,
         runfileLinksEnabled);
@@ -95,7 +97,25 @@ public SingleRunfilesSupplier(
     this(
         runfilesDir,
         runfiles,
-        () -> runfiles.getRunfilesInputs(/*eventHandler=*/ null, /*location=*/ null),
+        /*runfilesCachingEnabled=*/ false,
+        manifest,
+        buildRunfileLinks,
+        runfileLinksEnabled);
+  }
+
+  private SingleRunfilesSupplier(
+      PathFragment runfilesDir,
+      Runfiles runfiles,
+      boolean runfilesCachingEnabled,
+      @Nullable Artifact manifest,
+      boolean buildRunfileLinks,
+      boolean runfileLinksEnabled) {
+    this(
+        runfilesDir,
+        runfiles,
+        runfilesCachingEnabled
+            ? new RunfilesCacher(runfiles)
+            : () -> runfiles.getRunfilesInputs(/*eventHandler=*/ null, /*location=*/ null),
         manifest,
         buildRunfileLinks,
         runfileLinksEnabled);
@@ -117,6 +137,26 @@ private SingleRunfilesSupplier(
     this.runfileLinksEnabled = runfileLinksEnabled;
   }
 
+  @Override
+  public boolean equals(Object other) {
+    if (!(other instanceof SingleRunfilesSupplier)) {
+      return false;
+    }
+
+    SingleRunfilesSupplier that = (SingleRunfilesSupplier) other;
+    // Not dependent on runfilesInputs which is only used for enabling caching.
+    return (Objects.equals(runfilesDir, that.runfilesDir)
+        && Objects.equals(runfiles, that.runfiles)
+        && Objects.equals(manifest, that.manifest)
+        && (buildRunfileLinks == that.buildRunfileLinks)
+        && (runfileLinksEnabled == that.runfileLinksEnabled));
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(runfilesDir, runfiles, manifest, buildRunfileLinks, runfileLinksEnabled);
+  }
+
   @Override
   public NestedSet<Artifact> getArtifacts() {
     return runfiles.getAllArtifacts();

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/SpawnInputExpander.java

@@ -33,11 +33,13 @@
 import com.google.devtools.build.lib.actions.RunfilesSupplier;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
+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.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -95,7 +97,7 @@ public SpawnInputExpander(
     this.relSymlinkBehavior = relSymlinkBehavior;
   }
 
-  private void addMapping(
+  private static void addMapping(
       Map<PathFragment, ActionInput> inputMappings,
       PathFragment targetLocation,
       ActionInput input,
@@ -215,13 +217,12 @@ void addFilesetManifest(
       }
   }
 
-  private void addInputs(
+  private static void addInputs(
       Map<PathFragment, ActionInput> inputMap,
-      Spawn spawn,
+      NestedSet<? extends ActionInput> inputFiles,
       ArtifactExpander artifactExpander,
       PathFragment baseDirectory) {
-    List<ActionInput> inputs =
-        ActionInputHelper.expandArtifacts(spawn.getInputFiles(), artifactExpander);
+    List<ActionInput> inputs = ActionInputHelper.expandArtifacts(inputFiles, artifactExpander);
     for (ActionInput input : inputs) {
       addMapping(inputMap, input.getExecPath(), input, baseDirectory);
     }
@@ -243,7 +244,7 @@ public SortedMap<PathFragment, ActionInput> getInputMapping(
       MetadataProvider actionInputFileCache)
       throws IOException, ForbiddenActionInputException {
     TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
-    addInputs(inputMap, spawn, artifactExpander, baseDirectory);
+    addInputs(inputMap, spawn.getInputFiles(), artifactExpander, baseDirectory);
     addRunfilesToInputs(
         inputMap,
         spawn.getRunfilesSupplier(),
@@ -254,6 +255,126 @@ public SortedMap<PathFragment, ActionInput> getInputMapping(
     return inputMap;
   }
 
+  /** The interface for accessing part of the input hierarchy. */
+  public interface InputWalker {
+    SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
+        throws IOException, ForbiddenActionInputException;
+
+    void visitNonLeaves(InputVisitor visitor) throws IOException, ForbiddenActionInputException;
+  }
+
+  /** The interface for visiting part of the input hierarchy. */
+  public interface InputVisitor {
+    /**
+     * Visits a part of the input hierarchy.
+     *
+     * <p>{@code nodeKey} can be used as key when memoizing visited parts of the hierarchy.
+     */
+    void visit(Object nodeKey, InputWalker walker)
+        throws IOException, ForbiddenActionInputException;
+  }
+
+  /**
+   * Visits the input files hierarchy in a depth first manner.
+   *
+   * <p>Similar to {@link #getInputMapping} but allows for early exit, by not visiting children,
+   * when walking through the input hierarchy. By applying memoization, the retrieval process of the
+   * inputs can be speeded up.
+   *
+   * <p>{@code baseDirectory} is prepended to every path in the input key. This is useful if the
+   * mapping is used in a context where the directory relative to which the keys are interpreted is
+   * not the same as the execroot.
+   */
+  public void walkInputs(
+      Spawn spawn,
+      ArtifactExpander artifactExpander,
+      PathFragment baseDirectory,
+      MetadataProvider actionInputFileCache,
+      InputVisitor visitor)
+      throws IOException, ForbiddenActionInputException {
+    walkNestedSetInputs(baseDirectory, spawn.getInputFiles(), artifactExpander, visitor);
+
+    RunfilesSupplier runfilesSupplier = spawn.getRunfilesSupplier();
+    visitor.visit(
+        // The list of variables affecting the functional expressions below.
+        Arrays.asList(
+            // Assuming that artifactExpander and actionInputFileCache, different for each spawn,
+            // always expand the same way.
+            this, // For accessing addRunfilesToInputs.
+            runfilesSupplier,
+            baseDirectory),
+        new InputWalker() {
+          @Override
+          public SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
+              throws IOException, ForbiddenActionInputException {
+            TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
+            addRunfilesToInputs(
+                inputMap, runfilesSupplier, actionInputFileCache, artifactExpander, baseDirectory);
+            return inputMap;
+          }
+
+          @Override
+          public void visitNonLeaves(InputVisitor childVisitor) {}
+        });
+
+    Map<Artifact, ImmutableList<FilesetOutputSymlink>> filesetMappings = spawn.getFilesetMappings();
+    // filesetMappings is assumed to be very small, so no need to implement visitNonLeaves() for
+    // improved runtime.
+    visitor.visit(
+        // The list of variables affecting the functional expressions below.
+        Arrays.asList(
+            this, // For accessing addFilesetManifests.
+            filesetMappings,
+            baseDirectory),
+        new InputWalker() {
+          @Override
+          public SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
+              throws ForbiddenRelativeSymlinkException {
+            TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
+            addFilesetManifests(filesetMappings, inputMap, baseDirectory);
+            return inputMap;
+          }
+
+          @Override
+          public void visitNonLeaves(InputVisitor childVisitor) {}
+        });
+  }
+
+  /** Walks through one level of a {@link NestedSet} of {@link ActionInput}s. */
+  private void walkNestedSetInputs(
+      PathFragment baseDirectory,
+      NestedSet<? extends ActionInput> someInputFiles,
+      ArtifactExpander artifactExpander,
+      InputVisitor visitor)
+      throws IOException, ForbiddenActionInputException {
+    visitor.visit(
+        // addInputs is static so no need to add 'this' as dependent key.
+        Arrays.asList(
+            // Assuming that artifactExpander, different for each spawn, always expands the same
+            // way.
+            someInputFiles.toNode(), baseDirectory),
+        new InputWalker() {
+          @Override
+          public SortedMap<PathFragment, ActionInput> getLeavesInputMapping() {
+            TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
+            addInputs(
+                inputMap,
+                NestedSetBuilder.wrap(someInputFiles.getOrder(), someInputFiles.getLeaves()),
+                artifactExpander,
+                baseDirectory);
+            return inputMap;
+          }
+
+          @Override
+          public void visitNonLeaves(InputVisitor childVisitor)
+              throws IOException, ForbiddenActionInputException {
+            for (NestedSet<? extends ActionInput> subInputs : someInputFiles.getNonLeaves()) {
+              walkNestedSetInputs(baseDirectory, subInputs, artifactExpander, childVisitor);
+            }
+          }
+        });
+  }
+
   /**
    * Exception signaling that an input was not a regular file: most likely a directory. This
    * exception is currently never thrown in practice since we do not enforce "strict" mode.

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",
@@ -94,6 +95,7 @@ java_library(
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/protobuf:failure_details_java_proto",
         "//third_party:auth",
+        "//third_party:caffeine",
         "//third_party:flogger",
         "//third_party:guava",
         "//third_party:jsr305",