Add a flag to force Bazel to download certain artifacts when using --…

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

@@ -145,6 +145,9 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.Phaser;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Predicate;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import javax.annotation.Nullable;
 
 /**
@@ -172,6 +175,8 @@ public class RemoteExecutionService {
 
   private final AtomicBoolean shutdown = new AtomicBoolean(false);
   private final AtomicBoolean buildInterrupted = new AtomicBoolean(false);
+  private final boolean shouldForceDownloads;
+  private final Predicate<String> shouldForceDownloadPredicate;
 
   public RemoteExecutionService(
       Executor executor,
@@ -213,6 +218,20 @@ public RemoteExecutionService(
     this.captureCorruptedOutputsDir = captureCorruptedOutputsDir;
 
     this.scheduler = Schedulers.from(executor, /*interruptibleWorker=*/ true);
+
+    String regex = remoteOptions.remoteDownloadRegex;
+    // TODO(bazel-team): Consider adding a warning or more validation if the remoteDownloadRegex is
+    // used without RemoteOutputsMode.MINIMAL.
+    this.shouldForceDownloads =
+        !regex.isEmpty()
+            && (remoteOptions.remoteOutputsMode == RemoteOutputsMode.MINIMAL
+                || remoteOptions.remoteOutputsMode == RemoteOutputsMode.TOPLEVEL);
+    Pattern pattern = Pattern.compile(regex);
+    this.shouldForceDownloadPredicate =
+        path -> {
+          Matcher m = pattern.matcher(path);
+          return m.matches();
+        };
   }
 
   static Command buildCommand(
@@ -1011,24 +1030,10 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
             /* exitCode = */ result.getExitCode(),
             hasFilesToDownload(action.getSpawn().getOutputFiles(), filesToDownload));
 
-    if (downloadOutputs) {
-      HashSet<PathFragment> queuedFilePaths = new HashSet<>();
-
-      for (FileMetadata file : metadata.files()) {
-        PathFragment filePath = file.path().asFragment();
-        if (queuedFilePaths.add(filePath)) {
-          downloadsBuilder.add(downloadFile(action, file));
-        }
-      }
+    ImmutableList<ListenableFuture<FileMetadata>> forcedDownloads = ImmutableList.of();
 
-      for (Map.Entry<Path, DirectoryMetadata> entry : metadata.directories()) {
-        for (FileMetadata file : entry.getValue().files()) {
-          PathFragment filePath = file.path().asFragment();
-          if (queuedFilePaths.add(filePath)) {
-            downloadsBuilder.add(downloadFile(action, file));
-          }
-        }
-      }
+    if (downloadOutputs) {
+      downloadsBuilder.addAll(buildFilesToDownload(metadata, action));
     } else {
       checkState(
           result.getExitCode() == 0,
@@ -1039,6 +1044,10 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
             "Symlinks in action outputs are not yet supported by "
                 + "--experimental_remote_download_outputs=minimal");
       }
+      if (shouldForceDownloads) {
+        forcedDownloads =
+            buildFilesToDownloadWithPredicate(metadata, action, shouldForceDownloadPredicate);
+      }
     }
 
     FileOutErr tmpOutErr = outErr.childOutErr();
@@ -1053,6 +1062,19 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
     try (SilentCloseable c = Profiler.instance().profile("Remote.download")) {
       waitForBulkTransfer(downloads, /* cancelRemainingOnInterrupt= */ true);
     } catch (Exception e) {
+      // TODO(bazel-team): Consider adding better case-by-case exception handling instead of just
+      // rethrowing
+      captureCorruptedOutputs(e);
+      deletePartialDownloadedOutputs(result, tmpOutErr, e);
+      throw e;
+    }
+
+    // TODO(bazel-team): Unify this block with the equivalent block above.
+    try (SilentCloseable c = Profiler.instance().profile("Remote.forcedDownload")) {
+      waitForBulkTransfer(forcedDownloads, /* cancelRemainingOnInterrupt= */ true);
+    } catch (Exception e) {
+      // TODO(bazel-team): Consider adding better case-by-case exception handling instead of just
+      // rethrowing
       captureCorruptedOutputs(e);
       deletePartialDownloadedOutputs(result, tmpOutErr, e);
       throw e;
@@ -1083,6 +1105,13 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
       // might not be supported on all platforms
       createSymlinks(symlinks);
     } else {
+      // TODO(bazel-team): We should unify this if-block to rely on downloadOutputs above but, as of
+      // 2022-07-05,  downloadOuputs' semantics isn't exactly the same as build-without-the-bytes
+      // which is necessary for using remoteDownloadRegex.
+      if (!forcedDownloads.isEmpty()) {
+        moveOutputsToFinalLocation(forcedDownloads);
+      }
+
       ActionInput inMemoryOutput = null;
       Digest inMemoryOutputDigest = null;
       PathFragment inMemoryOutputPath = getInMemoryOutputPath(action.getSpawn());
@@ -1120,6 +1149,36 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
     return null;
   }
 
+  private ImmutableList<ListenableFuture<FileMetadata>> buildFilesToDownload(
+      ActionResultMetadata metadata, RemoteAction action) {
+    Predicate<String> alwaysTrue = unused -> true;
+    return buildFilesToDownloadWithPredicate(metadata, action, alwaysTrue);
+  }
+
+  private ImmutableList<ListenableFuture<FileMetadata>> buildFilesToDownloadWithPredicate(
+      ActionResultMetadata metadata, RemoteAction action, Predicate<String> predicate) {
+    HashSet<PathFragment> queuedFilePaths = new HashSet<>();
+    ImmutableList.Builder<ListenableFuture<FileMetadata>> builder = new ImmutableList.Builder<>();
+
+    for (FileMetadata file : metadata.files()) {
+      PathFragment filePath = file.path().asFragment();
+      if (queuedFilePaths.add(filePath) && predicate.test(file.path.toString())) {
+        builder.add(downloadFile(action, file));
+      }
+    }
+
+    for (Map.Entry<Path, DirectoryMetadata> entry : metadata.directories()) {
+      for (FileMetadata file : entry.getValue().files()) {
+        PathFragment filePath = file.path().asFragment();
+        if (queuedFilePaths.add(filePath) && predicate.test(file.path.toString())) {
+          builder.add(downloadFile(action, file));
+        }
+      }
+    }
+
+    return builder.build();
+  }
+
   private static String prettyPrint(ActionInput actionInput) {
     if (actionInput instanceof Artifact) {
       return ((Artifact) actionInput).prettyPrint();

src/main/java/com/google/devtools/build/lib/remote/options/RemoteOptions.java

@@ -579,6 +579,18 @@ public RemoteOutputsStrategyConverter() {
               + "`all` to print always.")
   public ExecutionMessagePrintMode remotePrintExecutionMessages;
 
+  @Option(
+      name = "experimental_remote_download_regex",
+      defaultValue = "",
+      documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
+      effectTags = {OptionEffectTag.AFFECTS_OUTPUTS},
+      help =
+          "Force Bazel to download the artifacts that match the given regexp. To be used in"
+              + " conjunction with --remote_download_minimal or --remote_download_toplevel to allow"
+              + " the client to request certain artifacts that might be needed locally (e.g. IDE"
+              + " support)")
+  public String remoteDownloadRegex;
+
   // The below options are not configurable by users, only tests.
   // This is part of the effort to reduce the overall number of flags.