Model the TopLevelArtifactContext as an argument to the CompletionFunction rather than a PRECOMPUTED value.

Having a stale TopLevelArtifactContext leads to invalidation of all the top level target nodes, causing time wasted due to a lot of cache hits for a null build.

--
MOS_MIGRATED_REVID=127585059

12 files changed, 161 insertions, 59 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
index 034651422b..8272143f9f 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
@@ -198,9 +198,13 @@ public final class BuildTool {
 
       // Execution phase.
       if (needsExecutionPhase(request.getBuildOptions())) {
-        env.getSkyframeExecutor().injectTopLevelContext(request.getTopLevelArtifactContext());
-        executionTool.executeBuild(request.getId(), analysisResult, result,
-            configurations, analysisResult.getPackageRoots());
+        executionTool.executeBuild(
+            request.getId(),
+            analysisResult,
+            result,
+            configurations,
+            analysisResult.getPackageRoots(),
+            request.getTopLevelArtifactContext());
       }
 
       String delayedErrorMsg = analysisResult.getError();
diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
index 0a8a0cbbb9..2d1e52c08d 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
@@ -55,6 +55,7 @@ import com.google.devtools.build.lib.analysis.BuildView;
 import com.google.devtools.build.lib.analysis.BuildView.AnalysisResult;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.SymlinkTreeActionContext;
+import com.google.devtools.build.lib.analysis.TopLevelArtifactContext;
 import com.google.devtools.build.lib.analysis.TopLevelArtifactHelper;
 import com.google.devtools.build.lib.analysis.WorkspaceStatusAction;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
@@ -336,7 +337,8 @@ public class ExecutionTool {
   void executeBuild(UUID buildId, AnalysisResult analysisResult,
       BuildResult buildResult,
       BuildConfigurationCollection configurations,
-      ImmutableMap<PackageIdentifier, Path> packageRoots)
+      ImmutableMap<PackageIdentifier, Path> packageRoots,
+      TopLevelArtifactContext topLevelArtifactContext)
       throws BuildFailedException, InterruptedException, TestExecException, AbruptExitException {
     Stopwatch timer = Stopwatch.createStarted();
     prepare(packageRoots);
@@ -438,7 +440,8 @@ public class ExecutionTool {
           executor,
           builtTargets,
           request.getBuildOptions().explanationPath != null,
-          env.getBlazeWorkspace().getLastExecutionTimeRange());
+          env.getBlazeWorkspace().getLastExecutionTimeRange(),
+          topLevelArtifactContext);
       buildCompleted = true;
     } catch (BuildFailedException | TestExecException e) {
       buildCompleted = true;
diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/SkyframeBuilder.java b/src/main/java/com/google/devtools/build/lib/buildtool/SkyframeBuilder.java
index afc68b17b2..ba0fbe6217 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/SkyframeBuilder.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/SkyframeBuilder.java
@@ -29,6 +29,7 @@ import com.google.devtools.build.lib.actions.MissingInputFileException;
 import com.google.devtools.build.lib.actions.ResourceManager;
 import com.google.devtools.build.lib.actions.TestExecException;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
+import com.google.devtools.build.lib.analysis.TopLevelArtifactContext;
 import com.google.devtools.build.lib.buildtool.buildevent.ExecutionProgressReceiverAvailableEvent;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.events.Reporter;
@@ -101,7 +102,8 @@ public class SkyframeBuilder implements Builder {
       Executor executor,
       Set<ConfiguredTarget> builtTargets,
       boolean explain,
-      @Nullable Range<Long> lastExecutionTimeRange)
+      @Nullable Range<Long> lastExecutionTimeRange,
+      TopLevelArtifactContext topLevelArtifactContext)
       throws BuildFailedException, AbruptExitException, TestExecException, InterruptedException {
     skyframeExecutor.prepareExecution(modifiedOutputFiles, lastExecutionTimeRange);
     skyframeExecutor.setFileCache(fileCache);
@@ -146,7 +148,8 @@ public class SkyframeBuilder implements Builder {
               finalizeActionsToOutputService,
               numJobs,
               actionCacheChecker,
-              executionProgressReceiver);
+              executionProgressReceiver,
+              topLevelArtifactContext);
       // progressReceiver is finished, so unsynchronized access to builtTargets is now safe.
       Optional<ExitCode> exitCode = processResult(reporter, result, keepGoing, skyframeExecutor);
 
@@ -185,7 +188,8 @@ public class SkyframeBuilder implements Builder {
                 finalizeActionsToOutputService,
                 numJobs,
                 actionCacheChecker,
-                null);
+                null,
+                topLevelArtifactContext);
         exitCode = processResult(reporter, result, keepGoing, skyframeExecutor);
         Preconditions.checkState(
             exitCode != null || !result.keyNames().isEmpty(),
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionArtifactCycleReporter.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionArtifactCycleReporter.java
index 91574fc16c..b71321cf7f 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionArtifactCycleReporter.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionArtifactCycleReporter.java
@@ -18,10 +18,10 @@ import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.actions.Action;
-import com.google.devtools.build.lib.analysis.LabelAndConfiguration;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.pkgcache.PackageProvider;
 import com.google.devtools.build.lib.skyframe.ArtifactValue.OwnedArtifact;
+import com.google.devtools.build.lib.skyframe.TargetCompletionValue.TargetCompletionKey;
 import com.google.devtools.build.skyframe.CycleInfo;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
@@ -50,9 +50,9 @@ public class ActionArtifactCycleReporter extends AbstractLabelCycleReporter {
       return "file: " + ((OwnedArtifact) arg).getArtifact().getRootRelativePathString();
     } else if (arg instanceof Action) {
       return "action: " + ((Action) arg).getMnemonic();
-    } else if (arg instanceof LabelAndConfiguration
+    } else if (arg instanceof TargetCompletionKey
         && skyFunctionName.equals(SkyFunctions.TARGET_COMPLETION)) {
-      return "configured target: " + ((LabelAndConfiguration) arg).getLabel();
+      return "configured target: " + ((TargetCompletionKey) arg).labelAndConfiguration().getLabel();
     }
     throw new IllegalStateException(
         "Argument is not Action, TargetCompletion,  or OwnedArtifact: " + arg);
@@ -60,7 +60,7 @@ public class ActionArtifactCycleReporter extends AbstractLabelCycleReporter {
 
   @Override
   protected Label getLabel(SkyKey key) {
-    Object arg = key.argument(); 
+    Object arg = key.argument();
     if (arg instanceof OwnedArtifact) {
       return ((OwnedArtifact) arg).getArtifact().getOwner();
     } else if (arg instanceof Action) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/AspectCompletionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/AspectCompletionValue.java
index 22b0ee042f..45c17faf99 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/AspectCompletionValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/AspectCompletionValue.java
@@ -13,8 +13,11 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.auto.value.AutoValue;
 import com.google.common.base.Function;
 import com.google.common.collect.Iterables;
+import com.google.devtools.build.lib.analysis.TopLevelArtifactContext;
+import com.google.devtools.build.lib.skyframe.AspectValue.AspectKey;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 
@@ -34,14 +37,29 @@ public class AspectCompletionValue implements SkyValue {
     return aspectValue;
   }
 
-  public static Iterable<SkyKey> keys(Collection<AspectValue> targets) {
+  public static Iterable<SkyKey> keys(
+      Collection<AspectValue> targets, final TopLevelArtifactContext ctx) {
     return Iterables.transform(
         targets,
         new Function<AspectValue, SkyKey>() {
           @Override
           public SkyKey apply(AspectValue aspectValue) {
-            return SkyKey.create(SkyFunctions.ASPECT_COMPLETION, aspectValue.getKey());
+            return SkyKey.create(
+                SkyFunctions.ASPECT_COMPLETION,
+                AspectCompletionKey.create(aspectValue.getKey(), ctx));
           }
         });
   }
+
+ @AutoValue
+ abstract static class AspectCompletionKey {
+    public static AspectCompletionKey create(
+        AspectKey aspectKey, TopLevelArtifactContext topLevelArtifactContext) {
+      return new AutoValue_AspectCompletionValue_AspectCompletionKey(
+          aspectKey, topLevelArtifactContext);
+    }
+
+    public abstract AspectKey aspectKey();
+    public abstract TopLevelArtifactContext topLevelArtifactContext();
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/Builder.java b/src/main/java/com/google/devtools/build/lib/skyframe/Builder.java
index e6f1c3dca9..d25341a4e7 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/Builder.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/Builder.java
@@ -20,6 +20,7 @@ import com.google.devtools.build.lib.actions.BuildFailedException;
 import com.google.devtools.build.lib.actions.Executor;
 import com.google.devtools.build.lib.actions.TestExecException;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
+import com.google.devtools.build.lib.analysis.TopLevelArtifactContext;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadCompatible;
 import com.google.devtools.build.lib.events.Reporter;
 import com.google.devtools.build.lib.util.AbruptExitException;
@@ -64,6 +65,8 @@ public interface Builder {
    *        valid even if a future action throws ActionExecutionException
    * @param lastExecutionTimeRange If not null, the start/finish time of the last build that
    *        run the execution phase.
+   * @param topLevelArtifactContext contains the the options which determine the artifacts to build
+   *        for the top-level targets.
    * @throws BuildFailedException if there were problems establishing the action execution
    *         environment, if the the metadata of any file  during the build could not be obtained,
    *         if any input files are missing, or if an action fails during execution
@@ -81,6 +84,7 @@ public interface Builder {
       Executor executor,
       Set<ConfiguredTarget> builtTargets,
       boolean explain,
-      @Nullable Range<Long> lastExecutionTimeRange)
+      @Nullable Range<Long> lastExecutionTimeRange,
+      TopLevelArtifactContext topLevelArtifactContext)
       throws BuildFailedException, AbruptExitException, InterruptedException, TestExecException;
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java
index f19a572f98..1f48c84f3b 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java
@@ -28,7 +28,9 @@ import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.skyframe.AspectCompletionValue.AspectCompletionKey;
 import com.google.devtools.build.lib.skyframe.AspectValue.AspectKey;
+import com.google.devtools.build.lib.skyframe.TargetCompletionValue.TargetCompletionKey;
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyFunctionException;
 import com.google.devtools.build.skyframe.SkyKey;
@@ -57,7 +59,21 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
     TValue getValueFromSkyKey(SkyKey skyKey, Environment env);
 
     /**
-     * Returns all artefacts that need to be built to complete the {@code value}
+     * Returns the options which determine the artifacts to build for the top-level targets.
+     * <p>
+     * For the Top level targets we made a conscious decision to include the TopLevelArtifactContext
+     * within the SkyKey as an argument to the CompletionFunction rather than a separate SkyKey.
+     * As a result we do have <num top level targets> extra SkyKeys for every unique
+     * TopLevelArtifactContexts used over the lifetime of Blaze. This is a minor tradeoff,
+     * since it significantly improves null build times when we're switching the
+     * TopLevelArtifactContexts frequently (common for IDEs), by reusing existing SkyKeys
+     * from earlier runs, instead of causing an eager invalidation
+     * were the TopLevelArtifactContext modeled as a separate SkyKey.
+     */
+    TopLevelArtifactContext getTopLevelArtifactContext(SkyKey skyKey);
+
+    /**
+     * Returns all artifacts that need to be built to complete the {@code value}
      */
     ArtifactsToBuild getAllArtifactsToBuild(TValue value, TopLevelArtifactContext context);
 
@@ -80,18 +96,30 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
      * Creates a failed completion value.
      */
     SkyValue createFailed(TValue value, NestedSet<Label> rootCauses);
+
+    /**
+     * Extracts a tag given the {@link SkyKey}.
+     */
+    String extractTag(SkyKey skyKey);
   }
 
   private static class TargetCompletor
       implements Completor<ConfiguredTargetValue, TargetCompletionValue> {
     @Override
     public ConfiguredTargetValue getValueFromSkyKey(SkyKey skyKey, Environment env) {
-      LabelAndConfiguration lac = (LabelAndConfiguration) skyKey.argument();
+      TargetCompletionKey tcKey = (TargetCompletionKey) skyKey.argument();
+      LabelAndConfiguration lac = tcKey.labelAndConfiguration();
       return (ConfiguredTargetValue)
           env.getValue(ConfiguredTargetValue.key(lac.getLabel(), lac.getConfiguration()));
     }
 
     @Override
+    public TopLevelArtifactContext getTopLevelArtifactContext(SkyKey skyKey) {
+      TargetCompletionKey tcKey = (TargetCompletionKey) skyKey.argument();
+      return tcKey.topLevelArtifactContext();
+    }
+
+    @Override
     public ArtifactsToBuild getAllArtifactsToBuild(
         ConfiguredTargetValue value, TopLevelArtifactContext topLevelContext) {
       return TopLevelArtifactHelper.getAllArtifactsToBuild(
@@ -126,16 +154,29 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
     public SkyValue createFailed(ConfiguredTargetValue value, NestedSet<Label> rootCauses) {
       return TargetCompleteEvent.createFailed(value.getConfiguredTarget(), rootCauses);
     }
+
+    @Override
+    public String extractTag(SkyKey skyKey) {
+      return Label.print(
+          ((TargetCompletionKey) skyKey.argument()).labelAndConfiguration().getLabel());
+    }
   }
 
   private static class AspectCompletor implements Completor<AspectValue, AspectCompletionValue> {
     @Override
     public AspectValue getValueFromSkyKey(SkyKey skyKey, Environment env) {
-      AspectKey aspectKey = (AspectKey) skyKey.argument();
+      AspectCompletionKey acKey = (AspectCompletionKey) skyKey.argument();
+      AspectKey aspectKey = acKey.aspectKey();
       return (AspectValue) env.getValue(AspectValue.key(aspectKey));
     }
 
     @Override
+    public TopLevelArtifactContext getTopLevelArtifactContext(SkyKey skyKey) {
+      AspectCompletionKey acKey = (AspectCompletionKey) skyKey.argument();
+      return acKey.topLevelArtifactContext();
+    }
+
+    @Override
     public ArtifactsToBuild getAllArtifactsToBuild(
         AspectValue value, TopLevelArtifactContext topLevelArtifactContext) {
       return TopLevelArtifactHelper.getAllArtifactsToBuild(value, topLevelArtifactContext);
@@ -172,6 +213,11 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
     public SkyValue createFailed(AspectValue value, NestedSet<Label> rootCauses) {
       return AspectCompleteEvent.createFailed(value, rootCauses);
     }
+
+    @Override
+    public String extractTag(SkyKey skyKey) {
+      return Label.print(((AspectCompletionKey) skyKey.argument()).aspectKey().getLabel());
+    }
   }
 
   public static SkyFunction targetCompletionFunction(AtomicReference<EventBus> eventBusRef) {
@@ -195,7 +241,7 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
   @Override
   public SkyValue compute(SkyKey skyKey, Environment env) throws CompletionFunctionException {
     TValue value = completor.getValueFromSkyKey(skyKey, env);
-    TopLevelArtifactContext topLevelContext = PrecomputedValue.TOP_LEVEL_CONTEXT.get(env);
+    TopLevelArtifactContext topLevelContext = completor.getTopLevelArtifactContext(skyKey);
     if (env.valuesMissing()) {
       return null;
     }
@@ -250,7 +296,7 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
 
   @Override
   public String extractTag(SkyKey skyKey) {
-    return Label.print(((LabelAndConfiguration) skyKey.argument()).getLabel());
+    return completor.extractTag(skyKey);
   }
 
   private static final class CompletionFunctionException extends SkyFunctionException {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValue.java
index 3fe34d4288..a5f7a6793d 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/PrecomputedValue.java
@@ -20,7 +20,6 @@ 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.analysis.BlazeDirectories;
-import com.google.devtools.build.lib.analysis.TopLevelArtifactContext;
 import com.google.devtools.build.lib.analysis.WorkspaceStatusAction;
 import com.google.devtools.build.lib.analysis.buildinfo.BuildInfoFactory;
 import com.google.devtools.build.lib.analysis.buildinfo.BuildInfoFactory.BuildInfoKey;
@@ -92,9 +91,6 @@ public final class PrecomputedValue implements SkyValue {
   static final Precomputed<ImmutableList<ActionAnalysisMetadata>> COVERAGE_REPORT_KEY =
       new Precomputed<>(SkyKey.create(SkyFunctions.PRECOMPUTED, "coverage_report_actions"));
 
-  static final Precomputed<TopLevelArtifactContext> TOP_LEVEL_CONTEXT =
-      new Precomputed<>(SkyKey.create(SkyFunctions.PRECOMPUTED, "top_level_context"));
-
   public static final Precomputed<Map<BuildInfoKey, BuildInfoFactory>> BUILD_INFO_FACTORIES =
       new Precomputed<>(SkyKey.create(SkyFunctions.PRECOMPUTED, "build_info_factories"));
 
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
index 4168ef143f..f9426d1b70 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
@@ -633,13 +633,6 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
     PrecomputedValue.DEFAULTS_PACKAGE_CONTENTS.set(injectable(), defaultsPackageContents);
   }
 
-  /**
-   * Injects the top-level artifact options.
-   */
-  public void injectTopLevelContext(TopLevelArtifactContext options) {
-    PrecomputedValue.TOP_LEVEL_CONTEXT.set(injectable(), options);
-  }
-
   public void injectWorkspaceStatusData() {
     PrecomputedValue.WORKSPACE_STATUS_KEY.set(injectable(),
         workspaceStatusActionFactory.createWorkspaceStatusAction(
@@ -1083,7 +1076,8 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
       boolean finalizeActionsToOutputService,
       int numJobs,
       ActionCacheChecker actionCacheChecker,
-      @Nullable EvaluationProgressReceiver executionProgressReceiver)
+      @Nullable EvaluationProgressReceiver executionProgressReceiver,
+      TopLevelArtifactContext topLevelArtifactContext)
       throws InterruptedException {
     checkActive();
     Preconditions.checkState(actionLogBufferPathGenerator != null);
@@ -1096,9 +1090,11 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
     try {
       progressReceiver.executionProgressReceiver = executionProgressReceiver;
       Iterable<SkyKey> artifactKeys = ArtifactValue.mandatoryKeys(artifactsToBuild);
-      Iterable<SkyKey> targetKeys = TargetCompletionValue.keys(targetsToBuild);
-      Iterable<SkyKey> aspectKeys = AspectCompletionValue.keys(aspects);
-      Iterable<SkyKey> testKeys = TestCompletionValue.keys(targetsToTest, exclusiveTesting);
+      Iterable<SkyKey> targetKeys =
+          TargetCompletionValue.keys(targetsToBuild, topLevelArtifactContext);
+      Iterable<SkyKey> aspectKeys = AspectCompletionValue.keys(aspects, topLevelArtifactContext);
+      Iterable<SkyKey> testKeys =
+          TestCompletionValue.keys(targetsToTest, topLevelArtifactContext, exclusiveTesting);
       return buildDriver.evaluate(
           Iterables.concat(artifactKeys, targetKeys, aspectKeys, testKeys),
           keepGoing,
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TargetCompletionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/TargetCompletionValue.java
index fcdcdd17eb..ef55b9f2e6 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/TargetCompletionValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TargetCompletionValue.java
@@ -13,10 +13,12 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.auto.value.AutoValue;
 import com.google.common.base.Function;
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.LabelAndConfiguration;
+import com.google.devtools.build.lib.analysis.TopLevelArtifactContext;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 
@@ -36,18 +38,38 @@ public class TargetCompletionValue implements SkyValue {
     return ct;
   }
 
-  public static SkyKey key(LabelAndConfiguration labelAndConfiguration) {
-    return SkyKey.create(SkyFunctions.TARGET_COMPLETION, labelAndConfiguration);
+  public static SkyKey key(
+      LabelAndConfiguration labelAndConfiguration,
+      TopLevelArtifactContext topLevelArtifactContext) {
+    return SkyKey.create(
+        SkyFunctions.TARGET_COMPLETION,
+        TargetCompletionKey.create(labelAndConfiguration, topLevelArtifactContext));
   }
 
-  public static Iterable<SkyKey> keys(Collection<ConfiguredTarget> targets) {
+  public static Iterable<SkyKey> keys(Collection<ConfiguredTarget> targets,
+      final TopLevelArtifactContext ctx) {
     return Iterables.transform(
         targets,
         new Function<ConfiguredTarget, SkyKey>() {
           @Override
           public SkyKey apply(ConfiguredTarget ct) {
-            return SkyKey.create(SkyFunctions.TARGET_COMPLETION, LabelAndConfiguration.of(ct));
+            return SkyKey.create(
+                SkyFunctions.TARGET_COMPLETION,
+                TargetCompletionKey.create(LabelAndConfiguration.of(ct), ctx));
           }
         });
   }
+
+  @AutoValue
+  abstract static class TargetCompletionKey {
+    public static TargetCompletionKey create(
+        LabelAndConfiguration labelAndConfiguration,
+        TopLevelArtifactContext topLevelArtifactContext) {
+      return new AutoValue_TargetCompletionValue_TargetCompletionKey(
+          labelAndConfiguration, topLevelArtifactContext);
+    }
+
+    public abstract LabelAndConfiguration labelAndConfiguration();
+    public abstract TopLevelArtifactContext topLevelArtifactContext();
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionFunction.java
index 13a2d95c67..cdb22dd3e9 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionFunction.java
@@ -16,6 +16,7 @@ package com.google.devtools.build.lib.skyframe;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.LabelAndConfiguration;
+import com.google.devtools.build.lib.analysis.TopLevelArtifactContext;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.rules.test.TestProvider;
 import com.google.devtools.build.skyframe.SkyFunction;
@@ -36,8 +37,9 @@ public final class TestCompletionFunction implements SkyFunction {
   public SkyValue compute(SkyKey skyKey, Environment env) {
     TestCompletionValue.TestCompletionKey key =
         (TestCompletionValue.TestCompletionKey) skyKey.argument();
-    LabelAndConfiguration lac = key.getLabelAndConfiguration();
-    if (env.getValue(TargetCompletionValue.key(lac)) == null) {
+    LabelAndConfiguration lac = key.labelAndConfiguration();
+    TopLevelArtifactContext ctx = key.topLevelArtifactContext();
+    if (env.getValue(TargetCompletionValue.key(lac, ctx)) == null) {
       return null;
     }
 
@@ -48,7 +50,7 @@ public final class TestCompletionFunction implements SkyFunction {
     }
 
     ConfiguredTarget ct = ctValue.getConfiguredTarget();
-    if (key.isExclusiveTesting()) {
+    if (key.exclusiveTesting()) {
       // Request test artifacts iteratively if testing exclusively.
       for (Artifact testArtifact : TestProvider.getTestStatusArtifacts(ct)) {
         if (env.getValue(ArtifactValue.key(testArtifact, /*isMandatory=*/true)) == null) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionValue.java
index b5321a801f..01f5698eb0 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionValue.java
@@ -13,10 +13,12 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.auto.value.AutoValue;
 import com.google.common.base.Function;
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.LabelAndConfiguration;
+import com.google.devtools.build.lib.analysis.TopLevelArtifactContext;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 
@@ -31,12 +33,18 @@ public class TestCompletionValue implements SkyValue {
 
   private TestCompletionValue() { }
 
-  public static SkyKey key(LabelAndConfiguration lac, boolean exclusive) {
-    return SkyKey.create(SkyFunctions.TEST_COMPLETION, new TestCompletionKey(lac, exclusive));
+  public static SkyKey key(
+      LabelAndConfiguration lac,
+      final TopLevelArtifactContext topLevelArtifactContext,
+      final boolean exclusiveTesting) {
+    return SkyKey.create(
+        SkyFunctions.TEST_COMPLETION,
+        TestCompletionKey.create(lac, topLevelArtifactContext, exclusiveTesting));
   }
 
   public static Iterable<SkyKey> keys(Collection<ConfiguredTarget> targets,
-                                      final boolean exclusive) {
+                                      final TopLevelArtifactContext topLevelArtifactContext,
+                                      final boolean exclusiveTesting) {
     return Iterables.transform(
         targets,
         new Function<ConfiguredTarget, SkyKey>() {
@@ -44,26 +52,25 @@ public class TestCompletionValue implements SkyValue {
           public SkyKey apply(ConfiguredTarget ct) {
             return SkyKey.create(
                 SkyFunctions.TEST_COMPLETION,
-                new TestCompletionKey(LabelAndConfiguration.of(ct), exclusive));
+                TestCompletionKey.create(
+                    LabelAndConfiguration.of(ct), topLevelArtifactContext, exclusiveTesting));
           }
         });
   }
-  
-  static class TestCompletionKey {
-    private final LabelAndConfiguration lac;
-    private final boolean exclusiveTesting;
 
-    TestCompletionKey(LabelAndConfiguration lac, boolean exclusive) {
-      this.lac = lac;
-      this.exclusiveTesting = exclusive;
-    }
+  @AutoValue
+  abstract static class TestCompletionKey {
 
-    public LabelAndConfiguration getLabelAndConfiguration() {
-      return lac;
+    public static TestCompletionKey create(
+        LabelAndConfiguration labelAndConfiguration,
+        TopLevelArtifactContext topLevelArtifactContext,
+        boolean exclusiveTesting) {
+      return new AutoValue_TestCompletionValue_TestCompletionKey(
+          labelAndConfiguration, topLevelArtifactContext, exclusiveTesting);
     }
 
-    public boolean isExclusiveTesting() {
-      return exclusiveTesting;
-    }
+    public abstract LabelAndConfiguration labelAndConfiguration();
+    public abstract TopLevelArtifactContext topLevelArtifactContext();
+    public abstract boolean exclusiveTesting();
   }
 }