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

3 files changed, 46 insertions, 42 deletions

diff --git a/src/test/java/com/google/devtools/build/lib/rules/objc/BazelJ2ObjcLibraryTest.java b/src/test/java/com/google/devtools/build/lib/rules/objc/BazelJ2ObjcLibraryTest.java
index 64e4b46eb7..6a72c21ef9 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/objc/BazelJ2ObjcLibraryTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/objc/BazelJ2ObjcLibraryTest.java
@@ -27,13 +27,13 @@ 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.ActionInputPrefetcher;
+import com.google.devtools.build.lib.actions.ActionTemplate.ActionTemplateExpansionException;
 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.actions.CommandAction;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
-import com.google.devtools.build.lib.analysis.actions.ActionTemplate.ActionTemplateExpansionException;
 import com.google.devtools.build.lib.analysis.actions.ParameterFileWriteAction;
 import com.google.devtools.build.lib.analysis.actions.SpawnAction;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunctionTest.java
index f8b6a9db3d..dba65ebbdc 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/ActionTemplateExpansionFunctionTest.java
@@ -24,6 +24,8 @@ import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionInputHelper;
 import com.google.devtools.build.lib.actions.ActionKeyContext;
+import com.google.devtools.build.lib.actions.ActionTemplate;
+import com.google.devtools.build.lib.actions.Actions;
 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.SpecialArtifactType;
@@ -33,18 +35,22 @@ import com.google.devtools.build.lib.actions.ArtifactPrefixConflictException;
 import com.google.devtools.build.lib.actions.ArtifactRoot;
 import com.google.devtools.build.lib.actions.MutableActionGraph.ActionConflictException;
 import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
+import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.actions.CustomCommandLine;
 import com.google.devtools.build.lib.analysis.actions.SpawnActionTemplate;
 import com.google.devtools.build.lib.analysis.actions.SpawnActionTemplate.OutputPathMapper;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.events.NullEventHandler;
+import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
+import com.google.devtools.build.lib.skyframe.ActionTemplateExpansionValue.ActionTemplateExpansionKey;
 import com.google.devtools.build.lib.testutil.FoundationTestCase;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.skyframe.EvaluationResult;
 import com.google.devtools.build.skyframe.InMemoryMemoizingEvaluator;
 import com.google.devtools.build.skyframe.MemoizingEvaluator;
-import com.google.devtools.build.skyframe.RecordingDifferencer;
 import com.google.devtools.build.skyframe.SequencedRecordingDifferencer;
 import com.google.devtools.build.skyframe.SequentialBuildDriver;
 import com.google.devtools.build.skyframe.SkyFunction;
@@ -60,12 +66,14 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
 
 /** Tests for {@link ActionTemplateExpansionFunction}. */
 @RunWith(JUnit4.class)
 public final class ActionTemplateExpansionFunctionTest extends FoundationTestCase  {
   private Map<Artifact, TreeArtifactValue> artifactValueMap;
   private SequentialBuildDriver driver;
+  private SequencedRecordingDifferencer differencer;
 
   @Before
   public void setUp() throws Exception  {
@@ -76,7 +84,7 @@ public final class ActionTemplateExpansionFunctionTest extends FoundationTestCas
                 rootDirectory.getFileSystem().getPath("/outputbase"),
                 ImmutableList.of(Root.fromPath(rootDirectory)),
                 BazelSkyframeExecutorConstants.BUILD_FILES_BY_PRIORITY));
-    RecordingDifferencer differencer = new SequencedRecordingDifferencer();
+    differencer = new SequencedRecordingDifferencer();
     MemoizingEvaluator evaluator =
         new InMemoryMemoizingEvaluator(
             ImmutableMap.<SkyFunctionName, SkyFunction>builder()
@@ -103,7 +111,7 @@ public final class ActionTemplateExpansionFunctionTest extends FoundationTestCas
     List<Action> actions = evaluate(spawnActionTemplate);
     assertThat(actions).hasSize(3);
 
-    ArtifactOwner owner = ActionTemplateExpansionValue.key(spawnActionTemplate);
+    ArtifactOwner owner = ActionTemplateExpansionValue.key(CTKEY, 0);
     int i = 0;
     for (Action action : actions) {
       String childName = "child" + i;
@@ -184,17 +192,26 @@ public final class ActionTemplateExpansionFunctionTest extends FoundationTestCas
     }
   }
 
+  private static final ConfiguredTargetKey CTKEY =
+      ConfiguredTargetKey.of(Label.parseAbsoluteUnchecked("//foo:foo"), null);
+
   private List<Action> evaluate(SpawnActionTemplate spawnActionTemplate) throws Exception {
-    SkyKey skyKey = ActionTemplateExpansionValue.key(spawnActionTemplate);
-    EvaluationResult<ActionTemplateExpansionValue> result = driver.evaluate(
-        ImmutableList.of(skyKey),
-        false,
-        SkyframeExecutor.DEFAULT_THREAD_COUNT,
-        NullEventHandler.INSTANCE);
+    ConfiguredTargetValue ctValue =
+        createConfiguredTargetValue(
+            spawnActionTemplate.getOutputTreeArtifact(), spawnActionTemplate);
+
+    differencer.inject(CTKEY, ctValue);
+    ActionTemplateExpansionKey templateKey = ActionTemplateExpansionValue.key(CTKEY, 0);
+    EvaluationResult<ActionTemplateExpansionValue> result =
+        driver.evaluate(
+            ImmutableList.of(templateKey),
+            false,
+            SkyframeExecutor.DEFAULT_THREAD_COUNT,
+            NullEventHandler.INSTANCE);
     if (result.hasError()) {
       throw result.getError().getException();
     }
-    ActionTemplateExpansionValue actionTemplateExpansionValue = result.get(skyKey);
+    ActionTemplateExpansionValue actionTemplateExpansionValue = result.get(templateKey);
     ImmutableList.Builder<Action> actionList = ImmutableList.builder();
     for (int i = 0; i < actionTemplateExpansionValue.getNumActions(); i++) {
       actionList.add(actionTemplateExpansionValue.getAction(i));
@@ -202,6 +219,16 @@ public final class ActionTemplateExpansionFunctionTest extends FoundationTestCas
     return actionList.build();
   }
 
+  private static ConfiguredTargetValue createConfiguredTargetValue(
+      Artifact artifact, ActionTemplate<?> actionTemplate) {
+    return new ConfiguredTargetValue(
+        Mockito.mock(ConfiguredTarget.class),
+        new Actions.GeneratingActions(
+            ImmutableList.of(actionTemplate), ImmutableMap.of(artifact, 0)),
+        NestedSetBuilder.<Package>stableOrder().build(),
+        /*removeActionsAfterEvaluation=*/ false);
+  }
+
   private SpecialArtifact createTreeArtifact(String path) {
     PathFragment execPath = PathFragment.create("out").getRelative(path);
     return new SpecialArtifact(
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactBuildTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactBuildTest.java
index 4928d73fa0..411fef60fb 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactBuildTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactBuildTest.java
@@ -22,10 +22,8 @@ import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
-import com.google.common.collect.Multimap;
 import com.google.common.hash.Hashing;
 import com.google.common.util.concurrent.Runnables;
 import com.google.devtools.build.lib.actions.Action;
@@ -46,12 +44,10 @@ import com.google.devtools.build.lib.actions.cache.MetadataHandler;
 import com.google.devtools.build.lib.actions.util.ActionsTestUtil;
 import com.google.devtools.build.lib.actions.util.TestAction;
 import com.google.devtools.build.lib.actions.util.TestAction.DummyAction;
-import com.google.devtools.build.lib.analysis.actions.ActionTemplate;
 import com.google.devtools.build.lib.analysis.actions.SpawnActionTemplate;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventKind;
 import com.google.devtools.build.lib.events.StoredEventHandler;
-import com.google.devtools.build.lib.skyframe.ActionTemplateExpansionValue.ActionTemplateExpansionKey;
 import com.google.devtools.build.lib.skyframe.serialization.testutils.ObjectCodecTester;
 import com.google.devtools.build.lib.testutil.TestUtils;
 import com.google.devtools.build.lib.vfs.FileStatus;
@@ -754,10 +750,7 @@ public class TreeArtifactBuildTest extends TimestampBuilderTestCase {
 
     actionTemplateExpansionFunction =
         new DummyActionTemplateExpansionFunction(
-            actionKeyContext,
-            ImmutableMultimap.<ActionTemplate<?>, Action>of(
-                actionTemplate, generateOutputAction,
-                actionTemplate, noGenerateOutputAction));
+            actionKeyContext, ImmutableList.of(generateOutputAction, noGenerateOutputAction));
 
     buildArtifact(artifact2);
   }
@@ -796,10 +789,7 @@ public class TreeArtifactBuildTest extends TimestampBuilderTestCase {
 
     actionTemplateExpansionFunction =
         new DummyActionTemplateExpansionFunction(
-            actionKeyContext,
-            ImmutableMultimap.<ActionTemplate<?>, Action>of(
-                actionTemplate, generateOutputAction,
-                actionTemplate, noGenerateOutputAction));
+            actionKeyContext, ImmutableList.of(generateOutputAction, noGenerateOutputAction));
 
     try {
       buildArtifact(artifact2);
@@ -843,10 +833,7 @@ public class TreeArtifactBuildTest extends TimestampBuilderTestCase {
 
     actionTemplateExpansionFunction =
         new DummyActionTemplateExpansionFunction(
-            actionKeyContext,
-            ImmutableMultimap.<ActionTemplate<?>, Action>of(
-                actionTemplate, generateOutputAction,
-                actionTemplate, throwingAction));
+            actionKeyContext, ImmutableList.of(generateOutputAction, throwingAction));
 
     try {
       buildArtifact(artifact2);
@@ -889,10 +876,7 @@ public class TreeArtifactBuildTest extends TimestampBuilderTestCase {
 
     actionTemplateExpansionFunction =
         new DummyActionTemplateExpansionFunction(
-            actionKeyContext,
-            ImmutableMultimap.<ActionTemplate<?>, Action>of(
-                actionTemplate, throwingAction,
-                actionTemplate, anotherThrowingAction));
+            actionKeyContext, ImmutableList.of(throwingAction, anotherThrowingAction));
 
     try {
       buildArtifact(artifact2);
@@ -1239,23 +1223,16 @@ public class TreeArtifactBuildTest extends TimestampBuilderTestCase {
   /** A dummy action template expansion function that just returns the injected actions */
   private static class DummyActionTemplateExpansionFunction implements SkyFunction {
     private final ActionKeyContext actionKeyContext;
-    private final Multimap<ActionTemplate<?>, Action> actionTemplateToActionMap;
+    private final List<Action> actions;
 
-    DummyActionTemplateExpansionFunction(
-        ActionKeyContext actionKeyContext,
-        Multimap<ActionTemplate<?>, Action> actionTemplateToActionMap) {
+    DummyActionTemplateExpansionFunction(ActionKeyContext actionKeyContext, List<Action> actions) {
       this.actionKeyContext = actionKeyContext;
-      this.actionTemplateToActionMap = actionTemplateToActionMap;
+      this.actions = actions;
     }
 
     @Override
     public SkyValue compute(SkyKey skyKey, Environment env) {
-      ActionTemplateExpansionKey key = (ActionTemplateExpansionKey) skyKey.argument();
-      ActionTemplate<?> actionTemplate = key.getActionTemplate();
-      return new ActionTemplateExpansionValue(
-          actionKeyContext,
-          Preconditions.checkNotNull(actionTemplateToActionMap.get(actionTemplate)),
-          false);
+      return new ActionTemplateExpansionValue(actionKeyContext, actions, false);
     }
 
     @Override