If --batch, --keep_going, --discard_analysis_cache, and the new --noexperimental_enable_critical_path_profiling flags are all specified, then Bazel will delete Actions from ActionLookupValues as they are executed in order to save memory.

Because an already-run action may output an artifact that is only requested later in the build, we need to maintain a way for the artifact to look up the action. But in most cases we don't need to keep the action itself, just its output metadata.

Some actions unfortunately are needed post-execution, and so we special-case them.

Also includes dependency change with description:
Move action out of key. This keeps action references from polluting the graph -- actions are just stored in one SkyValue, instead of being present in SkyKeys.

This does mean additional memory used: we have a separate ActionLookupData object per Action executed. That may reach ~24M for million-action builds.

PiperOrigin-RevId: 151756383
diff --git a/src/test/java/com/google/devtools/build/lib/actions/ActionLookupValueTest.java b/src/test/java/com/google/devtools/build/lib/actions/ActionLookupValueTest.java
new file mode 100644
index 0000000..f7392e2
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/actions/ActionLookupValueTest.java
@@ -0,0 +1,87 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// 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.actions;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.vfs.FileSystem;
+import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Basic tests for {@link ActionLookupValue}. */
+@RunWith(JUnit4.class)
+public class ActionLookupValueTest {
+
+  private FileSystem fs;
+
+  @Before
+  public void setUp() {
+    fs = new InMemoryFileSystem();
+  }
+
+  @Test
+  public void testActionPresentAfterEvaluation() {
+    Action action = mock(Action.class);
+    Artifact artifact = mock(Artifact.class);
+    when(action.getOutputs()).thenReturn(ImmutableSet.of(artifact));
+    when(action.canRemoveAfterExecution()).thenReturn(true);
+    ActionLookupValue underTest = new ActionLookupValue(action, false);
+    assertThat(underTest.getGeneratingActionIndex(artifact)).isEqualTo(0);
+    assertThat(underTest.getAction(0)).isSameAs(action);
+    underTest.actionEvaluated(0, action);
+    assertThat(underTest.getAction(0)).isSameAs(action);
+  }
+
+  @Test
+  public void testActionNotPresentAfterEvaluation() {
+    Root root = Root.asDerivedRoot(fs.getRootDirectory());
+    Action normalAction = mock(Action.class);
+    Artifact normalArtifact = new Artifact(new PathFragment("normal"), root);
+    when(normalAction.getOutputs()).thenReturn(ImmutableSet.of(normalArtifact));
+    when(normalAction.canRemoveAfterExecution()).thenReturn(true);
+    Action persistentAction = mock(Action.class);
+    Artifact persistentOutput = new Artifact(new PathFragment("persistent"), root);
+    when(persistentAction.getOutputs()).thenReturn(ImmutableSet.of(persistentOutput));
+    when(persistentAction.canRemoveAfterExecution()).thenReturn(false);
+    ActionLookupValue underTest =
+        new ActionLookupValue(
+            ImmutableList.<ActionAnalysisMetadata>of(normalAction, persistentAction), true);
+    assertThat(underTest.getGeneratingActionIndex(normalArtifact)).isEqualTo(0);
+    assertThat(underTest.getAction(0)).isSameAs(normalAction);
+    assertThat(underTest.getGeneratingActionIndex(persistentOutput)).isEqualTo(1);
+    assertThat(underTest.getAction(1)).isSameAs(persistentAction);
+    underTest.actionEvaluated(0, normalAction);
+    try {
+      underTest.getAction(0);
+      fail();
+    } catch (NullPointerException e) {
+      // Expected.
+    }
+    assertThat(underTest.getGeneratingActionIndex(persistentOutput)).isEqualTo(1);
+    assertThat(underTest.getAction(1)).isSameAs(persistentAction);
+    underTest.actionEvaluated(1, persistentAction);
+    // Action that said not to clear it won't be cleared.
+    assertThat(underTest.getGeneratingActionIndex(persistentOutput)).isEqualTo(1);
+    assertThat(underTest.getAction(1)).isSameAs(persistentAction);
+  }
+}
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/AspectValueTest.java b/src/test/java/com/google/devtools/build/lib/analysis/AspectValueTest.java
index 0714a0d..4b71f81 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/AspectValueTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/AspectValueTest.java
@@ -15,6 +15,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.testing.EqualsTester;
+import com.google.devtools.build.lib.actions.ActionLookupValue;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.util.AnalysisTestCase;
 import com.google.devtools.build.lib.analysis.util.TestAspects;
@@ -24,7 +25,6 @@
 import com.google.devtools.build.lib.packages.AspectDescriptor;
 import com.google.devtools.build.lib.packages.AspectParameters;
 import com.google.devtools.build.lib.packages.NativeAspectClass;
-import com.google.devtools.build.lib.skyframe.ActionLookupValue;
 import com.google.devtools.build.lib.skyframe.AspectValue;
 import com.google.devtools.build.lib.skyframe.AspectValue.AspectKey;
 import com.google.devtools.build.skyframe.SkyKey;
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/util/AnalysisTestUtil.java b/src/test/java/com/google/devtools/build/lib/analysis/util/AnalysisTestUtil.java
index de1247e..f821e46 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/util/AnalysisTestUtil.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/util/AnalysisTestUtil.java
@@ -147,7 +147,7 @@
     }
 
     @Override
-    public Iterable<ActionAnalysisMetadata> getRegisteredActions() {
+    public List<ActionAnalysisMetadata> getRegisteredActions() {
       return original.getRegisteredActions();
     }
 
@@ -329,7 +329,7 @@
     }
 
     @Override
-    public Iterable<ActionAnalysisMetadata> getRegisteredActions() {
+    public List<ActionAnalysisMetadata> getRegisteredActions() {
       return ImmutableList.of();
     }
 
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/util/BuildViewTestCase.java b/src/test/java/com/google/devtools/build/lib/analysis/util/BuildViewTestCase.java
index e827137..576a312 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/util/BuildViewTestCase.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/util/BuildViewTestCase.java
@@ -37,6 +37,7 @@
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
 import com.google.devtools.build.lib.actions.ActionGraph;
 import com.google.devtools.build.lib.actions.ActionInput;
+import com.google.devtools.build.lib.actions.ActionLookupValue;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactOwner;
 import com.google.devtools.build.lib.actions.MapBasedActionGraph;
@@ -120,7 +121,6 @@
 import com.google.devtools.build.lib.rules.extra.ExtraAction;
 import com.google.devtools.build.lib.rules.test.BaselineCoverageAction;
 import com.google.devtools.build.lib.rules.test.InstrumentedFilesProvider;
-import com.google.devtools.build.lib.skyframe.ActionLookupValue;
 import com.google.devtools.build.lib.skyframe.AspectValue;
 import com.google.devtools.build.lib.skyframe.ConfiguredTargetKey;
 import com.google.devtools.build.lib.skyframe.DiffAwareness;
@@ -1656,7 +1656,7 @@
     }
 
     @Override
-    public Iterable<ActionAnalysisMetadata> getRegisteredActions() {
+    public List<ActionAnalysisMetadata> getRegisteredActions() {
       throw new UnsupportedOperationException();
     }
 
diff --git a/src/test/java/com/google/devtools/build/lib/runtime/ExperimentalStateTrackerTest.java b/src/test/java/com/google/devtools/build/lib/runtime/ExperimentalStateTrackerTest.java
index 5954f43..ef48219 100644
--- a/src/test/java/com/google/devtools/build/lib/runtime/ExperimentalStateTrackerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/runtime/ExperimentalStateTrackerTest.java
@@ -23,6 +23,7 @@
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionCompletionEvent;
 import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
+import com.google.devtools.build.lib.actions.ActionLookupData;
 import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.ActionStartedEvent;
 import com.google.devtools.build.lib.actions.ActionStatusMessage;
@@ -145,7 +146,8 @@
     ExperimentalStateTracker stateTracker = new ExperimentalStateTracker(clock);
     stateTracker.actionStarted(new ActionStartedEvent(fastAction, 123456789));
     stateTracker.actionStarted(new ActionStartedEvent(slowAction, 123456999));
-    stateTracker.actionCompletion(new ActionCompletionEvent(20, fastAction));
+    stateTracker.actionCompletion(
+        new ActionCompletionEvent(20, fastAction, Mockito.mock(ActionLookupData.class)));
 
     LoggingTerminalWriter terminalWriter = new LoggingTerminalWriter(/*discardHighlight=*/ true);
     stateTracker.writeProgressBar(terminalWriter);
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 1b725b3..d24647a 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
@@ -16,6 +16,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.fail;
 
+import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
@@ -74,9 +75,10 @@
     MemoizingEvaluator evaluator =
         new InMemoryMemoizingEvaluator(
             ImmutableMap.<SkyFunctionName, SkyFunction>builder()
-                .put(SkyFunctions.ARTIFACT,
-                    new DummyArtifactFunction(artifactValueMap))
-                .put(SkyFunctions.ACTION_TEMPLATE_EXPANSION, new ActionTemplateExpansionFunction())
+                .put(SkyFunctions.ARTIFACT, new DummyArtifactFunction(artifactValueMap))
+                .put(
+                    SkyFunctions.ACTION_TEMPLATE_EXPANSION,
+                    new ActionTemplateExpansionFunction(Suppliers.ofInstance(false)))
                 .build(),
             differencer);
     driver = new SequentialBuildDriver(evaluator);
@@ -188,7 +190,12 @@
     if (result.hasError()) {
       throw result.getError().getException();
     }
-    return ImmutableList.copyOf(result.get(skyKey).getExpandedActions());
+    ActionTemplateExpansionValue actionTemplateExpansionValue = result.get(skyKey);
+    ImmutableList.Builder<Action> actionList = ImmutableList.builder();
+    for (int i = 0; i < actionTemplateExpansionValue.getNumActions(); i++) {
+      actionList.add(actionTemplateExpansionValue.getAction(i));
+    }
+    return actionList.build();
   }
 
   private Artifact createTreeArtifact(String path) {
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java
index 71bb32a..c729fc6 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTest.java
@@ -32,6 +32,8 @@
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata.MiddlemanType;
 import com.google.devtools.build.lib.actions.ActionInputHelper;
+import com.google.devtools.build.lib.actions.ActionLookupData;
+import com.google.devtools.build.lib.actions.ActionLookupValue;
 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;
@@ -464,9 +466,10 @@
 
   private void setGeneratingActions() {
     if (evaluator.getExistingValueForTesting(OWNER_KEY) == null) {
-      differencer.inject(ImmutableMap.of(
-          OWNER_KEY,
-          new ActionLookupValue(ImmutableList.<ActionAnalysisMetadata>copyOf(actions))));
+      differencer.inject(
+          ImmutableMap.of(
+              OWNER_KEY,
+              new ActionLookupValue(ImmutableList.<ActionAnalysisMetadata>copyOf(actions), false)));
     }
   }
 
@@ -483,11 +486,14 @@
   /** Value Builder for actions that just stats and stores the output file (which must exist). */
   private static class SimpleActionExecutionFunction implements SkyFunction {
     @Override
-    public SkyValue compute(SkyKey skyKey, Environment env) {
+    public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedException {
       Map<Artifact, FileValue> artifactData = new HashMap<>();
       Map<Artifact, TreeArtifactValue> treeArtifactData = new HashMap<>();
       Map<Artifact, FileArtifactValue> additionalOutputData = new HashMap<>();
-      Action action = (Action) skyKey.argument();
+      ActionLookupData actionLookupData = (ActionLookupData) skyKey.argument();
+      ActionLookupValue actionLookupValue =
+          (ActionLookupValue) env.getValue(actionLookupData.getActionLookupNode());
+      Action action = actionLookupValue.getAction(actionLookupData.getActionIndex());
       Artifact output = Iterables.getOnlyElement(action.getOutputs());
 
       try {
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTestCase.java b/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTestCase.java
index a598976..3632017 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTestCase.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/ArtifactFunctionTestCase.java
@@ -15,12 +15,13 @@
 
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
+import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
+import com.google.devtools.build.lib.actions.ActionLookupValue.ActionLookupKey;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
-import com.google.devtools.build.lib.skyframe.ActionLookupValue.ActionLookupKey;
 import com.google.devtools.build.lib.skyframe.ExternalFilesHelper.ExternalFileAction;
 import com.google.devtools.build.lib.skyframe.PackageLookupFunction.CrossRepositoryLabelViolationStrategy;
 import com.google.devtools.build.lib.skyframe.PackageLookupValue.BuildFileName;
@@ -110,7 +111,9 @@
                             TestRuleClassProvider.getRuleClassProvider(), root.getFileSystem()),
                         directories))
                 .put(SkyFunctions.EXTERNAL_PACKAGE, new ExternalPackageFunction())
-                .put(SkyFunctions.ACTION_TEMPLATE_EXPANSION, new ActionTemplateExpansionFunction())
+                .put(
+                    SkyFunctions.ACTION_TEMPLATE_EXPANSION,
+                    new ActionTemplateExpansionFunction(Suppliers.ofInstance(false)))
                 .build(),
             differencer);
     driver = new SequentialBuildDriver(evaluator);
@@ -149,12 +152,12 @@
 
   private static class SingletonActionLookupKey extends ActionLookupKey {
     @Override
-    SkyKey getSkyKeyInternal() {
+    protected SkyKey getSkyKeyInternal() {
       return OWNER_KEY;
     }
 
     @Override
-    SkyFunctionName getType() {
+    protected SkyFunctionName getType() {
       throw new UnsupportedOperationException();
     }
   }
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/FilesystemValueCheckerTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/FilesystemValueCheckerTest.java
index 3f42663..0dba1d9 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/FilesystemValueCheckerTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/FilesystemValueCheckerTest.java
@@ -25,6 +25,8 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.util.concurrent.Runnables;
 import com.google.devtools.build.lib.actions.Action;
+import com.google.devtools.build.lib.actions.ActionLookupValue;
+import com.google.devtools.build.lib.actions.ActionLookupValue.ActionLookupKey;
 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;
@@ -352,14 +354,16 @@
     FileSystemUtils.writeContentAsLatin1(out1.getPath(), "hello");
     FileSystemUtils.writeContentAsLatin1(out2.getPath(), "fizzlepop");
 
-    Action action1 =
-        new TestAction(
-            Runnables.doNothing(), ImmutableSet.<Artifact>of(), ImmutableSet.of(out1));
-    Action action2 =
-        new TestAction(
-            Runnables.doNothing(), ImmutableSet.<Artifact>of(), ImmutableSet.of(out2));
-    SkyKey actionKey1 = ActionExecutionValue.key(action1);
-    SkyKey actionKey2 = ActionExecutionValue.key(action2);
+    SkyKey actionLookupKey =
+        ActionLookupValue.key(
+            new ActionLookupKey() {
+              @Override
+              protected SkyFunctionName getType() {
+                return SkyFunctionName.FOR_TESTING;
+              }
+            });
+    SkyKey actionKey1 = ActionExecutionValue.key(actionLookupKey, 0);
+    SkyKey actionKey2 = ActionExecutionValue.key(actionLookupKey, 1);
     differencer.inject(
         ImmutableMap.<SkyKey, SkyValue>of(
             actionKey1,
@@ -429,27 +433,19 @@
     Artifact last = createTreeArtifact("zzzzzzzzzz");
     FileSystemUtils.createDirectoryAndParents(last.getPath());
 
-    Action action1 =
-        new TestAction(
-            Runnables.doNothing(), ImmutableSet.<Artifact>of(), ImmutableSet.of(out1));
-    Action action2 =
-        new TestAction(
-            Runnables.doNothing(), ImmutableSet.<Artifact>of(), ImmutableSet.of(out2));
-    Action actionEmpty =
-        new TestAction(
-            Runnables.doNothing(), ImmutableSet.<Artifact>of(), ImmutableSet.of(outEmpty));
-    Action actionUnchanging =
-        new TestAction(
-            Runnables.doNothing(), ImmutableSet.<Artifact>of(), ImmutableSet.of(outUnchanging));
-    Action lastAction =
-        new TestAction(
-            Runnables.doNothing(), ImmutableSet.<Artifact>of(), ImmutableSet.of(last));
-
-    SkyKey actionKey1 = ActionExecutionValue.key(action1);
-    SkyKey actionKey2 = ActionExecutionValue.key(action2);
-    SkyKey actionKeyEmpty = ActionExecutionValue.key(actionEmpty);
-    SkyKey actionKeyUnchanging = ActionExecutionValue.key(actionUnchanging);
-    SkyKey actionKeyLast = ActionExecutionValue.key(lastAction);
+    SkyKey actionLookupKey =
+        ActionLookupValue.key(
+            new ActionLookupKey() {
+              @Override
+              protected SkyFunctionName getType() {
+                return SkyFunctionName.FOR_TESTING;
+              }
+            });
+    SkyKey actionKey1 = ActionExecutionValue.key(actionLookupKey, 0);
+    SkyKey actionKey2 = ActionExecutionValue.key(actionLookupKey, 1);
+    SkyKey actionKeyEmpty = ActionExecutionValue.key(actionLookupKey, 2);
+    SkyKey actionKeyUnchanging = ActionExecutionValue.key(actionLookupKey, 3);
+    SkyKey actionKeyLast = ActionExecutionValue.key(actionLookupKey, 4);
     differencer.inject(
         ImmutableMap.<SkyKey, SkyValue>of(
             actionKey1,
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/SkyframeAwareActionTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/SkyframeAwareActionTest.java
index 02c2b41..66c469d 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/SkyframeAwareActionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/SkyframeAwareActionTest.java
@@ -388,7 +388,7 @@
         null);
 
     // Sanity check that our invalidation receiver is working correctly. We'll rely on it again.
-    SkyKey actionKey = ActionExecutionValue.key(action);
+    SkyKey actionKey = ActionExecutionValue.key(OWNER_KEY, 0);
     TrackingEvaluationProgressReceiver.EvaluatedEntry evaluatedAction =
         progressReceiver.getEvalutedEntry(actionKey);
     assertThat(evaluatedAction).isNotNull();
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java b/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
index 2ca64ea..8caead9 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
@@ -17,6 +17,7 @@
 
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
+import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -32,6 +33,8 @@
 import com.google.devtools.build.lib.actions.ActionExecutionStatusReporter;
 import com.google.devtools.build.lib.actions.ActionInputFileCache;
 import com.google.devtools.build.lib.actions.ActionLogBufferPathGenerator;
+import com.google.devtools.build.lib.actions.ActionLookupData;
+import com.google.devtools.build.lib.actions.ActionLookupValue;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.BuildFailedException;
 import com.google.devtools.build.lib.actions.Executor;
@@ -101,7 +104,7 @@
 public abstract class TimestampBuilderTestCase extends FoundationTestCase {
   protected static final ActionLookupValue.ActionLookupKey ALL_OWNER =
       new SingletonActionLookupKey();
-  private static final SkyKey OWNER_KEY = SkyKey.create(SkyFunctions.ACTION_LOOKUP, ALL_OWNER);
+  protected static final SkyKey OWNER_KEY = SkyKey.create(SkyFunctions.ACTION_LOOKUP, ALL_OWNER);
   protected static final Predicate<Action> ALWAYS_EXECUTE_FILTER = Predicates.alwaysTrue();
   protected static final String CYCLE_MSG = "Yarrrr, there be a cycle up in here";
 
@@ -118,7 +121,8 @@
     tsgm = new TimestampGranularityMonitor(clock);
     ResourceManager.instance().setAvailableResources(ResourceSet.createWithRamCpuIo(100, 1, 1));
     actions = new HashSet<>();
-    actionTemplateExpansionFunction = new ActionTemplateExpansionFunction();
+    actionTemplateExpansionFunction =
+        new ActionTemplateExpansionFunction(Suppliers.ofInstance(false));
   }
 
   protected void clearActions() {
@@ -218,7 +222,8 @@
       private void setGeneratingActions() {
         if (evaluator.getExistingValueForTesting(OWNER_KEY) == null) {
           differencer.inject(
-              ImmutableMap.of(OWNER_KEY, new ActionLookupValue(ImmutableList.copyOf(actions))));
+              ImmutableMap.of(
+                  OWNER_KEY, new ActionLookupValue(ImmutableList.copyOf(actions), false)));
         }
       }
 
@@ -443,12 +448,12 @@
 
   private static class SingletonActionLookupKey extends ActionLookupValue.ActionLookupKey {
     @Override
-    SkyKey getSkyKeyInternal() {
+    protected SkyKey getSkyKeyInternal() {
       return OWNER_KEY;
     }
 
     @Override
-    SkyFunctionName getType() {
+    protected SkyFunctionName getType() {
       throw new UnsupportedOperationException();
     }
   }
@@ -477,6 +482,9 @@
   private static final ActionCompletedReceiver EMPTY_COMPLETION_RECEIVER =
       new ActionCompletedReceiver() {
         @Override
-        public void actionCompleted(Action action) {}
+        public void actionCompleted(ActionLookupData actionLookupData) {}
+
+        @Override
+        public void noteActionEvaluationStarted(ActionLookupData actionLookupData, Action action) {}
       };
 }
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 099b307..99da8b9 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
@@ -1217,7 +1217,7 @@
       ActionTemplateExpansionKey key = (ActionTemplateExpansionKey) skyKey.argument();
       ActionTemplate<?> actionTemplate = key.getActionTemplate();
       return new ActionTemplateExpansionValue(
-          Preconditions.checkNotNull(actionTemplateToActionMap.get(actionTemplate)));
+          Preconditions.checkNotNull(actionTemplateToActionMap.get(actionTemplate)), false);
     }
 
     @Override
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactMetadataTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactMetadataTest.java
index 6817ccf..03d6618 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactMetadataTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/TreeArtifactMetadataTest.java
@@ -26,6 +26,8 @@
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
 import com.google.devtools.build.lib.actions.ActionInputHelper;
+import com.google.devtools.build.lib.actions.ActionLookupData;
+import com.google.devtools.build.lib.actions.ActionLookupValue;
 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;
@@ -221,9 +223,10 @@
 
   private void setGeneratingActions() {
     if (evaluator.getExistingValueForTesting(OWNER_KEY) == null) {
-      differencer.inject(ImmutableMap.of(
-          OWNER_KEY,
-          new ActionLookupValue(ImmutableList.<ActionAnalysisMetadata>copyOf(actions))));
+      differencer.inject(
+          ImmutableMap.of(
+              OWNER_KEY,
+              new ActionLookupValue(ImmutableList.<ActionAnalysisMetadata>copyOf(actions), false)));
     }
   }
 
@@ -239,10 +242,14 @@
 
   private class TreeArtifactExecutionFunction implements SkyFunction {
     @Override
-    public SkyValue compute(SkyKey skyKey, Environment env) throws SkyFunctionException {
+    public SkyValue compute(SkyKey skyKey, Environment env)
+        throws SkyFunctionException, InterruptedException {
       Map<Artifact, FileValue> fileData = new HashMap<>();
       Map<TreeFileArtifact, FileArtifactValue> treeArtifactData = new HashMap<>();
-      Action action = (Action) skyKey.argument();
+      ActionLookupData actionLookupData = (ActionLookupData) skyKey.argument();
+      ActionLookupValue actionLookupValue =
+          (ActionLookupValue) env.getValue(actionLookupData.getActionLookupNode());
+      Action action = actionLookupValue.getAction(actionLookupData.getActionIndex());
       Artifact output = Iterables.getOnlyElement(action.getOutputs());
       for (PathFragment subpath : testTreeArtifactContents) {
         try {