Use ToplevelArtifactsDownloader to download toplevel artifacts

src/main/java/com/google/devtools/build/lib/remote/BUILD

@@ -46,6 +46,7 @@ java_library(
         ":ReferenceCountedChannel",
         ":Retrier",
         ":abstract_action_input_prefetcher",
+        ":toplevel_artifacts_downloader",
         "//src/main/java/com/google/devtools/build/lib:build-request-options",
         "//src/main/java/com/google/devtools/build/lib:runtime",
         "//src/main/java/com/google/devtools/build/lib:runtime/command_line_path_factory",

src/main/java/com/google/devtools/build/lib/remote/RemoteActionContextProvider.java

@@ -47,7 +47,6 @@ final class RemoteActionContextProvider {
   @Nullable private final ListeningScheduledExecutorService retryScheduler;
   private final DigestUtil digestUtil;
   @Nullable private final Path logDir;
-  private ImmutableSet<ActionInput> filesToDownload = ImmutableSet.of();
   private TempPathGenerator tempPathGenerator;
   private RemoteExecutionService remoteExecutionService;
   @Nullable private RemoteActionInputFetcher actionInputFetcher;
@@ -157,7 +156,6 @@ private RemoteExecutionService getRemoteExecutionService() {
               checkNotNull(env.getOptions().getOptions(RemoteOptions.class)),
               remoteCache,
               remoteExecutor,
-              filesToDownload,
               tempPathGenerator,
               captureCorruptedOutputsDir);
       env.getEventBus().register(remoteExecutionService);
@@ -213,15 +211,16 @@ RemoteExecutionClient getRemoteExecutionClient() {
     return remoteExecutor;
   }
 
-  void setFilesToDownload(ImmutableSet<ActionInput> topLevelOutputs) {
-    this.filesToDownload = Preconditions.checkNotNull(topLevelOutputs, "filesToDownload");
-  }
-
   void setTempPathGenerator(TempPathGenerator tempPathGenerator) {
     this.tempPathGenerator = tempPathGenerator;
   }
 
   public void afterCommand() {
+    // actionInputFetcher uses remoteCache to prefetch inputs, so must shut it down before
+    // remoteCache.
+    if (actionInputFetcher != null) {
+      actionInputFetcher.shutdown();
+    }
     if (remoteExecutionService != null) {
       remoteExecutionService.shutdown();
     } else {
@@ -232,9 +231,5 @@ public void afterCommand() {
         remoteExecutor.close();
       }
     }
-
-    if (actionInputFetcher != null) {
-      actionInputFetcher.shutdown();
-    }
   }
 }

src/main/java/com/google/devtools/build/lib/remote/RemoteActionFileSystem.java

@@ -508,7 +508,7 @@ public long getNodeId() {
   }
 
   @Nullable
-  private RemoteFileArtifactValue getRemoteMetadata(PathFragment path) {
+  protected RemoteFileArtifactValue getRemoteMetadata(PathFragment path) {
     if (!isOutput(path)) {
       return null;
     }

src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java

@@ -25,7 +25,6 @@
 import static com.google.devtools.build.lib.remote.util.Utils.getFromFuture;
 import static com.google.devtools.build.lib.remote.util.Utils.getInMemoryOutputPath;
 import static com.google.devtools.build.lib.remote.util.Utils.grpcAwareErrorMessage;
-import static com.google.devtools.build.lib.remote.util.Utils.hasFilesToDownload;
 import static com.google.devtools.build.lib.remote.util.Utils.shouldDownloadAllSpawnOutputs;
 import static com.google.devtools.build.lib.remote.util.Utils.shouldUploadLocalResultsToRemoteCache;
 import static com.google.devtools.build.lib.remote.util.Utils.waitForBulkTransfer;
@@ -163,7 +162,6 @@ public class RemoteExecutionService {
   private final RemoteOptions remoteOptions;
   @Nullable private final RemoteCache remoteCache;
   @Nullable private final RemoteExecutionClient remoteExecutor;
-  private final ImmutableSet<PathFragment> filesToDownload;
   private final TempPathGenerator tempPathGenerator;
   @Nullable private final Path captureCorruptedOutputsDir;
   private final Cache<Object, MerkleTree> merkleTreeCache;
@@ -187,7 +185,6 @@ public RemoteExecutionService(
       RemoteOptions remoteOptions,
       @Nullable RemoteCache remoteCache,
       @Nullable RemoteExecutionClient remoteExecutor,
-      ImmutableSet<ActionInput> filesToDownload,
       TempPathGenerator tempPathGenerator,
       @Nullable Path captureCorruptedOutputsDir) {
     this.reporter = reporter;
@@ -208,11 +205,6 @@ public RemoteExecutionService(
     }
     this.merkleTreeCache = merkleTreeCacheBuilder.build();
 
-    ImmutableSet.Builder<PathFragment> filesToDownloadBuilder = ImmutableSet.builder();
-    for (ActionInput actionInput : filesToDownload) {
-      filesToDownloadBuilder.add(actionInput.getExecPath());
-    }
-    this.filesToDownload = filesToDownloadBuilder.build();
     this.tempPathGenerator = tempPathGenerator;
     this.captureCorruptedOutputsDir = captureCorruptedOutputsDir;
 
@@ -1018,11 +1010,8 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
     ImmutableList.Builder<ListenableFuture<FileMetadata>> downloadsBuilder =
         ImmutableList.builder();
     RemoteOutputsMode remoteOutputsMode = remoteOptions.remoteOutputsMode;
-    boolean downloadOutputs =
-        shouldDownloadAllSpawnOutputs(
-            remoteOutputsMode,
-            /* exitCode = */ result.getExitCode(),
-            hasFilesToDownload(action.getSpawn().getOutputFiles(), filesToDownload));
+    boolean downloadOutputs = shouldDownloadAllSpawnOutputs(remoteOutputsMode,
+        result.getExitCode());
 
     // Download into temporary paths, then move everything at the end.
     // This avoids holding the output lock while downloading, which would prevent the local branch

src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java

@@ -27,26 +27,17 @@
 import com.google.common.base.Strings;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.GoogleLogger;
 import com.google.common.util.concurrent.ListeningScheduledExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
-import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
-import com.google.devtools.build.lib.actions.ActionGraph;
-import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.AnalysisResult;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
-import com.google.devtools.build.lib.analysis.FilesToRunProvider;
-import com.google.devtools.build.lib.analysis.RunfilesSupport;
-import com.google.devtools.build.lib.analysis.TopLevelArtifactContext;
-import com.google.devtools.build.lib.analysis.TopLevelArtifactHelper;
 import com.google.devtools.build.lib.analysis.config.BuildOptions;
 import com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget;
-import com.google.devtools.build.lib.analysis.test.TestProvider;
 import com.google.devtools.build.lib.authandtls.AuthAndTLSOptions;
 import com.google.devtools.build.lib.authandtls.CallCredentialsProvider;
 import com.google.devtools.build.lib.authandtls.GoogleAuthUtils;
@@ -57,14 +48,12 @@
 import com.google.devtools.build.lib.buildeventstream.LocalFilesArtifactUploader;
 import com.google.devtools.build.lib.buildtool.BuildRequest;
 import com.google.devtools.build.lib.buildtool.BuildRequestOptions;
-import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.Reporter;
 import com.google.devtools.build.lib.exec.ExecutionOptions;
 import com.google.devtools.build.lib.exec.ExecutorBuilder;
 import com.google.devtools.build.lib.exec.ModuleActionContextRegistry;
 import com.google.devtools.build.lib.exec.SpawnStrategyRegistry;
-import com.google.devtools.build.lib.packages.TargetUtils;
 import com.google.devtools.build.lib.remote.RemoteServerCapabilities.ServerCapabilitiesRequirement;
 import com.google.devtools.build.lib.remote.common.RemoteCacheClient;
 import com.google.devtools.build.lib.remote.common.RemoteExecutionClient;
@@ -111,10 +100,8 @@
 import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
-import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -140,6 +127,7 @@ public final class RemoteModule extends BlazeModule {
   @Nullable private ExecutorService executorService;
   @Nullable private RemoteActionContextProvider actionContextProvider;
   @Nullable private RemoteActionInputFetcher actionInputFetcher;
+  @Nullable private ToplevelArtifactsDownloader toplevelArtifactsDownloader;
   @Nullable private RemoteOptions remoteOptions;
   @Nullable private RemoteOutputService remoteOutputService;
   @Nullable private TempPathGenerator tempPathGenerator;
@@ -715,77 +703,12 @@ private static void handleInitFailure(
                 remoteExecutionCode));
   }
 
-  private static ImmutableList<Artifact> getRunfiles(ConfiguredTarget buildTarget) {
-    FilesToRunProvider runfilesProvider = buildTarget.getProvider(FilesToRunProvider.class);
-    if (runfilesProvider == null) {
-      return ImmutableList.of();
-    }
-    RunfilesSupport runfilesSupport = runfilesProvider.getRunfilesSupport();
-    if (runfilesSupport == null) {
-      return ImmutableList.of();
-    }
-    ImmutableList.Builder<Artifact> runfilesBuilder = ImmutableList.builder();
-    for (Artifact runfile : runfilesSupport.getRunfiles().getArtifacts().toList()) {
-      if (runfile.isSourceArtifact()) {
-        continue;
-      }
-      runfilesBuilder.add(runfile);
-    }
-    return runfilesBuilder.build();
-  }
-
-  private static ImmutableList<ActionInput> getTestOutputs(ConfiguredTarget testTarget) {
-    TestProvider testProvider = testTarget.getProvider(TestProvider.class);
-    if (testProvider == null) {
-      return ImmutableList.of();
-    }
-    return testProvider.getTestParams().getOutputs();
-  }
-
-  private static NestedSet<Artifact> getArtifactsToBuild(
-      ConfiguredTarget buildTarget, TopLevelArtifactContext topLevelArtifactContext) {
-    return TopLevelArtifactHelper.getAllArtifactsToBuild(buildTarget, topLevelArtifactContext)
-        .getImportantArtifacts();
-  }
-
-  private static boolean isTestRule(ConfiguredTarget configuredTarget) {
-    if (configuredTarget instanceof RuleConfiguredTarget) {
-      RuleConfiguredTarget ruleConfiguredTarget = (RuleConfiguredTarget) configuredTarget;
-      return TargetUtils.isTestRuleName(ruleConfiguredTarget.getRuleClassString());
-    }
-    return false;
-  }
-
   @Override
   public void afterAnalysis(
       CommandEnvironment env,
       BuildRequest request,
       BuildOptions buildOptions,
       AnalysisResult analysisResult) {
-    // The actionContextProvider may be null if remote execution is disabled or if there was an
-    // error during initialization.
-    if (remoteOptions != null
-        && remoteOptions.remoteOutputsMode.downloadToplevelOutputsOnly()
-        && actionContextProvider != null) {
-      boolean isTestCommand = env.getCommandName().equals("test");
-      TopLevelArtifactContext artifactContext = request.getTopLevelArtifactContext();
-      Set<ActionInput> filesToDownload = new HashSet<>();
-      for (ConfiguredTarget configuredTarget : analysisResult.getTargetsToBuild()) {
-        if (isTestCommand && isTestRule(configuredTarget)) {
-          // When running a test download the test.log and test.xml. These are never symlinks.
-          filesToDownload.addAll(getTestOutputs(configuredTarget));
-        } else {
-          fetchSymlinkDependenciesRecursively(
-              analysisResult.getActionGraph(),
-              filesToDownload,
-              getArtifactsToBuild(configuredTarget, artifactContext).toList());
-          fetchSymlinkDependenciesRecursively(
-              analysisResult.getActionGraph(), filesToDownload, getRunfiles(configuredTarget));
-        }
-      }
-      actionContextProvider.setFilesToDownload(ImmutableSet.copyOf(filesToDownload));
-    }
-
     if (remoteOptions != null
         && remoteOptions.remoteBuildEventUploadMode == RemoteBuildEventUploadMode.ALL
         && remoteOptions.incompatibleRemoteBuildEventUploadRespectNoCache) {
@@ -826,44 +749,6 @@ private void parseNoCacheOutputs(AnalysisResult analysisResult) {
     }
   }
 
-  // This is a short-term fix for top-level outputs that are symlinks. Unfortunately, we cannot
-  // reliably tell after analysis whether actions will create symlinks (the RE protocol allows any
-  // action to generate and return symlinks), but at least we can handle basic C++ rules with this
-  // change.
-  // TODO(ulfjack): I think we should separate downloading files from action execution. That would
-  // also resolve issues around action invalidation - we currently invalidate actions to trigger
-  // downloads of top-level outputs when the top-level targets change.
-  private static void fetchSymlinkDependenciesRecursively(
-      ActionGraph actionGraph, Set<ActionInput> builder, List<Artifact> inputs) {
-    for (Artifact input : inputs) {
-      // Only fetch recursively if we don't have the file to avoid visiting nodes multiple times.
-      if (builder.add(input)) {
-        fetchSymlinkDependenciesRecursively(actionGraph, builder, input);
-      }
-    }
-  }
-
-  private static void fetchSymlinkDependenciesRecursively(
-      ActionGraph actionGraph, Set<ActionInput> builder, Artifact artifact) {
-    if (!(actionGraph.getGeneratingAction(artifact) instanceof ActionExecutionMetadata)) {
-      // The top-level artifact could be a tree artifact, in which case the generating action may
-      // be an ActionTemplate, which does not implement ActionExecutionMetadata. We don't handle
-      // this case right now, so exit.
-      return;
-    }
-    ActionExecutionMetadata action =
-        (ActionExecutionMetadata) actionGraph.getGeneratingAction(artifact);
-    if (action.mayInsensitivelyPropagateInputs()) {
-      List<Artifact> inputs = action.getInputs().toList();
-      if (inputs.size() > 5) {
-        logger.atWarning().log(
-            "Action with a lot of inputs insensitively propagates them; this could be performance"
-                + " problem");
-      }
-      fetchSymlinkDependenciesRecursively(actionGraph, builder, inputs);
-    }
-  }
-
   private static void cleanAndCreateRemoteLogsDir(Path logDir) throws AbruptExitException {
     try {
       // Clean out old logs files.
@@ -921,6 +806,7 @@ public void afterCommand() throws AbruptExitException {
     remoteDownloaderSupplier.set(null);
     actionContextProvider = null;
     actionInputFetcher = null;
+    toplevelArtifactsDownloader = null;
     remoteOptions = null;
     remoteOutputService = null;
     tempPathGenerator = null;
@@ -1042,6 +928,18 @@ public void executorInit(CommandEnvironment env, BuildRequest request, ExecutorB
       builder.setActionInputPrefetcher(actionInputFetcher);
       remoteOutputService.setActionInputFetcher(actionInputFetcher);
       actionContextProvider.setActionInputFetcher(actionInputFetcher);
+
+      if (remoteOutputsMode.downloadToplevelOutputsOnly()) {
+        toplevelArtifactsDownloader = new ToplevelArtifactsDownloader(env.getSkyframeExecutor()
+            .getEvaluator(), actionInputFetcher, (path) -> {
+          FileSystem fileSystem = path.getFileSystem();
+          Preconditions.checkState(fileSystem instanceof RemoteActionFileSystem,
+              "fileSystem must be an instance of RemoteActionFileSystem");
+          return ((RemoteActionFileSystem) path.getFileSystem()).getRemoteMetadata(
+              path.asFragment());
+        });
+        env.getEventBus().register(toplevelArtifactsDownloader);
+      }
     }
   }
 

src/main/java/com/google/devtools/build/lib/remote/util/Utils.java

@@ -190,35 +190,16 @@ private static Instant timestampToInstant(Timestamp timestamp) {
     return Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos());
   }
 
-  /** Returns {@code true} if all spawn outputs should be downloaded to disk. */
-  public static boolean shouldDownloadAllSpawnOutputs(
-      RemoteOutputsMode remoteOutputsMode, int exitCode, boolean hasTopLevelOutputs) {
+  /**
+   * Returns {@code true} if all spawn outputs should be downloaded to disk.
+   */
+  public static boolean shouldDownloadAllSpawnOutputs(RemoteOutputsMode remoteOutputsMode,
+      int exitCode) {
     return remoteOutputsMode.downloadAllOutputs()
         ||
         // In case the action failed, download all outputs. It might be helpful for debugging
         // and there is no point in injecting output metadata of a failed action.
-        exitCode != 0
-        ||
-        // If one output of a spawn is a top level output then download all outputs. Spawns
-        // are typically structured in a way that either all or no outputs are top level and
-        // it's much simpler to implement under this assumption.
-        (remoteOutputsMode.downloadToplevelOutputsOnly() && hasTopLevelOutputs);
-  }
-
-  /** Returns {@code true} if outputs contains one or more top level outputs. */
-  public static boolean hasFilesToDownload(
-      Collection<? extends ActionInput> outputs, ImmutableSet<PathFragment> filesToDownload) {
-    if (filesToDownload.isEmpty()) {
-      return false;
-    }
-
-    for (ActionInput output : outputs) {
-      if (filesToDownload.contains(output.getExecPath())) {
-        return true;
-      }
-    }
-
-    return false;
+        exitCode != 0;
   }
 
   private static String statusName(int code) {