Cache merkle trees

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,28 @@ 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) &&
+      Objects.equals(buildRunfileLinks, that.buildRunfileLinks) &&
+      Objects.equals(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,13 @@ 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);
+        ActionInputHelper.expandArtifacts(inputFiles, artifactExpander);
     for (ActionInput input : inputs) {
       addMapping(inputMap, input.getExecPath(), input, baseDirectory);
     }
@@ -243,7 +245,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 +256,120 @@ public SortedMap<PathFragment, ActionInput> getInputMapping(
     return inputMap;
   }
 
+  public interface InputWalker {
+    SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
+        throws IOException, ForbiddenActionInputException;
+    void visitNonLeaves(InputVisitor visitor)
+        throws IOException, ForbiddenActionInputException;
+  }
+
+  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 retreival 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() {
+          public SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
+              throws IOException, ForbiddenActionInputException {
+            TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
+            addRunfilesToInputs(
+                inputMap,
+                runfilesSupplier, actionInputFileCache, artifactExpander, baseDirectory);
+            return inputMap;
+          }
+
+          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() {
+          public SortedMap<PathFragment, ActionInput> getLeavesInputMapping()
+              throws ForbiddenRelativeSymlinkException {
+            TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
+            addFilesetManifests(filesetMappings, inputMap, baseDirectory);
+            return inputMap;
+          }
+
+          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() {
+          public SortedMap<PathFragment, ActionInput> getLeavesInputMapping() {
+            TreeMap<PathFragment, ActionInput> inputMap = new TreeMap<>();
+            addInputs(
+                inputMap,
+                NestedSetBuilder.wrap(someInputFiles.getOrder(), someInputFiles.getLeaves()),
+                artifactExpander,
+                baseDirectory);
+            return inputMap;
+          }
+
+          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",