Stop storing ActionTemplate in a SkyKey: it's too heavyweight. Use the same mechanism as for normal actions, have the ActionTemplateExpansionFunction look the template up when needed.

PiperOrigin-RevId: 185861672

7 files changed, 96 insertions, 83 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionLookupValue.java b/src/main/java/com/google/devtools/build/lib/actions/ActionLookupValue.java
index 000f359315..33555a20c7 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ActionLookupValue.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionLookupValue.java
@@ -122,18 +122,19 @@ public class ActionLookupValue implements SkyValue {
     return Preconditions.checkNotNull(actions.get(index), "null action: %s %s", index, this);
   }
 
+  public ActionTemplate<?> getActionTemplate(int index) {
+    ActionAnalysisMetadata result = getActionAnalysisMetadata(index);
+    Preconditions.checkState(
+        result instanceof ActionTemplate, "Not action template: %s %s %s", result, index, this);
+    return (ActionTemplate<?>) result;
+  }
+
   /**
-   * Returns the {@link ActionAnalysisMetadata} at index {@code index} if it is present and
-   * <i>not</i> an {@link Action}. Tree artifacts need their {@code ActionTemplate}s in order to
-   * generate the correct actions, but in general most actions are not needed after they are
-   * executed and may not even be available.
+   * Returns if the action at {@code index} is an {@link ActionTemplate} so that tree artifacts can
+   * take the proper action.
    */
-  public ActionAnalysisMetadata getIfPresentAndNotAction(int index) {
-    ActionAnalysisMetadata actionAnalysisMetadata = actions.get(index);
-    if (!(actionAnalysisMetadata instanceof Action)) {
-      return actionAnalysisMetadata;
-    }
-    return null;
+  public boolean isActionTemplate(int index) {
+    return actions.get(index) instanceof ActionTemplate;
   }
 
   /** To be used only when checking consistency of the action graph -- not by other values. */
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/ActionTemplate.java b/src/main/java/com/google/devtools/build/lib/actions/ActionTemplate.java
index 5fbc5ea5c8..2a68628e75 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/ActionTemplate.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionTemplate.java
@@ -11,29 +11,26 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-package com.google.devtools.build.lib.analysis.actions;
+package com.google.devtools.build.lib.actions;
 
-import com.google.devtools.build.lib.actions.Action;
-import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
-import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
-import com.google.devtools.build.lib.actions.ArtifactOwner;
 
 /**
  * A placeholder action that, at execution time, expands into a list of {@link Action}s to be
  * executed.
  *
- * <p>ActionTemplate is for users who want to dynamically register Actions operating on
- * individual {@link TreeFileArtifact} inside input and output TreeArtifacts at execution time.
+ * <p>ActionTemplate is for users who want to dynamically register Actions operating on individual
+ * {@link TreeFileArtifact} inside input and output TreeArtifacts at execution time.
  *
  * <p>It takes in one TreeArtifact and generates one TreeArtifact. The following happens at
  * execution time for ActionTemplate:
+ *
  * <ol>
  *   <li>Input TreeArtifact is resolved.
  *   <li>For each individual {@link TreeFileArtifact} inside input TreeArtifact, generate an output
  *       {@link TreeFileArtifact} inside output TreeArtifact.
- *   <li>For each pair of input and output {@link TreeFileArtifact}s, generate an associated
- *       {@link Action}.
+ *   <li>For each pair of input and output {@link TreeFileArtifact}s, generate an associated {@link
+ *       Action}.
  *   <li>All expanded {@link Action}s are executed and their output {@link TreeFileArtifact}s
  *       collected.
  *   <li>Output TreeArtifact is resolved.
@@ -41,13 +38,14 @@ import com.google.devtools.build.lib.actions.ArtifactOwner;
  *
  * <p>Implementations of ActionTemplate must follow the contract of this interface and also make
  * sure:
+ *
  * <ol>
  *   <li>ActionTemplate instances should be immutable and side-effect free.
  *   <li>ActionTemplate inputs and outputs are supersets of the inputs and outputs of expanded
- *       actions, excluding inputs discovered at execution time. This ensures the ActionTemplate
- *       can properly represent the expanded actions at analysis time, and the action graph
- *       at analysis time is correct. This is important because the action graph is walked in a lot
- *       of places for correctness checks and build analysis.
+ *       actions, excluding inputs discovered at execution time. This ensures the ActionTemplate can
+ *       properly represent the expanded actions at analysis time, and the action graph at analysis
+ *       time is correct. This is important because the action graph is walked in a lot of places
+ *       for correctness checks and build analysis.
  *   <li>The outputs of expanded actions must be under the output TreeArtifact and must not have
  *       artifact or artifact path prefix conflicts.
  * </ol>
@@ -72,10 +70,10 @@ public interface ActionTemplate<T extends Action> extends ActionAnalysisMetadata
    *
    * @param inputTreeFileArtifacts the list of {@link TreeFileArtifact}s inside input TreeArtifact
    *     resolved at execution time
-   * @param artifactOwner the {@link ArtifactOwner} of the generated output
-   *     {@link TreeFileArtifact}s
-   * @return a list of expanded {@link Action}s to execute, one for each input
-   *     {@link TreeFileArtifact}
+   * @param artifactOwner the {@link ArtifactOwner} of the generated output {@link
+   *     TreeFileArtifact}s
+   * @return a list of expanded {@link Action}s to execute, one for each input {@link
+   *     TreeFileArtifact}
    */
   Iterable<T> generateActionForInputArtifacts(
       Iterable<TreeFileArtifact> inputTreeFileArtifacts, ArtifactOwner artifactOwner)
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnActionTemplate.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnActionTemplate.java
index 1b944d44e8..c3d1228d96 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnActionTemplate.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnActionTemplate.java
@@ -20,6 +20,7 @@ import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionInputHelper;
 import com.google.devtools.build.lib.actions.ActionOwner;
+import com.google.devtools.build.lib.actions.ActionTemplate;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionTemplate.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionTemplate.java
index a3f6cc0d86..7cc7dc32bf 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionTemplate.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileActionTemplate.java
@@ -21,11 +21,11 @@ import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionExecutionException;
 import com.google.devtools.build.lib.actions.ActionInputHelper;
 import com.google.devtools.build.lib.actions.ActionOwner;
+import com.google.devtools.build.lib.actions.ActionTemplate;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.ArtifactOwner;
-import com.google.devtools.build.lib.analysis.actions.ActionTemplate;
 import com.google.devtools.build.lib.analysis.config.InvalidConfigurationException;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunction.java
index 6822dafb59..e3081aaf81 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunction.java
@@ -20,13 +20,13 @@ import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
 import com.google.devtools.build.lib.actions.ActionKeyContext;
 import com.google.devtools.build.lib.actions.ActionLookupValue;
+import com.google.devtools.build.lib.actions.ActionTemplate;
+import com.google.devtools.build.lib.actions.ActionTemplate.ActionTemplateExpansionException;
 import com.google.devtools.build.lib.actions.Actions;
 import com.google.devtools.build.lib.actions.Actions.GeneratingActions;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.ArtifactPrefixConflictException;
 import com.google.devtools.build.lib.actions.MutableActionGraph.ActionConflictException;
-import com.google.devtools.build.lib.analysis.actions.ActionTemplate;
-import com.google.devtools.build.lib.analysis.actions.ActionTemplate.ActionTemplateExpansionException;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.skyframe.ActionTemplateExpansionValue.ActionTemplateExpansionKey;
 import com.google.devtools.build.skyframe.SkyFunction;
@@ -57,7 +57,12 @@ public class ActionTemplateExpansionFunction implements SkyFunction {
   public SkyValue compute(SkyKey skyKey, Environment env)
       throws ActionTemplateExpansionFunctionException, InterruptedException {
     ActionTemplateExpansionKey key = (ActionTemplateExpansionKey) skyKey.argument();
-    ActionTemplate actionTemplate = key.getActionTemplate();
+    ActionLookupValue value = (ActionLookupValue) env.getValue(key.getActionLookupKey());
+    if (value == null) {
+      // Shouldn't actually happen in practice, but tolerate.
+      return null;
+    }
+    ActionTemplate<?> actionTemplate = value.getActionTemplate(key.getActionIndex());
 
     // Requests the TreeArtifactValue object for the input TreeArtifact.
     SkyKey artifactValueKey = ArtifactSkyKey.key(actionTemplate.getInputTreeArtifact(), true);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionValue.java
index 0897116f7b..b8feb28ffa 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionValue.java
@@ -13,13 +13,12 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
-import com.google.common.base.Preconditions;
+import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionKeyContext;
 import com.google.devtools.build.lib.actions.ActionLookupValue;
-import com.google.devtools.build.lib.analysis.actions.ActionTemplate;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.skyframe.SkyFunctionName;
@@ -38,19 +37,17 @@ public final class ActionTemplateExpansionValue extends ActionLookupValue {
   private static final Interner<ActionTemplateExpansionKey> interner =
       BlazeInterners.newWeakInterner();
 
-  static ActionTemplateExpansionKey key(ActionTemplate<?> actionTemplate) {
-    return interner.intern(new ActionTemplateExpansionKey(actionTemplate));
+  static ActionTemplateExpansionKey key(ActionLookupKey actionLookupKey, int actionIndex) {
+    return interner.intern(new ActionTemplateExpansionKey(actionLookupKey, actionIndex));
   }
 
   static final class ActionTemplateExpansionKey extends ActionLookupKey {
-    private final ActionTemplate<?> actionTemplate;
+    private final ActionLookupKey actionLookupKey;
+    private final int actionIndex;
 
-    private ActionTemplateExpansionKey(ActionTemplate<?> actionTemplate) {
-      Preconditions.checkNotNull(
-          actionTemplate,
-          "Passed in action template cannot be null: %s",
-          actionTemplate);
-      this.actionTemplate = actionTemplate;
+    private ActionTemplateExpansionKey(ActionLookupKey actionLookupKey, int actionIndex) {
+      this.actionLookupKey = actionLookupKey;
+      this.actionIndex = actionIndex;
     }
 
     @Override
@@ -60,32 +57,45 @@ public final class ActionTemplateExpansionValue extends ActionLookupValue {
 
     @Override
     public Label getLabel() {
-      return actionTemplate.getOwner().getLabel();
+      return actionLookupKey.getLabel();
     }
 
-    /** Returns the associated {@link ActionTemplate} */
-    ActionTemplate<?> getActionTemplate() {
-      return actionTemplate;
+    ActionLookupKey getActionLookupKey() {
+      return actionLookupKey;
+    }
+
+    /**
+     * Index of the action in question in the node keyed by {@link #getActionLookupKey}. Should be
+     * passed to {@link ActionLookupValue#getAction}.
+     */
+    int getActionIndex() {
+      return actionIndex;
     }
 
     @Override
     public int hashCode() {
-      return actionTemplate.hashCode();
+      return 37 * actionLookupKey.hashCode() + actionIndex;
     }
 
     @Override
     public boolean equals(Object obj) {
-     if (this == obj) {
-       return true;
-     }
-      if (obj == null) {
-        return false;
+      if (this == obj) {
+        return true;
       }
       if (!(obj instanceof ActionTemplateExpansionKey)) {
         return false;
       }
-      ActionTemplateExpansionKey other = (ActionTemplateExpansionKey) obj;
-      return actionTemplate.equals(other.actionTemplate);
+      ActionTemplateExpansionKey that = (ActionTemplateExpansionKey) obj;
+      return this.actionIndex == that.actionIndex
+          && this.actionLookupKey.equals(that.actionLookupKey);
+    }
+
+    @Override
+    public String toString() {
+      return MoreObjects.toStringHelper(this)
+          .add("actionLookupKey", actionLookupKey)
+          .add("actionIndex", actionIndex)
+          .toString();
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java
index 6e6e4321ea..5cef0fc2c6 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java
@@ -29,12 +29,10 @@ import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.ArtifactOwner;
 import com.google.devtools.build.lib.actions.MissingInputFileException;
-import com.google.devtools.build.lib.analysis.actions.ActionTemplate;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.util.Pair;
-import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyFunctionException;
@@ -84,26 +82,21 @@ class ArtifactFunction implements SkyFunction {
 
     // If the action is an ActionTemplate, we need to expand the ActionTemplate into concrete
     // actions, execute those actions in parallel and then aggregate the action execution results.
-    if (artifact.isTreeArtifact()) {
-      ActionAnalysisMetadata actionMetadata =
-          actionLookupValue.getIfPresentAndNotAction(actionIndex);
-      if (actionMetadata instanceof ActionTemplate) {
-        // Create the directory structures for the output TreeArtifact first.
-        try {
-          FileSystemUtils.createDirectoryAndParents(artifact.getPath());
-        } catch (IOException e) {
-          env.getListener()
-              .handle(
-                  Event.error(
-                      String.format(
-                          "Failed to create output directory for TreeArtifact %s: %s",
-                          artifact, e.getMessage())));
-          throw new ArtifactFunctionException(e, Transience.TRANSIENT);
-        }
-
-        return createTreeArtifactValueFromActionTemplate(
-            (ActionTemplate<?>) actionMetadata, artifact, env);
+    if (artifact.isTreeArtifact() && actionLookupValue.isActionTemplate(actionIndex)) {
+      // Create the directory structures for the output TreeArtifact first.
+      try {
+        artifact.getPath().createDirectoryAndParents();
+      } catch (IOException e) {
+        env.getListener()
+            .handle(
+                Event.error(
+                    String.format(
+                        "Failed to create output directory for TreeArtifact %s: %s",
+                        artifact, e.getMessage())));
+        throw new ArtifactFunctionException(e, Transience.TRANSIENT);
       }
+
+      return createTreeArtifactValueFromActionKey(actionLookupKey, actionIndex, artifact, env);
     }
     ActionExecutionValue actionValue =
         (ActionExecutionValue) env.getValue(ActionExecutionValue.key(actionLookupKey, actionIndex));
@@ -133,12 +126,15 @@ class ArtifactFunction implements SkyFunction {
     return createSimpleFileArtifactValue(artifact, actionValue);
   }
 
-  private static TreeArtifactValue createTreeArtifactValueFromActionTemplate(
-      final ActionTemplate<?> actionTemplate, final Artifact treeArtifact, Environment env)
-          throws InterruptedException {
+  private static TreeArtifactValue createTreeArtifactValueFromActionKey(
+      ActionLookupKey actionLookupKey,
+      int actionIndex,
+      final Artifact treeArtifact,
+      Environment env)
+      throws InterruptedException {
     // Request the list of expanded actions from the ActionTemplate.
     ActionTemplateExpansionValue.ActionTemplateExpansionKey templateKey =
-        ActionTemplateExpansionValue.key(actionTemplate);
+        ActionTemplateExpansionValue.key(actionLookupKey, actionIndex);
     ActionTemplateExpansionValue expansionValue =
         (ActionTemplateExpansionValue) env.getValue(templateKey);
 
@@ -166,9 +162,10 @@ class ArtifactFunction implements SkyFunction {
           (ActionExecutionValue)
               Preconditions.checkNotNull(
                   expandedActionValueMap.get(expandedActionExecutionKeys.get(i)),
-                  "Missing tree value: %s %s %s %s",
+                  "Missing tree value: %s %s %s %s %s",
                   treeArtifact,
-                  actionTemplate,
+                  actionLookupKey,
+                  actionIndex,
                   expansionValue,
                   expandedActionValueMap);
       Iterable<TreeFileArtifact> treeFileArtifacts =
@@ -180,11 +177,12 @@ class ArtifactFunction implements SkyFunction {
                     public boolean apply(Artifact artifact) {
                       Preconditions.checkState(
                           artifact.hasParent(),
-                          "No parent: %s %s %s %s",
+                          "No parent: %s %s %s %s %s",
                           artifact,
                           treeArtifact,
                           actionExecutionValue,
-                          actionTemplate);
+                          actionLookupKey,
+                          actionIndex);
                       return artifact.getParent().equals(treeArtifact);
                     }
                   }),