Add ability for a Spawn to declare a required local output

Some native actions run both remotely and locally. For example, the CppCompileAction runs compilation remotely but does .d file pruning locally. This change adds the necessary infrastructure for an action to tell a Spawn which of its outputs it needs on the local filesytem.
bazelbuild · Jan 28, 2019 · e4502f1 · e4502f1
1 parent 2f82b33
commit e4502f1
Show file tree

Hide file tree

Showing 16 changed files with 87 additions and 1 deletion.
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionContext.java b/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionContext.java
@@ -32,6 +32,7 @@
 import com.google.devtools.common.options.OptionsProvider;
 import java.io.Closeable;
 import java.io.IOException;
+import java.util.Collection;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -72,6 +73,8 @@ public enum ShowSubcommands {
 
   private final ArtifactPathResolver pathResolver;
 
+  private final ImmutableList<Artifact> requiredLocalOutputs;
+
   private ActionExecutionContext(
       Executor executor,
       MetadataProvider actionInputFileCache,
@@ -82,6 +85,7 @@ private ActionExecutionContext(
       ExtendedEventHandler eventHandler,
       Map<String, String> clientEnv,
       ImmutableMap<Artifact, ImmutableList<FilesetOutputSymlink>> topLevelFilesets,
+      ImmutableList<Artifact> requiredLocalOutputs,
       @Nullable ArtifactExpander artifactExpander,
       @Nullable Environment env,
       @Nullable FileSystem actionFileSystem,
@@ -102,6 +106,7 @@ private ActionExecutionContext(
     this.pathResolver = ArtifactPathResolver.createPathResolver(actionFileSystem,
         // executor is only ever null in testing.
         executor == null ? null : executor.getExecRoot());
+    this.requiredLocalOutputs = Preconditions.checkNotNull(requiredLocalOutputs);
   }
 
   public ActionExecutionContext(
@@ -114,6 +119,7 @@ public ActionExecutionContext(
       ExtendedEventHandler eventHandler,
       Map<String, String> clientEnv,
       ImmutableMap<Artifact, ImmutableList<FilesetOutputSymlink>> topLevelFilesets,
+      ImmutableList<Artifact> requiredLocalOutputs,
       ArtifactExpander artifactExpander,
       @Nullable FileSystem actionFileSystem,
       @Nullable Object skyframeDepsResult) {
@@ -127,6 +133,7 @@ public ActionExecutionContext(
         eventHandler,
         clientEnv,
         topLevelFilesets,
+        requiredLocalOutputs,
         artifactExpander,
         /*env=*/ null,
         actionFileSystem,
@@ -154,6 +161,7 @@ public static ActionExecutionContext forInputDiscovery(
         eventHandler,
         clientEnv,
         ImmutableMap.of(),
+        /*requiredLocalOutputs=*/ ImmutableList.of(),
         /*artifactExpander=*/ null,
         env,
         actionFileSystem,
@@ -307,6 +315,18 @@ public ActionKeyContext getActionKeyContext() {
     return actionKeyContext;
   }
 
+  /**
+   * Returns the collection of files that this command must write and make available via
+   * the local {@link com.google.devtools.build.lib.vfs.FileSystem}. The returned output
+   * artifacts are a subset of the action's output artifacts.
+   *
+   * <p>This is for use with remote execution, where as an optimization we don't want to
+   * download all output files.
+   */
+  public Collection<Artifact> getRequiredLocalOutputs() {
+    return requiredLocalOutputs;
+  }
+
   @Override
   public void close() throws IOException {
     fileOutErr.close();
@@ -327,6 +347,29 @@ public ActionExecutionContext withFileOutErr(FileOutErr fileOutErr) {
         eventHandler,
         clientEnv,
         topLevelFilesets,
+        requiredLocalOutputs,
+        artifactExpander,
+        env,
+        actionFileSystem,
+        skyframeDepsResult);
+  }
+
+  /**
+   * Allows us to create a new context that provides a list of output artifacts that need to
+   * be staged on the local filesystem.
+   */
+  public ActionExecutionContext withRequiredLocalOutputs(Collection<? extends Artifact> requiredLocalOutputs) {
+    return new ActionExecutionContext(
+        executor,
+        actionInputFileCache,
+        actionInputPrefetcher,
+        actionKeyContext,
+        metadataHandler,
+        fileOutErr,
+        eventHandler,
+        clientEnv,
+        topLevelFilesets,
+        ImmutableList.copyOf(requiredLocalOutputs),
         artifactExpander,
         env,
         actionFileSystem,

diff --git a/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java b/src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java
@@ -21,6 +21,7 @@
 import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionStatusMessage;
+import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.ArtifactExpander;
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
@@ -43,6 +44,7 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.io.IOException;
 import java.time.Duration;
+import java.util.Collection;
 import java.util.List;
 import java.util.SortedMap;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -271,5 +273,10 @@ public void report(ProgressStatus state, String name) {
           break;
       }
     }
+
+    @Override
+    public Collection<Artifact> getRequiredLocalOutputs() {
+      return actionExecutionContext.getRequiredLocalOutputs();
+    }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/exec/SpawnRunner.java b/src/main/java/com/google/devtools/build/lib/exec/SpawnRunner.java
@@ -13,7 +13,9 @@
 // limitations under the License.
 package com.google.devtools.build.lib.exec;
 
+import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.ActionInput;
+import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.Artifact.ArtifactExpander;
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.ExecException;
@@ -25,6 +27,7 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.io.IOException;
 import java.time.Duration;
+import java.util.Collection;
 import java.util.SortedMap;
 
 /**
@@ -129,6 +132,7 @@ public enum ProgressStatus {
    * by different threads, so they MUST not call any shared non-thread-safe objects.
    */
   interface SpawnExecutionContext {
+
     /**
      * Returns a unique id for this spawn, to be used for logging. Note that a single spawn may be
      * passed to multiple {@link SpawnRunner} implementations, so any log entries should also
@@ -199,6 +203,18 @@ SortedMap<PathFragment, ActionInput> getInputMapping(boolean expandTreeArtifacts
 
     /** Reports a progress update to the Spawn strategy. */
     void report(ProgressStatus state, String name);
+
+    /**
+     * Returns the collection of files that this command must write and make available via
+     * the local {@link com.google.devtools.build.lib.vfs.FileSystem}. The returned output
+     * artifacts are a subset of the action's output artifacts.
+     *
+     * <p>This is for use with remote execution, where as an optimization we don't want to
+     * download all output files.
+     */
+    default Collection<Artifact> getRequiredLocalOutputs() {
+      return ImmutableList.of();
+    }
   }
 
   /**
@@ -234,4 +250,4 @@ SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
 
   /* Name of the SpawnRunner. */
   String getName();
-}
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
@@ -601,6 +601,7 @@ public ActionExecutionContext getContext(
             : progressSuppressingEventHandler,
         clientEnv,
         topLevelFilesets,
+        ImmutableList.of(),
         new ArtifactExpanderImpl(expandedInputs, expandedFilesets),
         actionFileSystem,
         skyframeDepsResult);

diff --git a/src/test/java/com/google/devtools/build/lib/actions/ExecutableSymlinkActionTest.java b/src/test/java/com/google/devtools/build/lib/actions/ExecutableSymlinkActionTest.java
@@ -18,6 +18,7 @@
 import static com.google.devtools.build.lib.actions.util.ActionsTestUtil.NULL_ACTION_OWNER;
 import static org.junit.Assert.fail;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.util.DummyExecutor;
 import com.google.devtools.build.lib.analysis.actions.SymlinkAction;
@@ -66,6 +67,7 @@ private ActionExecutionContext createContext() {
         executor.getEventHandler(),
         ImmutableMap.<String, String>of(),
         ImmutableMap.of(),
+        /*requiredLocalOutputs=*/ ImmutableList.of(),
         null,
         null,
         null);

diff --git a/src/test/java/com/google/devtools/build/lib/actions/util/ActionsTestUtil.java b/src/test/java/com/google/devtools/build/lib/actions/util/ActionsTestUtil.java
@@ -147,6 +147,7 @@ public static ActionExecutionContext createContext(
         executor != null ? executor.getEventHandler() : null,
         ImmutableMap.copyOf(clientEnv),
         ImmutableMap.of(),
+        ImmutableList.of(),
         actionGraph == null
             ? createDummyArtifactExpander()
             : ActionInputHelper.actionGraphArtifactExpander(actionGraph),
@@ -187,6 +188,7 @@ public static ActionExecutionContext createContext(ExtendedEventHandler eventHan
         eventHandler,
         ImmutableMap.of(),
         ImmutableMap.of(),
+        ImmutableList.of(),
         createDummyArtifactExpander(),
         /*actionFileSystem=*/ null,
         /*skyframeDepsResult=*/ null);

diff --git a/src/test/java/com/google/devtools/build/lib/analysis/actions/FileWriteActionTestCase.java b/src/test/java/com/google/devtools/build/lib/analysis/actions/FileWriteActionTestCase.java
@@ -16,6 +16,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.devtools.build.lib.actions.util.ActionsTestUtil.NULL_ACTION_OWNER;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
@@ -71,6 +72,7 @@ public final void createExecutorAndContext() throws Exception {
             executor.getEventHandler(),
             ImmutableMap.<String, String>of(),
             ImmutableMap.of(),
+            ImmutableList.of(),
             null,
             null,
             null);

diff --git a/src/test/java/com/google/devtools/build/lib/analysis/actions/ParamFileWriteActionTest.java b/src/test/java/com/google/devtools/build/lib/analysis/actions/ParamFileWriteActionTest.java
@@ -203,6 +203,7 @@ public void expand(Artifact artifact, Collection<? super Artifact> output) {
         executor.getEventHandler(),
         ImmutableMap.<String, String>of(),
         ImmutableMap.of(),
+        ImmutableList.of(),
         artifactExpander,
         /*actionFileSystem=*/ null,
         /*skyframeDepsResult=*/ null);

diff --git a/src/test/java/com/google/devtools/build/lib/analysis/actions/SymlinkActionTest.java b/src/test/java/com/google/devtools/build/lib/analysis/actions/SymlinkActionTest.java
@@ -16,6 +16,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.devtools.build.lib.actions.util.ActionsTestUtil.NULL_ACTION_OWNER;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionInputPrefetcher;
@@ -89,6 +90,7 @@ public void testSymlink() throws Exception {
                 executor.getEventHandler(),
                 ImmutableMap.<String, String>of(),
                 ImmutableMap.of(),
+                ImmutableList.of(),
                 null,
                 /*actionFileSystem=*/ null,
                 /*skyframeDepsResult=*/ null));

diff --git a/...test/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionActionTest.java b/...test/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionActionTest.java
@@ -192,6 +192,7 @@ private ActionExecutionContext createContext(Executor executor) {
         executor.getEventHandler(),
         ImmutableMap.of(),
         ImmutableMap.of(),
+        ImmutableList.of(),
         null,
         /*actionFileSystem=*/ null,
         /*skyframeDepsResult=*/ null);

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
@@ -2231,6 +2231,7 @@ public ActionExecutionContext build() {
           reporter,
           clientEnv,
           ImmutableMap.of(),
+          ImmutableList.of(),
           artifactExpander,
           /*actionFileSystem=*/ null,
           /*skyframeDepsResult*/ null);

diff --git a/src/test/java/com/google/devtools/build/lib/exec/StandaloneTestStrategyTest.java b/src/test/java/com/google/devtools/build/lib/exec/StandaloneTestStrategyTest.java
@@ -96,6 +96,7 @@ public FakeActionExecutionContext(
           /*eventHandler=*/ null,
           /* clientEnv= */ ImmutableMap.of(),
           /* topLevelFilesets= */ ImmutableMap.of(),
+          /* requiredLocalOutputs= */ ImmutableList.of(),
           /* artifactExpander= */ null,
           /* actionFileSystem= */ null,
           /* skyframeDepsResult= */ null);

diff --git a/src/test/java/com/google/devtools/build/lib/rules/cpp/CreateIncSymlinkActionTest.java b/src/test/java/com/google/devtools/build/lib/rules/cpp/CreateIncSymlinkActionTest.java
@@ -17,6 +17,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.devtools.build.lib.actions.util.ActionsTestUtil.NULL_ACTION_OWNER;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionKeyContext;
@@ -125,6 +126,7 @@ private ActionExecutionContext makeDummyContext() {
         executor.getEventHandler(),
         ImmutableMap.of(),
         ImmutableMap.of(),
+        ImmutableList.of(),
         null,
         null,
         null);

diff --git a/src/test/java/com/google/devtools/build/lib/rules/cpp/LtoBackendActionTest.java b/src/test/java/com/google/devtools/build/lib/rules/cpp/LtoBackendActionTest.java
@@ -15,6 +15,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.actions.AbstractAction;
@@ -90,6 +91,7 @@ public final void createExecutorAndContext() throws Exception {
             executor.getEventHandler(),
             ImmutableMap.<String, String>of(),
             ImmutableMap.of(),
+            ImmutableList.of(),
             null,
             /*actionFileSystem=*/ null,
             /*skyframeDepsResult=*/ null);

diff --git a/src/test/java/com/google/devtools/build/lib/rules/objc/BazelJ2ObjcLibraryTest.java b/src/test/java/com/google/devtools/build/lib/rules/objc/BazelJ2ObjcLibraryTest.java
@@ -794,6 +794,7 @@ public void testJ2ObjCCustomModuleMap() throws Exception {
             null,
             ImmutableMap.<String, String>of(),
             ImmutableMap.of(),
+            ImmutableList.of(),
             DUMMY_ARTIFACT_EXPANDER,
             null,
             null);
@@ -847,6 +848,7 @@ public void testModuleMapFromGenJarTreeArtifact() throws Exception {
             null,
             ImmutableMap.of(),
             ImmutableMap.of(),
+            ImmutableList.of(),
             DUMMY_ARTIFACT_EXPANDER,
             null,
             null);

diff --git a/src/test/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategyTest.java b/src/test/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategyTest.java
@@ -195,6 +195,7 @@ private ActionExecutionContext createContext() {
         executor.getEventHandler(),
         ImmutableMap.<String, String>of(),
         ImmutableMap.of(),
+        ImmutableList.of(),
         SIMPLE_ARTIFACT_EXPANDER,
         /*actionFileSystem=*/ null,
         /*skyframeDepsResult=*/ null);