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

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

@@ -147,6 +147,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;
 
 /**
@@ -174,6 +177,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,
@@ -215,6 +220,16 @@ public RemoteExecutionService(
     this.captureCorruptedOutputsDir = captureCorruptedOutputsDir;
 
     this.scheduler = Schedulers.from(executor, /*interruptibleWorker=*/ true);
+
+    String regex = remoteOptions.experimentalForceDownloadsRegex;
+    // TODO(bazel-team): Consider adding a warning or more validation if the experimentalForceDownloadsRegex is used
+    // without RemoteOutputsMode.MINIMAL
+    this.shouldForceDownloads = !regex.isEmpty() && remoteOptions.remoteOutputsMode == RemoteOutputsMode.MINIMAL;
+    Pattern pattern = Pattern.compile(regex);
+    this.shouldForceDownloadPredicate = path -> {
+      Matcher m = pattern.matcher(path);
+      return m.matches();
+    };
   }
 
   static Command buildCommand(
@@ -1021,24 +1036,9 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
             /* exitCode = */ result.getExitCode(),
             hasFilesToDownload(action.getSpawn().getOutputFiles(), filesToDownload));
 
+    ImmutableList.Builder<ListenableFuture<FileMetadata>> forcedDownloadsBuilder = ImmutableList.builder();
     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));
-        }
-      }
-
-      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));
-          }
-        }
-      }
+      downloadsBuilder.addAll(buildFilesToDownload(metadata, action));
     } else {
       checkState(
           result.getExitCode() == 0,
@@ -1049,6 +1049,9 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
             "Symlinks in action outputs are not yet supported by "
                 + "--experimental_remote_download_outputs=minimal");
       }
+      if (shouldForceDownloads) {
+        forcedDownloadsBuilder.addAll(buildFilesToDownloadWithPredicate(metadata, action, shouldForceDownloadPredicate));
+      }
     }
 
     FileOutErr tmpOutErr = outErr.childOutErr();
@@ -1068,6 +1071,15 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
       throw e;
     }
 
+    ImmutableList<ListenableFuture<FileMetadata>> forcedDownloads = forcedDownloadsBuilder.build();
+    try (SilentCloseable c = Profiler.instance().profile("Remote.forcedDownload")) {
+      waitForBulkTransfer(forcedDownloads, /* cancelRemainingOnInterrupt= */ true);
+    } catch (Exception e) {
+      captureCorruptedOutputs(e);
+      deletePartialDownloadedOutputs(result, tmpOutErr, e);
+      throw e;
+    }
+
     FileOutErr.dump(tmpOutErr, outErr);
 
     // Ensure that we are the only ones writing to the output files when using the dynamic spawn
@@ -1079,6 +1091,10 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
     tmpOutErr.clearOut();
     tmpOutErr.clearErr();
 
+    if (!forcedDownloads.isEmpty()) {
+      moveOutputsToFinalLocation(forcedDownloads);
+    }
+
     if (downloadOutputs) {
       moveOutputsToFinalLocation(downloads);
 
@@ -1133,6 +1149,35 @@ 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

@@ -574,6 +574,16 @@ public RemoteOutputsStrategyConverter() {
       help = "Maximum number of open files allowed during BEP artifact upload.")
   public int maximumOpenFiles;
 
+  @Option(
+          name = "experimental_force_downloads_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 to allow the client to request certain artifacts that might be needed" +
+                  "locally (e.g. IDE support)")
+  public String experimentalForceDownloadsRegex;
+
   // The below options are not configurable by users, only tests.
   // This is part of the effort to reduce the overall number of flags.
 

src/test/shell/bazel/remote/remote_execution_test.sh

@@ -3347,6 +3347,63 @@ EOF
   expect_log "2 processes: 1 internal, 1 remote"
 }
 
+function test_forced_downloads() {
+  mkdir -p a
+
+  cat > a/BUILD <<'EOF'
+java_library(
+    name = "lib",
+    srcs = ["Library.java"],
+)
+java_test(
+    name = "test",
+    srcs = ["JavaTest.java"],
+    test_class = "JavaTest",
+    deps = [":lib"],
+)
+EOF
+
+  cat > a/Library.java <<'EOF'
+public class Library {
+  public static boolean TEST = true;
+}
+EOF
+
+  cat > a/JavaTest.java <<'EOF'
+import org.junit.Assert;
+import org.junit.Test;
+public class JavaTest {
+    @Test
+    public void test() { Assert.assertTrue(Library.TEST); }
+}
+EOF
+  bazel test \
+      --remote_executor=grpc://localhost:${worker_port} \
+      --remote_download_minimal \
+      //a:test >& $TEST_log || fail "Failed to build"
+
+  [[ ! -e "bazel-bin/a/liblib.jar" ]] || fail "bazel-bin/a/liblib.jar shouldn't exist"
+  [[ ! -e "bazel-bin/a/liblib.jdeps" ]] || fail "bazel-bin/a/liblib.jdeps shouldn't exist"
+
+  bazel clean && bazel test \
+        --remote_executor=grpc://localhost:${worker_port} \
+        --remote_download_minimal \
+        --experimental_force_downloads_regex=".*" \
+        //a:test >& $TEST_log || fail "Failed to build"
+
+  [[ -e "bazel-bin/a/liblib.jar" ]] || fail "bazel-bin/a/liblib.jar file does not exist!"
+  [[ -e "bazel-bin/a/liblib.jdeps" ]] || fail "bazel-bin/a/liblib.jdeps file does not exist!"
+
+  bazel clean && bazel test \
+    --remote_executor=grpc://localhost:${worker_port} \
+    --remote_download_minimal \
+    --experimental_force_downloads_regex=".*jar$" \
+    //a:test >& $TEST_log || fail "Failed to build"
+
+  [[ -e "bazel-bin/a/liblib.jar" ]] || fail "bazel-bin/a/liblib.jar file does not exist!"
+  [[ ! -e "bazel-bin/a/liblib.jdeps" ]] || fail "bazel-bin/a/liblib.jdeps shouldn't exist"
+}
+
 function test_grpc_connection_errors_are_propagated() {
   # Test that errors when creating grpc connection are propagated instead of crashing Bazel.
   # https://github.com/bazelbuild/bazel/issues/13724