Include full tree artifact in inputs when prefetcher doesn't support partial tree artifacts.

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 bazelbuild#16333 to enable prefetching partial tree artifacts, and remove this workaround.

Fixes bazelbuild#16987.

PiperOrigin-RevId: 495050207
Change-Id: I7d29713085d2cf84ce5302394fc18ff2a96ec4be

Author: tjgq

src/main/java/com/google/devtools/build/lib/actions/ActionInputPrefetcher.java

@@ -28,6 +28,11 @@ public ListenableFuture<Void> prefetchFiles(
           // Do nothing.
           return immediateVoidFuture();
         }
+
+        @Override
+        public boolean supportsPartialTreeArtifactInputs() {
+          return true;
+        }
       };
 
   /**
@@ -39,4 +44,10 @@ public ListenableFuture<Void> prefetchFiles(
    */
   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();
 }

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

@@ -61,6 +61,12 @@ class RemoteActionInputFetcher extends AbstractActionInputPrefetcher {
     this.remoteCache = Preconditions.checkNotNull(remoteCache);
   }
 
+  @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)) {

src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java

@@ -1253,7 +1253,8 @@ private <S extends ActionInputMapSink, R> R accumulateInputs(
           topLevelFilesets,
           input,
           value,
-          env);
+          env,
+          skyframeActionExecutor.supportsPartialTreeArtifactInputs());
     }
 
     if (actionExecutionFunctionExceptionHandler != null) {

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 @@ static void addToMap(
       Map<Artifact, ImmutableList<FilesetOutputSymlink>> topLevelFilesets,
       Artifact key,
       SkyValue value,
-      Environment env)
+      Environment env,
+      boolean prefetcherSupportsPartialTreeArtifacts)
       throws InterruptedException {
     addToMap(
         inputMap,
@@ -58,7 +60,8 @@ static void addToMap(
         key,
         value,
         env,
-        MetadataConsumerForMetrics.NO_OP);
+        MetadataConsumerForMetrics.NO_OP,
+        prefetcherSupportsPartialTreeArtifacts);
   }
 
   /**
@@ -74,7 +77,8 @@ static void addToMap(
       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 @@ static void addToMap(
           /*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);

src/main/java/com/google/devtools/build/lib/skyframe/CompletionFunction.java

@@ -205,7 +205,8 @@ public SkyValue compute(SkyKey skyKey, Environment env)
                 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,
@@ -220,7 +221,8 @@ public SkyValue compute(SkyKey skyKey, Environment env)
                   input,
                   artifactValue,
                   env,
-                  MetadataConsumerForMetrics.NO_OP);
+                  MetadataConsumerForMetrics.NO_OP,
+                  skyframeActionExecutor.supportsPartialTreeArtifactInputs());
             }
           }
         }

src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java

@@ -314,6 +314,10 @@ boolean useArchivedTreeArtifacts(ActionAnalysisMetadata action) {
         .test(action.getMnemonic());
   }
 
+  boolean supportsPartialTreeArtifactInputs() {
+    return actionInputPrefetcher.supportsPartialTreeArtifactInputs();
+  }
+
   boolean publishTargetSummaries() {
     return options.getOptions(BuildEventProtocolOptions.class).publishTargetSummary;
   }

src/test/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutorTest.java

@@ -879,6 +879,7 @@ public void testSharedActionsRacing() throws Exception {
     // 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 @@ public void sharedActionTemplate() throws Exception {
     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 @@ public void analysisEventsNotStoredInExecution() throws Exception {
             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));

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

@@ -70,14 +70,7 @@ else
   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 @@ def _tree_impl(ctx):
 
 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() {