Add a flag to force Bazel to download certain artifacts when using --remote_download_minimal (bazelbuild#15870)

When using --remote_download_minimal Bazel only downloads to the client the minimum number of artifacts necessary for the build to succeed. This can be sometimes problematic when local development (i.e. IDE) requires some artifacts to do local work (e.g. syntax completion). We add a flag --experimental_remote_download_regex that allow specifying a regular expression that will force certain artifacts to be forced to downloaded.

Tested locally that syntax errors disappear in IntelliJ when compiling Bazel itself with --remote_download_minimal and --experimental_remote_download_regex=".*jar$", also included an integration test with several cases.

Signed-off-by: Luis Fernando Pino Duque <[email protected]>

Closes bazelbuild#15638.

PiperOrigin-RevId: 460672954
Change-Id: I3a4f18d046d2d91603a276f16173ad2a253480dd

Co-authored-by: Luis Fernando Pino Duque <[email protected]>

Author: sgowroji

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.