Add experimental_remote_cache_key_ignore_stamping to skip volatile …

…stamping files on compute shared cache key
bazelbuild · Sep 8, 2022 · 733b624 · 733b624
1 parent a380a15
commit 733b624
Show file tree

Hide file tree

Showing 4 changed files with 113 additions and 52 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
@@ -165,7 +165,8 @@ public class RemoteExecutionService {
   private final ImmutableSet<PathFragment> filesToDownload;
   private final TempPathGenerator tempPathGenerator;
   @Nullable private final Path captureCorruptedOutputsDir;
-  private final Cache<Object, MerkleTree> merkleTreeCache;
+  private final MerkleTreeCache merkleTreeCache;
+  private final CacheKeyTreeCache cacheKeyTreeCache;
   private final Set<String> reportedErrors = new HashSet<>();
   private final Phaser backgroundTaskPhaser = new Phaser(1);
 
@@ -207,7 +208,8 @@ public RemoteExecutionService(
     if (remoteOptions.remoteMerkleTreeCacheSize != 0) {
       merkleTreeCacheBuilder.maximumSize(remoteOptions.remoteMerkleTreeCacheSize);
     }
-    this.merkleTreeCache = merkleTreeCacheBuilder.build();
+    this.merkleTreeCache = new MerkleTreeCache(merkleTreeCacheBuilder);
+    this.cacheKeyTreeCache = new CacheKeyTreeCache(merkleTreeCacheBuilder);
 
     ImmutableSet.Builder<PathFragment> filesToDownloadBuilder = ImmutableSet.builder();
     for (ActionInput actionInput : filesToDownload) {
@@ -381,62 +383,105 @@ private SortedMap<PathFragment, ActionInput> buildOutputDirMap(Spawn spawn) {
     return outputDirMap;
   }
 
-  private MerkleTree buildInputMerkleTree(Spawn spawn, SpawnExecutionContext context)
+  protected class MerkleTreeCache {
+    private final Cache<Object, MerkleTree> merkleTreeCache;
+
+    public MerkleTreeCache(Caffeine<Object, Object> cacheBuilder) {
+      this.merkleTreeCache = cacheBuilder.build();
+    }
+
+    public MerkleTree buildInputMerkleTree(Spawn spawn, SpawnExecutionContext context)
       throws IOException, ForbiddenActionInputException {
-    // Add output directories to inputs so that they are created as empty directories by the
-    // executor. The spec only requires the executor to create the parent directory of an output
-    // directory, which differs from the behavior of both local and sandboxed execution.
-    SortedMap<PathFragment, ActionInput> outputDirMap = buildOutputDirMap(spawn);
-    if (remoteOptions.remoteMerkleTreeCache) {
-      MetadataProvider metadataProvider = context.getMetadataProvider();
+      // Add output directories to inputs so that they are created as empty directories by the
+      // executor. The spec only requires the executor to create the parent directory of an output
+      // directory, which differs from the behavior of both local and sandboxed execution.
+      SortedMap<PathFragment, ActionInput> outputDirMap = buildOutputDirMap(spawn);
+      if (remoteOptions.remoteMerkleTreeCache) {
+        MetadataProvider metadataProvider = context.getMetadataProvider();
+        ConcurrentLinkedQueue<MerkleTree> subMerkleTrees = new ConcurrentLinkedQueue<>();
+        remotePathResolver.walkInputs(
+            spawn,
+            context,
+            (Object nodeKey, InputWalker walker) -> {
+              subMerkleTrees.add(buildMerkleTreeVisitor(nodeKey, walker, metadataProvider));
+            });
+        if (!outputDirMap.isEmpty()) {
+          subMerkleTrees.add(MerkleTree.build(outputDirMap, metadataProvider, execRoot, digestUtil));
+        }
+        return MerkleTree.merge(subMerkleTrees, digestUtil);
+      } else {
+        SortedMap<PathFragment, ActionInput> inputMap = remotePathResolver.getInputMapping(context);
+        if (!outputDirMap.isEmpty()) {
+          // The map returned by getInputMapping is mutable, but must not be mutated here as it is
+          // shared with all other strategies.
+          SortedMap<PathFragment, ActionInput> newInputMap = new TreeMap<>();
+          newInputMap.putAll(inputMap);
+          newInputMap.putAll(outputDirMap);
+          inputMap = newInputMap;
+        }
+        return MerkleTree.build(filterInputs(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;
+    }
+
+    public MerkleTree uncachedBuildMerkleTreeVisitor(
+        InputWalker walker, MetadataProvider metadataProvider)
+        throws IOException, ForbiddenActionInputException {
       ConcurrentLinkedQueue<MerkleTree> subMerkleTrees = new ConcurrentLinkedQueue<>();
-      remotePathResolver.walkInputs(
-          spawn,
-          context,
-          (Object nodeKey, InputWalker walker) -> {
-            subMerkleTrees.add(buildMerkleTreeVisitor(nodeKey, walker, metadataProvider));
+      subMerkleTrees.add(
+        MerkleTree.build(filterInputs(walker.getLeavesInputMapping()), metadataProvider, execRoot, digestUtil));
+      walker.visitNonLeaves(
+          (Object subNodeKey, InputWalker subWalker) -> {
+            subMerkleTrees.add(buildMerkleTreeVisitor(subNodeKey, subWalker, metadataProvider));
           });
-      if (!outputDirMap.isEmpty()) {
-        subMerkleTrees.add(MerkleTree.build(outputDirMap, metadataProvider, execRoot, digestUtil));
-      }
       return MerkleTree.merge(subMerkleTrees, digestUtil);
-    } else {
-      SortedMap<PathFragment, ActionInput> inputMap = remotePathResolver.getInputMapping(context);
-      if (!outputDirMap.isEmpty()) {
-        // The map returned by getInputMapping is mutable, but must not be mutated here as it is
-        // shared with all other strategies.
-        SortedMap<PathFragment, ActionInput> newInputMap = new TreeMap<>();
-        newInputMap.putAll(inputMap);
-        newInputMap.putAll(outputDirMap);
-        inputMap = newInputMap;
-      }
-      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);
+    protected SortedMap<PathFragment, ActionInput> filterInputs(SortedMap<PathFragment, ActionInput> inputs) {
+      return inputs;
     }
-    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);
+    InputWalker walker, MetadataProvider metadataProvider)
+    throws IOException, ForbiddenActionInputException {
+    return merkleTreeCache.uncachedBuildMerkleTreeVisitor(walker, metadataProvider);
+  }
+
+  protected class CacheKeyTreeCache extends MerkleTreeCache {
+    public CacheKeyTreeCache(Caffeine<Object, Object> cacheBuilder) {
+      super(cacheBuilder);
+    }
+
+    @Override
+    protected SortedMap<PathFragment, ActionInput> filterInputs(SortedMap<PathFragment, ActionInput> inputs) {
+      SortedMap<PathFragment, ActionInput> result = new TreeMap<>();
+      for (Entry<PathFragment, ActionInput> entry : inputs.entrySet()) {
+        ActionInput input = entry.getValue();
+        if (!isConstantMetadata(input)) {
+          result.put(entry.getKey(), input);
+        }
+      }
+      return result;
+    }
+
+    private boolean isConstantMetadata(ActionInput input) {
+      if (input instanceof Artifact) {
+        return ((Artifact) input).isConstantMetadata();
+      }
+      return false;
+    }
   }
 
   @Nullable
@@ -458,7 +503,7 @@ private static ByteString buildSalt(Spawn spawn) {
   /** Creates a new {@link RemoteAction} instance from spawn. */
   public RemoteAction buildRemoteAction(Spawn spawn, SpawnExecutionContext context)
       throws IOException, UserExecException, ForbiddenActionInputException {
-    final MerkleTree merkleTree = buildInputMerkleTree(spawn, context);
+    final MerkleTree merkleTree = merkleTreeCache.buildInputMerkleTree(spawn, context);
 
     // Get the remote platform properties.
     Platform platform = PlatformUtils.getPlatformProto(spawn, remoteOptions);
@@ -480,7 +525,11 @@ public RemoteAction buildRemoteAction(Spawn spawn, SpawnExecutionContext context
             Spawns.mayBeCachedRemotely(spawn),
             buildSalt(spawn));
 
-    ActionKey actionKey = digestUtil.computeActionKey(action);
+    Digest cacheKeyDigest = merkleTree.getRootDigest();
+    if (remoteOptions.remoteCacheKeyIgnoreStamping) {
+      cacheKeyDigest = cacheKeyTreeCache.buildInputMerkleTree(spawn, context).getRootDigest();
+    }
+    ActionKey actionKey = digestUtil.computeActionKey(action, cacheKeyDigest);
 
     RequestMetadata metadata =
         TracingMetadataUtils.buildMetadata(

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
@@ -83,6 +83,14 @@ public final class RemoteOptions extends CommonRemoteOptions {
               + "disable TLS.")
   public String remoteExecutor;
 
+  @Option(
+    name = "experimental_remote_cache_key_ignore_stamping",
+    defaultValue = "false",
+    documentationCategory = OptionDocumentationCategory.REMOTE,
+    effectTags = {OptionEffectTag.UNKNOWN},
+    help = "Don't use volatile stamping data in shared cache key.")
+  public boolean remoteCacheKeyIgnoreStamping;
+
   @Option(
       name = "experimental_remote_execution_keepalive",
       defaultValue = "false",

diff --git a/src/main/java/com/google/devtools/build/lib/remote/util/DigestUtil.java b/src/main/java/com/google/devtools/build/lib/remote/util/DigestUtil.java
@@ -85,8 +85,12 @@ public Digest computeAsUtf8(String str) {
     return compute(str.getBytes(UTF_8));
   }
 
-  public ActionKey computeActionKey(Action action) {
-    return new ActionKey(compute(action));
+  public ActionKey computeActionKey(Action action, Digest inputDigest) {
+    Action cacheKeyAction = Action.newBuilder()
+      .mergeFrom(action)
+      .setInputRootDigest(inputDigest)
+      .build();
+    return new ActionKey(compute(cacheKeyAction));
   }
 
   /**

diff --git a/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java b/src/test/java/com/google/devtools/build/lib/remote/GrpcCacheClientTest.java
@@ -496,7 +496,7 @@ private ActionResult upload(
   private ActionResult uploadDirectory(RemoteCache remoteCache, List<Path> outputs)
       throws Exception {
     Action action = Action.getDefaultInstance();
-    ActionKey actionKey = DIGEST_UTIL.computeActionKey(action);
+    ActionKey actionKey = DIGEST_UTIL.computeActionKey(action, action.getInputRootDigest());
     Command cmd = Command.getDefaultInstance();
     return upload(remoteCache, actionKey, action, cmd, outputs);
   }