Request test artifacts to be built in parallel with running the test.

In cases where not all test artifacts are needed to run the test, this allows for greater parallelism in the build.

We need to inject the request for the test to be run into the TargetCompletion function, so that it can properly process any errors. Thanks to nharmata@ for suggesting this.

MEMORY: One additional Skyframe node and 2 edges for each test to be run (sharded tests still only count as one). So with 500 test targets, 100K is a conservative estimate of memory usage.
PiperOrigin-RevId: 166256545
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 8808252..16977df 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
@@ -32,11 +32,13 @@
  */
 public class ActionArtifactCycleReporter extends AbstractLabelCycleReporter {
   @SuppressWarnings("unchecked")
-  private static final Predicate<SkyKey> IS_ARTIFACT_OR_ACTION_SKY_KEY = Predicates.or(
-      SkyFunctions.isSkyFunction(SkyFunctions.ARTIFACT),
-      SkyFunctions.isSkyFunction(SkyFunctions.ACTION_EXECUTION),
-      SkyFunctions.isSkyFunction(SkyFunctions.TARGET_COMPLETION),
-      SkyFunctions.isSkyFunction(SkyFunctions.TEST_COMPLETION));
+  private static final Predicate<SkyKey> IS_ARTIFACT_OR_ACTION_SKY_KEY =
+      Predicates.or(
+          SkyFunctions.isSkyFunction(SkyFunctions.ARTIFACT),
+          SkyFunctions.isSkyFunction(SkyFunctions.ACTION_EXECUTION),
+          SkyFunctions.isSkyFunction(SkyFunctions.TARGET_COMPLETION),
+          SkyFunctions.isSkyFunction(SkyFunctions.TEST_COMPLETION),
+          SkyFunctions.isSkyFunction(SkyFunctions.TEST_EXECUTION));
 
   ActionArtifactCycleReporter(PackageProvider packageProvider) {
     super(packageProvider);
@@ -58,6 +60,10 @@
     } else if (arg instanceof TestCompletionKey
         && skyFunctionName.equals(SkyFunctions.TEST_COMPLETION)) {
       return  "test target: " + ((TestCompletionKey) arg).labelAndConfiguration().getLabel();
+    } else if (arg instanceof TestExecutionFunction.TestExecutionKey
+        && skyFunctionName.equals(SkyFunctions.TEST_EXECUTION)) {
+      return "test target: "
+          + ((TestExecutionFunction.TestExecutionKey) arg).getLabelAndConfiguration().getLabel();
     }
     throw new IllegalStateException(
         "Argument is not Action, TargetCompletion, TestCompletion or OwnedArtifact: " + arg);
@@ -76,6 +82,9 @@
     } else if (arg instanceof TestCompletionKey
         && key.functionName().equals(SkyFunctions.TEST_COMPLETION)) {
       return  ((TestCompletionKey) arg).labelAndConfiguration().getLabel();
+    } else if (arg instanceof TestExecutionFunction.TestExecutionKey
+        && key.functionName().equals(SkyFunctions.TEST_EXECUTION)) {
+      return ((TestExecutionFunction.TestExecutionKey) arg).getLabelAndConfiguration().getLabel();
     }
     throw new IllegalStateException(
         "Argument is not Action, TargetCompletion, TestCompletion or OwnedArtifact: " + arg);
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 8744f1a..5e10018 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
@@ -13,9 +13,10 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterables;
 import com.google.common.eventbus.EventBus;
 import com.google.devtools.build.lib.actions.ActionExecutionException;
-import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.MissingInputFileException;
 import com.google.devtools.build.lib.analysis.AspectCompleteEvent;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
@@ -55,6 +56,13 @@
     TValue getValueFromSkyKey(SkyKey skyKey, Environment env) throws InterruptedException;
 
     /**
+     * Gets any additional {@link SkyKey}s that should be requested in this {@link SkyFunction} for
+     * increased parallelism, together with the owner for each. The owner is only used in case of a
+     * {@link MissingInputFileException} thrown when trying to retrieve the key.
+     */
+    ImmutableMap<SkyKey, Label> getExtraSkyKeysWithOwnersToRequest(SkyKey skyKey);
+
+    /**
      * 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
@@ -107,6 +115,15 @@
     }
 
     @Override
+    public ImmutableMap<SkyKey, Label> getExtraSkyKeysWithOwnersToRequest(SkyKey skyKey) {
+      TargetCompletionKey tcKey = (TargetCompletionKey) skyKey.argument();
+      SkyKey testExecutionSkyKey = tcKey.testExecutionSkyKey();
+      return testExecutionSkyKey == null
+          ? ImmutableMap.of()
+          : ImmutableMap.of(testExecutionSkyKey, tcKey.labelAndConfiguration().getLabel());
+    }
+
+    @Override
     public TopLevelArtifactContext getTopLevelArtifactContext(SkyKey skyKey) {
       TargetCompletionKey tcKey = (TargetCompletionKey) skyKey.argument();
       return tcKey.topLevelArtifactContext();
@@ -165,6 +182,11 @@
     }
 
     @Override
+    public ImmutableMap<SkyKey, Label> getExtraSkyKeysWithOwnersToRequest(SkyKey skyKey) {
+      return ImmutableMap.of();
+    }
+
+    @Override
     public TopLevelArtifactContext getTopLevelArtifactContext(SkyKey skyKey) {
       AspectCompletionKey acKey = (AspectCompletionKey) skyKey.argument();
       return acKey.topLevelArtifactContext();
@@ -241,10 +263,14 @@
       return null;
     }
 
+    ImmutableMap<SkyKey, Label> extraSkyKeysToOwner =
+        completor.getExtraSkyKeysWithOwnersToRequest(skyKey);
     Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps =
         env.getValuesOrThrow(
-            ArtifactSkyKey.mandatoryKeys(
-                completor.getAllArtifactsToBuild(value, topLevelContext).getAllArtifacts()),
+            Iterables.concat(
+                ArtifactSkyKey.mandatoryKeys(
+                    completor.getAllArtifactsToBuild(value, topLevelContext).getAllArtifacts()),
+                extraSkyKeysToOwner.keySet()),
             MissingInputFileException.class,
             ActionExecutionException.class);
 
@@ -254,12 +280,14 @@
     NestedSetBuilder<Cause> rootCausesBuilder = NestedSetBuilder.stableOrder();
     for (Map.Entry<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>>
         depsEntry : inputDeps.entrySet()) {
-      Artifact input = ArtifactSkyKey.artifact(depsEntry.getKey());
       try {
         depsEntry.getValue().get();
       } catch (MissingInputFileException e) {
         missingCount++;
-        final Label inputOwner = input.getOwner();
+        Label inputOwner = extraSkyKeysToOwner.get(depsEntry.getKey());
+        if (inputOwner == null) {
+          inputOwner = ArtifactSkyKey.artifact(depsEntry.getKey()).getOwner();
+        }
         if (inputOwner != null) {
           Cause cause = new LabelCause(inputOwner);
           rootCausesBuilder.add(cause);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java
index 4f7479b..e7acc50 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java
@@ -78,6 +78,7 @@
       SkyFunctionName.create("TARGET_COMPLETION");
   public static final SkyFunctionName ASPECT_COMPLETION =
       SkyFunctionName.create("ASPECT_COMPLETION");
+  public static final SkyFunctionName TEST_EXECUTION = SkyFunctionName.create("TEST_EXECUTION");
   public static final SkyFunctionName TEST_COMPLETION = SkyFunctionName.create("TEST_COMPLETION");
   public static final SkyFunctionName BUILD_CONFIGURATION =
       SkyFunctionName.create("BUILD_CONFIGURATION");
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 f61ba9c..77bc25b 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
@@ -435,6 +435,7 @@
     map.put(SkyFunctions.TARGET_COMPLETION, CompletionFunction.targetCompletionFunction(eventBus));
     map.put(SkyFunctions.ASPECT_COMPLETION, CompletionFunction.aspectCompletionFunction(eventBus));
     map.put(SkyFunctions.TEST_COMPLETION, new TestCompletionFunction());
+    map.put(SkyFunctions.TEST_EXECUTION, new TestExecutionFunction());
     map.put(SkyFunctions.ARTIFACT, new ArtifactFunction(allowedMissingInputs));
     map.put(
         SkyFunctions.BUILD_INFO_COLLECTION,
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 476f69d..6c84563 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
@@ -14,15 +14,15 @@
 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.LegacySkyKey;
+import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import java.util.Collection;
+import javax.annotation.Nullable;
 
 /**
  * The value of a TargetCompletion. Currently this just stores a ConfiguredTarget.
@@ -40,36 +40,38 @@
 
   public static SkyKey key(
       LabelAndConfiguration labelAndConfiguration,
-      TopLevelArtifactContext topLevelArtifactContext) {
-    return LegacySkyKey.create(
-        SkyFunctions.TARGET_COMPLETION,
-        TargetCompletionKey.create(labelAndConfiguration, topLevelArtifactContext));
+      TopLevelArtifactContext topLevelArtifactContext,
+      SkyKey testExecutionSkyKey) {
+    return TargetCompletionKey.create(
+        labelAndConfiguration, topLevelArtifactContext, testExecutionSkyKey);
   }
 
   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 LegacySkyKey.create(
-                SkyFunctions.TARGET_COMPLETION,
-                TargetCompletionKey.create(LabelAndConfiguration.of(ct), ctx));
-          }
-        });
+        targets, ct -> TargetCompletionKey.create(LabelAndConfiguration.of(ct), ctx, null));
   }
 
   @AutoValue
-  abstract static class TargetCompletionKey {
+  abstract static class TargetCompletionKey implements SkyKey {
     public static TargetCompletionKey create(
         LabelAndConfiguration labelAndConfiguration,
-        TopLevelArtifactContext topLevelArtifactContext) {
+        TopLevelArtifactContext topLevelArtifactContext,
+        @Nullable SkyKey testExecutionSkyKey) {
       return new AutoValue_TargetCompletionValue_TargetCompletionKey(
-          labelAndConfiguration, topLevelArtifactContext);
+          labelAndConfiguration, topLevelArtifactContext, testExecutionSkyKey);
     }
 
-    public abstract LabelAndConfiguration labelAndConfiguration();
+    abstract LabelAndConfiguration labelAndConfiguration();
+
     public abstract TopLevelArtifactContext topLevelArtifactContext();
+
+    @Nullable
+    abstract SkyKey testExecutionSkyKey();
+
+    @Override
+    public SkyFunctionName functionName() {
+      return SkyFunctions.TARGET_COMPLETION;
+    }
   }
 }
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 dcdf47e..1dc46ce 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
@@ -13,11 +13,8 @@
 // limitations under the License.
 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.analysis.test.TestProvider;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyKey;
@@ -35,35 +32,20 @@
         (TestCompletionValue.TestCompletionKey) skyKey.argument();
     LabelAndConfiguration lac = key.labelAndConfiguration();
     TopLevelArtifactContext ctx = key.topLevelArtifactContext();
-    if (env.getValue(TargetCompletionValue.key(lac, ctx)) == null) {
+    env.getValue(
+        TargetCompletionValue.key(
+            lac, ctx, TestExecutionFunction.key(lac, key.exclusiveTesting())));
+    if (env.valuesMissing()) {
       return null;
     }
-
-    ConfiguredTargetValue ctValue = (ConfiguredTargetValue)
-        env.getValue(ConfiguredTargetValue.key(lac.getLabel(), lac.getConfiguration()));
-    if (ctValue == null) {
-      return null;
-    }
-
-    ConfiguredTarget ct = ctValue.getConfiguredTarget();
-    if (key.exclusiveTesting()) {
-      // Request test artifacts iteratively if testing exclusively.
-      for (Artifact testArtifact : TestProvider.getTestStatusArtifacts(ct)) {
-        if (env.getValue(ArtifactSkyKey.key(testArtifact, /*isMandatory=*/ true)) == null) {
-          return null;
-        }
-      }
-    } else {
-      env.getValues(ArtifactSkyKey.mandatoryKeys(TestProvider.getTestStatusArtifacts(ct)));
-      if (env.valuesMissing()) {
-        return null;
-      }
-    }
     return TestCompletionValue.TEST_COMPLETION_MARKER;
   }
 
   @Override
   public String extractTag(SkyKey skyKey) {
-    return Label.print(((LabelAndConfiguration) skyKey.argument()).getLabel());
+    return Label.print(
+        ((TestCompletionValue.TestCompletionKey) skyKey.argument())
+            .labelAndConfiguration()
+            .getLabel());
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TestExecutionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/TestExecutionFunction.java
new file mode 100644
index 0000000..a752c7d
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TestExecutionFunction.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.skyframe;
+
+import com.google.auto.value.AutoValue;
+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.test.TestProvider;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.skyframe.SkyFunction;
+import com.google.devtools.build.skyframe.SkyFunctionName;
+import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
+import javax.annotation.Nullable;
+
+/** {@link SkyFunction} to ensure that a test has executed. */
+class TestExecutionFunction implements SkyFunction {
+  private static final SkyValue TEST_EXECUTION_MARKER = new SkyValue() {};
+
+  static SkyKey key(LabelAndConfiguration lac, boolean exclusiveTesting) {
+    return TestExecutionKey.create(lac, exclusiveTesting);
+  }
+
+  @AutoValue
+  abstract static class TestExecutionKey implements SkyKey {
+    abstract LabelAndConfiguration getLabelAndConfiguration();
+
+    abstract boolean exclusiveTesting();
+
+    static TestExecutionKey create(LabelAndConfiguration lac, boolean exclusiveTesting) {
+      return new AutoValue_TestExecutionFunction_TestExecutionKey(lac, exclusiveTesting);
+    }
+
+    @Override
+    public SkyFunctionName functionName() {
+      return SkyFunctions.TEST_EXECUTION;
+    }
+  }
+
+  @Nullable
+  @Override
+  public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedException {
+    TestExecutionKey key = (TestExecutionKey) skyKey.argument();
+    LabelAndConfiguration lac = key.getLabelAndConfiguration();
+    ConfiguredTargetValue ctValue =
+        (ConfiguredTargetValue)
+            env.getValue(ConfiguredTargetValue.key(lac.getLabel(), lac.getConfiguration()));
+    if (ctValue == null) {
+      return null;
+    }
+
+    ConfiguredTarget ct = ctValue.getConfiguredTarget();
+    if (key.exclusiveTesting()) {
+      // Request test artifacts iteratively if testing exclusively.
+      for (Artifact testArtifact : TestProvider.getTestStatusArtifacts(ct)) {
+        if (env.getValue(ArtifactSkyKey.key(testArtifact, /*isMandatory=*/ true)) == null) {
+          return null;
+        }
+      }
+    } else {
+      env.getValues(ArtifactSkyKey.mandatoryKeys(TestProvider.getTestStatusArtifacts(ct)));
+      if (env.valuesMissing()) {
+        return null;
+      }
+    }
+    return TEST_EXECUTION_MARKER;
+  }
+
+  @Override
+  public String extractTag(SkyKey skyKey) {
+    return Label.print(
+        ((TestExecutionKey) skyKey.argument()).getLabelAndConfiguration().getLabel());
+  }
+}