From 86dee6d2ecb269e0c41a97718812054394ee51a4 Mon Sep 17 00:00:00 2001
From: Chi Wang <chiwang@google.com>
Date: Thu, 8 Dec 2022 02:33:08 -0800
Subject: [PATCH] Correctly match regex with tree artifact

It doesn't work when setting `--experimental_remote_download_regex` to match files inside tree artifact after https://github.com/bazelbuild/bazel/commit/e01e7f51dd19f39ce3bc0718cec20ed6474de733. This change fixes that.

Fixes #16922.

Closes #16949.

PiperOrigin-RevId: 493838296
Change-Id: I6eceffbffce30949173d10120a9120c6c608a983
---
 .../remote/AbstractActionInputPrefetcher.java | 21 ++++-
 ...ildWithoutTheBytesIntegrationTestBase.java | 81 +++++++++++++++++++
 2 files changed, 98 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java b/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java
index 44685de468281b..9590f2c57cfa5a 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/AbstractActionInputPrefetcher.java
@@ -548,11 +548,15 @@ public void finalizeAction(Action action, MetadataHandler metadataHandler) {
         inputsToDownload.add(output);
       }
 
-      for (Pattern pattern : patternsToDownload) {
-        if (pattern.matcher(output.getExecPathString()).matches()) {
-          outputsToDownload.add(output);
-          break;
+      if (output.isTreeArtifact()) {
+        var children = metadataHandler.getTreeArtifactChildren((SpecialArtifact) output);
+        for (var file : children) {
+          if (outputMatchesPattern(file)) {
+            outputsToDownload.add(file);
+          }
         }
+      } else if (outputMatchesPattern(output)) {
+        outputsToDownload.add(output);
       }
     }
 
@@ -565,6 +569,15 @@ public void finalizeAction(Action action, MetadataHandler metadataHandler) {
     }
   }
 
+  private boolean outputMatchesPattern(Artifact output) {
+    for (var pattern : patternsToDownload) {
+      if (pattern.matcher(output.getExecPathString()).matches()) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   public void flushOutputTree() throws InterruptedException {
     downloadCache.awaitInProgressTasks();
   }
diff --git a/src/test/java/com/google/devtools/build/lib/remote/BuildWithoutTheBytesIntegrationTestBase.java b/src/test/java/com/google/devtools/build/lib/remote/BuildWithoutTheBytesIntegrationTestBase.java
index f53137b7879dd7..cdab72bae779bd 100644
--- a/src/test/java/com/google/devtools/build/lib/remote/BuildWithoutTheBytesIntegrationTestBase.java
+++ b/src/test/java/com/google/devtools/build/lib/remote/BuildWithoutTheBytesIntegrationTestBase.java
@@ -101,6 +101,87 @@ public void downloadOutputsWithRegex() throws Exception {
     assertOutputsDoNotExist("//:foobar");
   }
 
+  @Test
+  public void downloadOutputsWithRegex_treeOutput_regexMatchesTreeFile() throws Exception {
+    // Disable on Windows since it fails for unknown reasons.
+    // TODO(chiwang): Enable it on windows.
+    if (OS.getCurrent() == OS.WINDOWS) {
+      return;
+    }
+
+    writeOutputDirRule();
+    write(
+        "BUILD",
+        "load(':output_dir.bzl', 'output_dir')",
+        "output_dir(",
+        "  name = 'foo',",
+        "  manifest = ':manifest',",
+        ")");
+    write("manifest", "file-1", "file-2", "file-3");
+    addOptions("--experimental_remote_download_regex=.*foo/file-2$");
+
+    buildTarget("//:foo");
+    waitDownloads();
+
+    assertValidOutputFile("foo/file-2", "file-2\n");
+    assertOutputDoesNotExist("foo/file-1");
+    assertOutputDoesNotExist("foo/file-3");
+  }
+
+  @Test
+  public void downloadOutputsWithRegex_treeOutput_regexMatchesTreeRoot() throws Exception {
+    writeOutputDirRule();
+    write(
+        "BUILD",
+        "load(':output_dir.bzl', 'output_dir')",
+        "output_dir(",
+        "  name = 'foo',",
+        "  manifest = ':manifest',",
+        ")");
+    write("manifest", "file-1", "file-2", "file-3");
+    addOptions("--experimental_remote_download_regex=.*foo$");
+
+    buildTarget("//:foo");
+    waitDownloads();
+
+    assertThat(getOutputPath("foo").exists()).isTrue();
+    assertOutputDoesNotExist("foo/file-1");
+    assertOutputDoesNotExist("foo/file-2");
+    assertOutputDoesNotExist("foo/file-3");
+  }
+
+  @Test
+  public void downloadOutputsWithRegex_regexMatchParentPath_filesNotDownloaded() throws Exception {
+    write(
+        "BUILD",
+        "genrule(",
+        "  name = 'file-1',",
+        "  srcs = [],",
+        "  outs = ['foo/file-1'],",
+        "  cmd = 'echo file-1 > $@',",
+        ")",
+        "genrule(",
+        "  name = 'file-2',",
+        "  srcs = [],",
+        "  outs = ['foo/file-2'],",
+        "  cmd = 'echo file-2 > $@',",
+        ")",
+        "genrule(",
+        "  name = 'file-3',",
+        "  srcs = [],",
+        "  outs = ['foo/file-3'],",
+        "  cmd = 'echo file-3 > $@',",
+        ")");
+    addOptions("--experimental_remote_download_regex=.*foo$");
+
+    buildTarget("//:file-1", "//:file-2", "//:file-3");
+    waitDownloads();
+
+    assertOutputDoesNotExist("foo/file-1");
+    assertOutputDoesNotExist("foo/file-2");
+    assertOutputDoesNotExist("foo/file-3");
+  }
+
   @Test
   public void intermediateOutputsAreInputForLocalActions_prefetchIntermediateOutputs()
       throws Exception {