Fix false sharing in action cache due to incorrect key computation for custom

Starlark command lines.

Actions constructing arguments in Starlark can define them as a lazy
evaluation, which will be run when we are about to execute the action, using
`Args.add_all`.

`Args.add_all` offers a `expand_directories` feature which causes all of the
directories in the args list to be replaced with the files in the directory at
the time of evaluation.

When this option is used in conjunction with `map_each`, the behavior is to
first expand the directory and than run the provided function on each of the
elements of the expanded array.

`ArtifactExpander`, used to expand the directories, is not available before
execution phase. Given `Action.getKey` can be called before that (and is for
instance for shared actions detection), we currently compute the fingerprint
without expanding directories. In fact, current implementation, calls the
provided function on the directory itself rather than expanded list of files in
it. As a result of that, if we have 2 functions which return the same value for
the directory, but differ for the files in it, the key will not denote them as
different. That leads to false sharing of actions through the action cache,
which is based on action fingerprints.

Action cache checks happen during execution phase, when we technically have the
ability to expand the directories. Expand the fingerprinting function to take
advantage of `ArtifactExpander` when it is provided to produce the key based on
the result of actual function execution. Rework the `ActionExecutionFunction`
to provide `ArtifactExpander` to logic computing keys for action caches.

PiperOrigin-RevId: 323090803

Author: alexjski

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

@@ -19,6 +19,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.actions.Artifact.ArtifactExpander;
 import com.google.devtools.build.lib.actions.cache.ActionCache;
 import com.google.devtools.build.lib.actions.cache.DigestUtils;
 import com.google.devtools.build.lib.actions.cache.MetadataHandler;
@@ -249,6 +250,7 @@ public Token getTokenIfNeedToExecute(
       Map<String, String> clientEnv,
       EventHandler handler,
       MetadataHandler metadataHandler,
+      ArtifactExpander artifactExpander,
       Map<String, String> remoteDefaultPlatformProperties) {
     // TODO(bazel-team): (2010) For RunfilesAction/SymlinkAction and similar actions that
     // produce only symlinks we should not check whether inputs are valid at all - all that matters
@@ -292,6 +294,7 @@ public Token getTokenIfNeedToExecute(
         entry,
         handler,
         metadataHandler,
+        artifactExpander,
         actionInputs,
         clientEnv,
         remoteDefaultPlatformProperties)) {
@@ -307,11 +310,12 @@ public Token getTokenIfNeedToExecute(
     return null;
   }
 
-  protected boolean mustExecute(
+  private boolean mustExecute(
       Action action,
       @Nullable ActionCache.Entry entry,
       EventHandler handler,
       MetadataHandler metadataHandler,
+      ArtifactExpander artifactExpander,
       NestedSet<Artifact> actionInputs,
       Map<String, String> clientEnv,
       Map<String, String> remoteDefaultPlatformProperties) {
@@ -336,9 +340,7 @@ protected boolean mustExecute(
       reportChanged(handler, action);
       actionCache.accountMiss(MissReason.DIFFERENT_FILES);
       return true;
-    } else if (!entry
-        .getActionKey()
-        .equals(action.getKey(actionKeyContext, /*artifactExpander=*/ null))) {
+    } else if (!entry.getActionKey().equals(action.getKey(actionKeyContext, artifactExpander))) {
       reportCommand(handler, action);
       actionCache.accountMiss(MissReason.DIFFERENT_ACTION_KEY);
       return true;
@@ -382,6 +384,7 @@ public void updateActionCache(
       Action action,
       Token token,
       MetadataHandler metadataHandler,
+      ArtifactExpander artifactExpander,
       Map<String, String> clientEnv,
       Map<String, String> remoteDefaultPlatformProperties)
       throws IOException {
@@ -397,7 +400,7 @@ public void updateActionCache(
         computeUsedEnv(action, clientEnv, remoteDefaultPlatformProperties);
     ActionCache.Entry entry =
         new ActionCache.Entry(
-            action.getKey(actionKeyContext, /*artifactExpander=*/ null),
+            action.getKey(actionKeyContext, artifactExpander),
             usedEnvironment,
             action.discoversInputs());
     for (Artifact output : action.getOutputs()) {

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

@@ -28,35 +28,44 @@ public abstract class ActionKeyCacher implements ActionAnalysisMetadata {
   @Override
   public final String getKey(
       ActionKeyContext actionKeyContext, @Nullable ArtifactExpander artifactExpander) {
+    // Only cache the key when it is given all necessary information to compute a correct key.
+    // Practically, most of the benefit of the cache comes from execution, which does provide the
+    // artifactExpander.
+    if (artifactExpander == null) {
+      return computeActionKey(actionKeyContext, null);
+    }
+
     if (cachedKey == null) {
       synchronized (this) {
         if (cachedKey == null) {
-          try {
-            Fingerprint fp = new Fingerprint();
-            // TODO(b/153904017): Make use of the provided artifactExpander and only cache if
-            // present.
-            computeKey(actionKeyContext, /*artifactExpander=*/ null, fp);
-
-            // Add a bool indicating whether the execution platform was set.
-            fp.addBoolean(getExecutionPlatform() != null);
-            if (getExecutionPlatform() != null) {
-              // Add the execution platform information.
-              getExecutionPlatform().addTo(fp);
-            }
-
-            fp.addStringMap(getExecProperties());
-
-            // Compute the actual key and store it.
-            cachedKey = fp.hexDigestAndReset();
-          } catch (CommandLineExpansionException e) {
-            cachedKey = KEY_ERROR;
-          }
+          cachedKey = computeActionKey(actionKeyContext, artifactExpander);
         }
       }
     }
     return cachedKey;
   }
 
+  private String computeActionKey(
+      ActionKeyContext actionKeyContext, @Nullable ArtifactExpander artifactExpander) {
+    try {
+      Fingerprint fp = new Fingerprint();
+      computeKey(actionKeyContext, artifactExpander, fp);
+
+      // Add a bool indicating whether the execution platform was set.
+      fp.addBoolean(getExecutionPlatform() != null);
+      if (getExecutionPlatform() != null) {
+        // Add the execution platform information.
+        getExecutionPlatform().addTo(fp);
+      }
+
+      fp.addStringMap(getExecProperties());
+      // Compute the actual key and store it.
+      return fp.hexDigestAndReset();
+    } catch (CommandLineExpansionException e) {
+      return KEY_ERROR;
+    }
+  }
+
   /**
    * See the javadoc for {@link Action} and {@link ActionAnalysisMetadata#getKey} for the contract
    * of this method.

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

@@ -22,6 +22,7 @@
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.VisibleForSerialization;
 import com.google.devtools.build.lib.util.Fingerprint;
+import javax.annotation.Nullable;
 
 /** A representation of a list of arguments. */
 public abstract class CommandLine {
@@ -52,7 +53,18 @@ public Iterable<String> arguments(ArtifactExpander artifactExpander)
     return arguments();
   }
 
-  public void addToFingerprint(ActionKeyContext actionKeyContext, Fingerprint fingerprint)
+  /**
+   * Adds the command line to the provided {@link Fingerprint}.
+   *
+   * <p>Some of the implementations may require the to expand provided directory in order to produce
+   * a unique key. Consequently, the result of calling this function can be different depending on
+   * whether the {@link ArtifactExpander} is provided. Moreover, without it, the produced key may
+   * not always be unique.
+   */
+  public void addToFingerprint(
+      ActionKeyContext actionKeyContext,
+      @Nullable ArtifactExpander artifactExpander,
+      Fingerprint fingerprint)
       throws CommandLineExpansionException {
     for (String s : arguments()) {
       fingerprint.addString(s);

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

@@ -180,19 +180,22 @@ ExpandedCommandLines expand(
     return new ExpandedCommandLines(arguments.build(), paramFiles);
   }
 
-  public void addToFingerprint(ActionKeyContext actionKeyContext, Fingerprint fingerprint)
+  public void addToFingerprint(
+      ActionKeyContext actionKeyContext,
+      @Nullable ArtifactExpander artifactExpander,
+      Fingerprint fingerprint)
       throws CommandLineExpansionException {
     // Optimize for simple case of single command line
     if (commandLines instanceof CommandLine) {
       CommandLine commandLine = (CommandLine) commandLines;
-      commandLine.addToFingerprint(actionKeyContext, fingerprint);
+      commandLine.addToFingerprint(actionKeyContext, artifactExpander, fingerprint);
       return;
     }
     List<CommandLineAndParamFileInfo> commandLines = getCommandLines();
     for (CommandLineAndParamFileInfo pair : commandLines) {
       CommandLine commandLine = pair.commandLine;
       ParamFileInfo paramFileInfo = pair.paramFileInfo;
-      commandLine.addToFingerprint(actionKeyContext, fingerprint);
+      commandLine.addToFingerprint(actionKeyContext, artifactExpander, fingerprint);
       if (paramFileInfo != null) {
         addParamFileInfoToFingerprint(paramFileInfo, fingerprint);
       }

src/main/java/com/google/devtools/build/lib/analysis/actions/CustomCommandLine.java

@@ -1288,7 +1288,10 @@ private Object substituteTreeFileArtifactArgvFragment(Object arg) {
 
   @Override
   @SuppressWarnings("unchecked")
-  public void addToFingerprint(ActionKeyContext actionKeyContext, Fingerprint fingerprint) {
+  public void addToFingerprint(
+      ActionKeyContext actionKeyContext,
+      @Nullable ArtifactExpander artifactExpander,
+      Fingerprint fingerprint) {
     int count = arguments.size();
     for (int i = 0; i < count; ) {
       Object arg = arguments.get(i++);

src/main/java/com/google/devtools/build/lib/analysis/actions/ParameterFileWriteAction.java

@@ -182,6 +182,6 @@ protected void computeKey(
     fp.addString(GUID);
     fp.addString(String.valueOf(makeExecutable));
     fp.addString(type.toString());
-    commandLine.addToFingerprint(actionKeyContext, fp);
+    commandLine.addToFingerprint(actionKeyContext, artifactExpander, fp);
   }
 }

src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java

@@ -441,7 +441,7 @@ protected void computeKey(
       Fingerprint fp)
       throws CommandLineExpansionException {
     fp.addString(GUID);
-    commandLines.addToFingerprint(actionKeyContext, fp);
+    commandLines.addToFingerprint(actionKeyContext, artifactExpander, fp);
     fp.addString(getMnemonic());
     // We don't need the toolManifests here, because they are a subset of the inputManifests by
     // definition and the output of an action shouldn't change whether something is considered a