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

1 files changed, 33 insertions, 5 deletions

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);