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

7 files changed, 165 insertions, 55 deletions

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 8808252848..16977df4dd 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 @@ import com.google.devtools.build.skyframe.SkyKey;
  */
 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 @@ public class ActionArtifactCycleReporter extends AbstractLabelCycleReporter {
     } 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 @@ public class ActionArtifactCycleReporter extends AbstractLabelCycleReporter {
     } 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 8744f1a975..5e100181f8 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 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
     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 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
     }
 
     @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 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
     }
 
     @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 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
       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 @@ public final class CompletionFunction<TValue extends SkyValue, TResult extends S
     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 4f7479b5d6..e7acc502f4 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 @@ public final class SkyFunctions {
       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 f61ba9c54a..77bc25bfea 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 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
     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 476f69da44..6c84563619 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 class TargetCompletionValue implements SkyValue {
 
   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 dcdf47e11d..1dc46cedfe 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 @@ public final class TestCompletionFunction implements SkyFunction {
         (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 0000000000..a752c7dc9e
--- /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());
+  }
+}