[6.0.0] Include full tree artifact in inputs when prefetcher doesn't support partial tree artifacts. (#17013)

The actions generated by SpawnActionTemplate and CppCompileActionTemplate contain individual TreeFileArtifacts, but not the respective TreeArtifact, in their inputs. This causes the input prefetcher to crash when building without the bytes, as it is only able to fetch entire tree artifacts and expects their input metadata to be available.

This PR introduces an ActionInputPrefetcher#supportsPartialTreeArtifactInputs method signaling whether prefetching partial tree artifacts is supported; if not, the entire tree artifact is included in the inputs of any such actions.

This is a temporary workaround to unblock the 6.0.0 release. The long-term plan is to fix #16333 to enable prefetching partial tree artifacts, and remove this workaround.

Fixes #16987.

PiperOrigin-RevId: 495050207
Change-Id: I7d29713085d2cf84ce5302394fc18ff2a96ec4be
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java b/src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java
index 0435d17..fb2b07c 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java
@@ -28,6 +28,11 @@
           // Do nothing.
           return immediateVoidFuture();
         }
+
+        @Override
+        public boolean supportsPartialTreeArtifactInputs() {
+          return true;
+        }
       };
 
   /**
@@ -39,4 +44,10 @@
    */
   ListenableFuture<Void> prefetchFiles(
       Iterable<? extends ActionInput> inputs, MetadataProvider metadataProvider);
+
+  /**
+   * Whether the prefetcher is able to fetch individual files in a tree artifact without fetching
+   * the entire tree artifact.
+   */
+  boolean supportsPartialTreeArtifactInputs();
 }
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java
index 5ae39e9..d128e83 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java
@@ -62,6 +62,12 @@
   }
 
   @Override
+  public boolean supportsPartialTreeArtifactInputs() {
+    // This prefetcher is unable to fetch only individual files inside a tree artifact.
+    return false;
+  }
+
+  @Override
   protected void prefetchVirtualActionInput(VirtualActionInput input) throws IOException {
     if (!(input instanceof EmptyActionInput)) {
       Path outputPath = execRoot.getRelative(input.getExecPath());
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
index 952e2dd..4a116c0 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
@@ -1253,7 +1253,8 @@
           topLevelFilesets,
           input,
           value,
-          env);
+          env,
+          skyframeActionExecutor.supportsPartialTreeArtifactInputs());
     }
 
     if (actionExecutionFunctionExceptionHandler != null) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionInputMapHelper.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionInputMapHelper.java
index f024131..4126592 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionInputMapHelper.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionInputMapHelper.java
@@ -25,6 +25,7 @@
 import com.google.devtools.build.lib.actions.Artifact.ArchivedTreeArtifact;
 import com.google.devtools.build.lib.actions.Artifact.DerivedArtifact;
 import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
+import com.google.devtools.build.lib.actions.Artifact.TreeFileArtifact;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.FilesetOutputSymlink;
 import com.google.devtools.build.lib.analysis.actions.SymlinkAction;
@@ -47,7 +48,8 @@
       Map<Artifact, ImmutableList<FilesetOutputSymlink>> topLevelFilesets,
       Artifact key,
       SkyValue value,
-      Environment env)
+      Environment env,
+      boolean prefetcherSupportsPartialTreeArtifacts)
       throws InterruptedException {
     addToMap(
         inputMap,
@@ -58,7 +60,8 @@
         key,
         value,
         env,
-        MetadataConsumerForMetrics.NO_OP);
+        MetadataConsumerForMetrics.NO_OP,
+        prefetcherSupportsPartialTreeArtifacts);
   }
 
   /**
@@ -74,7 +77,8 @@
       Artifact key,
       SkyValue value,
       Environment env,
-      MetadataConsumerForMetrics consumer)
+      MetadataConsumerForMetrics consumer,
+      boolean prefetcherSupportsPartialTreeArtifacts)
       throws InterruptedException {
     if (value instanceof RunfilesArtifactValue) {
       // Note: we don't expand the .runfiles/MANIFEST file into the inputs. The reason for that
@@ -120,16 +124,35 @@
           /*depOwner=*/ key);
       consumer.accumulate(treeArtifactValue);
     } else if (value instanceof ActionExecutionValue) {
-      FileArtifactValue metadata = ((ActionExecutionValue) value).getExistingFileArtifactValue(key);
-      inputMap.put(key, metadata, key);
-      if (key.isFileset()) {
-        ImmutableList<FilesetOutputSymlink> filesets = getFilesets(env, (SpecialArtifact) key);
-        if (filesets != null) {
-          topLevelFilesets.put(key, filesets);
-          consumer.accumulate(filesets);
-        }
+      if (!prefetcherSupportsPartialTreeArtifacts && key instanceof TreeFileArtifact) {
+        // If we're unable to prefetch individual files in a tree artifact, include the full tree
+        // artifact in the action inputs. This makes actions that consume partial tree artifacts
+        // (such as the ones generated by SpawnActionTemplate or CppCompileActionTemplate) less
+        // efficient, but is needed until https://github.com/bazelbuild/bazel/issues/16333 is fixed.
+        SpecialArtifact treeArtifact = key.getParent();
+        TreeArtifactValue treeArtifactValue =
+            ((ActionExecutionValue) value).getTreeArtifactValue(treeArtifact);
+        expandTreeArtifactAndPopulateArtifactData(
+            treeArtifact,
+            treeArtifactValue,
+            expandedArtifacts,
+            archivedTreeArtifacts,
+            inputMap,
+            /* depOwner= */ treeArtifact);
+        consumer.accumulate(treeArtifactValue);
       } else {
-        consumer.accumulate(metadata);
+        FileArtifactValue metadata =
+            ((ActionExecutionValue) value).getExistingFileArtifactValue(key);
+        inputMap.put(key, metadata, key);
+        if (key.isFileset()) {
+          ImmutableList<FilesetOutputSymlink> filesets = getFilesets(env, (SpecialArtifact) key);
+          if (filesets != null) {
+            topLevelFilesets.put(key, filesets);
+            consumer.accumulate(filesets);
+          }
+        } else {
+          consumer.accumulate(metadata);
+        }
       }
     } else {
       Preconditions.checkArgument(value instanceof FileArtifactValue, "Unexpected value %s", value);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java
index df7cb75..1dfec25 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java
@@ -218,7 +218,8 @@
                 input,
                 artifactValue,
                 env,
-                currentConsumer);
+                currentConsumer,
+                skyframeActionExecutor.supportsPartialTreeArtifactInputs());
             if (!allArtifactsAreImportant && importantArtifactSet.contains(input)) {
               // Calling #addToMap a second time with `input` and `artifactValue` will perform no-op
               // updates to the secondary collections passed in (eg. expandedArtifacts,
@@ -233,7 +234,8 @@
                   input,
                   artifactValue,
                   env,
-                  MetadataConsumerForMetrics.NO_OP);
+                  MetadataConsumerForMetrics.NO_OP,
+                  skyframeActionExecutor.supportsPartialTreeArtifactInputs());
             }
           }
         }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
index 6329c46..342d977 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
@@ -320,6 +320,10 @@
         .test(action.getMnemonic());
   }
 
+  boolean supportsPartialTreeArtifactInputs() {
+    return actionInputPrefetcher.supportsPartialTreeArtifactInputs();
+  }
+
   boolean publishTargetSummaries() {
     return options.getOptions(BuildEventProtocolOptions.class).publishTargetSummary;
   }
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutorTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutorTest.java
index c806bc3..5f4bedc 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutorTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutorTest.java
@@ -879,6 +879,7 @@
     // is running. This way, both actions will check the action cache beforehand and try to update
     // the action cache post-build.
     final CountDownLatch inputsRequested = new CountDownLatch(2);
+    skyframeExecutor.configureActionExecutor(/* fileCache= */ null, ActionInputPrefetcher.NONE);
     skyframeExecutor
         .getEvaluator()
         .injectGraphTransformerForTesting(
@@ -1231,6 +1232,7 @@
     ActionTemplate<DummyAction> template2 =
         new DummyActionTemplate(baseOutput, sharedOutput2, ActionOwner.SYSTEM_ACTION_OWNER);
     ActionLookupValue shared2Ct = createActionLookupValue(template2, shared2);
+    skyframeExecutor.configureActionExecutor(/* fileCache= */ null, ActionInputPrefetcher.NONE);
     // Inject the "configured targets" into the graph.
     skyframeExecutor
         .getDifferencerForTesting()
@@ -1645,6 +1647,7 @@
             createActionLookupValue(action2, lc2),
             null,
             NestedSetBuilder.create(Order.STABLE_ORDER, Event.warn("analysis warning 2")));
+    skyframeExecutor.configureActionExecutor(/* fileCache= */ null, ActionInputPrefetcher.NONE);
     skyframeExecutor
         .getDifferencerForTesting()
         .inject(ImmutableMap.of(lc1, ctValue1, lc2, ctValue2));
diff --git a/src/test/shell/bazel/remote/build_without_the_bytes_test.sh b/src/test/shell/bazel/remote/build_without_the_bytes_test.sh
index 477a2f0..3680f84 100755
--- a/src/test/shell/bazel/remote/build_without_the_bytes_test.sh
+++ b/src/test/shell/bazel/remote/build_without_the_bytes_test.sh
@@ -70,14 +70,7 @@
   declare -r EXE_EXT=""
 fi
 
-function test_cc_tree() {
-  if [[ "$PLATFORM" == "darwin" ]]; then
-    # TODO(b/37355380): This test is disabled due to RemoteWorker not supporting
-    # setting SDKROOT and DEVELOPER_DIR appropriately, as is required of
-    # action executors in order to select the appropriate Xcode toolchain.
-    return 0
-  fi
-
+function setup_cc_tree() {
   mkdir -p a
   cat > a/BUILD <<EOF
 load(":tree.bzl", "mytree")
@@ -93,11 +86,40 @@
 
 mytree = rule(implementation = _tree_impl)
 EOF
+}
+
+function test_cc_tree_remote_executor() {
+  if [[ "$PLATFORM" == "darwin" ]]; then
+    # TODO(b/37355380): This test is disabled due to RemoteWorker not supporting
+    # setting SDKROOT and DEVELOPER_DIR appropriately, as is required of
+    # action executors in order to select the appropriate Xcode toolchain.
+    return 0
+  fi
+
+  setup_cc_tree
+
   bazel build \
       --remote_executor=grpc://localhost:${worker_port} \
       --remote_download_minimal \
       //a:tree_cc >& "$TEST_log" \
-      || fail "Failed to build //a:tree_cc with minimal downloads"
+      || fail "Failed to build //a:tree_cc with remote executor and minimal downloads"
+}
+
+function test_cc_tree_remote_cache() {
+  if [[ "$PLATFORM" == "darwin" ]]; then
+    # TODO(b/37355380): This test is disabled due to RemoteWorker not supporting
+    # setting SDKROOT and DEVELOPER_DIR appropriately, as is required of
+    # action executors in order to select the appropriate Xcode toolchain.
+    return 0
+  fi
+
+  setup_cc_tree
+
+  bazel build \
+      --remote_cache=grpc://localhost:${worker_port} \
+      --remote_download_minimal \
+      //a:tree_cc >& "$TEST_log" \
+      || fail "Failed to build //a:tree_cc with remote cache and minimal downloads"
 }
 
 function test_cc_include_scanning_and_minimal_downloads() {