BEP events retain ActionInputMap holding only important artifacts.

To support IDE use-cases, 06c9566 changed target completion events to keep
a reference to the ActionInputMap in order to keep around metadata specifying if
an output file is a directory or not. This worked but caused an unexpected
increase in heap high-watermark, causing OOMs.

The unexpected heap usage was traced to large _validation output groups holding
tens of thousands of artifacts per target. This output group is not included in
BEP so the original change was in fact retaining far more memory than necessary
to implement the intended optimization.

With this change, a second ActionInputMap is prepared which includes only the
"important" artifacts and is sized appropriately. Important artifacts are those
which are not produced by internal output groups and are exactly the artifacts
whose URIs are written to BEP. Maintaining this second ActionInputMap costs
additional latency in CompletionFunction proportional to the number of
important artifacts, though care is taken to avoid doubling the runtime of
CompletionFunction in the case where all artifacts are important.

PiperOrigin-RevId: 378793974

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

@@ -20,6 +20,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.FilesetManifest.RelativeSymlinkBehaviorWithoutError;
+import com.google.devtools.build.lib.bugreport.BugReport;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.Map;
@@ -47,7 +48,11 @@ public class CompletionContext {
   private final ArtifactPathResolver pathResolver;
   private final Map<Artifact, ImmutableCollection<? extends Artifact>> expandedArtifacts;
   private final Map<Artifact, ImmutableList<FilesetOutputSymlink>> expandedFilesets;
-  private final ActionInputMap inputMap;
+  // Only contains the metadata for 'important' artifacts of the Target/Aspect that completed. Any
+  // 'unimportant' artifacts produced by internal output groups (most importantly, _validation) will
+  // not be included to avoid retaining many GB on the heap. This ActionInputMap must only be
+  // consulted with respect to known-important artifacts (eg. artifacts referenced in BEP).
+  private final ActionInputMap importantInputMap;
   private final boolean expandFilesets;
   private final boolean fullyResolveFilesetLinks;
 
@@ -57,14 +62,14 @@ public class CompletionContext {
       Map<Artifact, ImmutableCollection<? extends Artifact>> expandedArtifacts,
       Map<Artifact, ImmutableList<FilesetOutputSymlink>> expandedFilesets,
       ArtifactPathResolver pathResolver,
-      ActionInputMap inputMap,
+      ActionInputMap importantInputMap,
       boolean expandFilesets,
       boolean fullyResolveFilesetLinks) {
     this.execRoot = execRoot;
     this.expandedArtifacts = expandedArtifacts;
     this.expandedFilesets = expandedFilesets;
     this.pathResolver = pathResolver;
-    this.inputMap = inputMap;
+    this.importantInputMap = importantInputMap;
     this.expandFilesets = expandFilesets;
     this.fullyResolveFilesetLinks = fullyResolveFilesetLinks;
   }
@@ -75,6 +80,7 @@ public static CompletionContext create(
       boolean expandFilesets,
       boolean fullyResolveFilesetSymlinks,
       ActionInputMap inputMap,
+      ActionInputMap importantInputMap,
       PathResolverFactory pathResolverFactory,
       Path execRoot,
       String workspaceName) {
@@ -88,7 +94,7 @@ public static CompletionContext create(
         expandedArtifacts,
         expandedFilesets,
         pathResolver,
-        inputMap,
+        importantInputMap,
         expandFilesets,
         fullyResolveFilesetSymlinks);
   }
@@ -98,8 +104,11 @@ public ArtifactPathResolver pathResolver() {
   }
 
   /** Returns true if the given artifact is guaranteed to be a file (and not a directory). */
-  public boolean isOutputFile(Artifact artifact) {
-    FileArtifactValue metadata = inputMap.getMetadata(artifact);
+  public boolean isGuaranteedToBeOutputFile(Artifact artifact) {
+    FileArtifactValue metadata = importantInputMap.getMetadata(artifact);
+    // If we have no metadata for an output file that will be reported in BEP, return that the
+    // output is not guaranteed to be a file. (We expect this to happen for baseline_coverage.dat
+    // files when coverage is enabled.)
     if (metadata == null) {
       return false;
     }
@@ -124,7 +133,15 @@ public void visitArtifacts(Iterable<Artifact> artifacts, ArtifactReceiver receiv
                   : RelativeSymlinkBehaviorWithoutError.RESOLVE);
         }
       } else if (artifact.isTreeArtifact()) {
-        if (FileArtifactValue.OMITTED_FILE_MARKER.equals(inputMap.getMetadata(artifact))) {
+        FileArtifactValue treeArtifactMetadata = importantInputMap.getMetadata(artifact);
+        if (treeArtifactMetadata == null) {
+          BugReport.sendBugReport(
+              new IllegalStateException(
+                  String.format(
+                      "missing artifact metadata for tree artifact: %s",
+                      artifact.toDebugString())));
+        }
+        if (FileArtifactValue.OMITTED_FILE_MARKER.equals(treeArtifactMetadata)) {
           // Expansion can be missing for omitted tree artifacts -- skip the whole tree.
           continue;
         }
@@ -163,6 +180,7 @@ private void visitFileset(
   /** A function that accepts an {@link Artifact}. */
   public interface ArtifactReceiver {
     void accept(Artifact artifact);
+
     void acceptFilesetMapping(Artifact fileset, PathFragment relName, Path targetFile);
   }
 

src/main/java/com/google/devtools/build/lib/analysis/TopLevelArtifactHelper.java

@@ -81,9 +81,12 @@ public boolean isIncomplete() {
   @Immutable
   public static final class ArtifactsToBuild {
     private final ImmutableMap<String, ArtifactsInOutputGroup> artifacts;
+    private final boolean allOutputGroupsImportant;
 
-    private ArtifactsToBuild(ImmutableMap<String, ArtifactsInOutputGroup> artifacts) {
+    private ArtifactsToBuild(
+        ImmutableMap<String, ArtifactsInOutputGroup> artifacts, boolean allOutputGroupsImportant) {
       this.artifacts = checkNotNull(artifacts);
+      this.allOutputGroupsImportant = allOutputGroupsImportant;
     }
 
     /** Returns the artifacts that the user should know about. */
@@ -115,6 +118,14 @@ public NestedSet<Artifact> getAllArtifacts() {
     public ImmutableMap<String, ArtifactsInOutputGroup> getAllArtifactsByOutputGroup() {
       return artifacts;
     }
+
+    /**
+     * Returns if all of the output groups returned by {@link #getAllArtifactsByOutputGroup()} are
+     * "important" - implying that all artifacts will be reported in BEP events.
+     */
+    public boolean areAllOutputGroupsImportant() {
+      return allOutputGroupsImportant;
+    }
   }
 
   private TopLevelArtifactHelper() {
@@ -193,7 +204,7 @@ static ArtifactsToBuild getAllArtifactsToBuild(
       TopLevelArtifactContext context) {
     ImmutableMap.Builder<String, ArtifactsInOutputGroup> allOutputGroups =
         ImmutableMap.builderWithExpectedSize(context.outputGroups().size());
-
+    boolean allOutputGroupsImportant = true;
     for (String outputGroup : context.outputGroups()) {
       NestedSetBuilder<Artifact> results = NestedSetBuilder.stableOrder();
 
@@ -213,13 +224,16 @@ static ArtifactsToBuild getAllArtifactsToBuild(
       boolean isImportantGroup =
           !outputGroup.startsWith(OutputGroupInfo.HIDDEN_OUTPUT_GROUP_PREFIX);
 
+      allOutputGroupsImportant &= isImportantGroup;
+
       ArtifactsInOutputGroup artifacts =
           new ArtifactsInOutputGroup(isImportantGroup, /*incomplete=*/ false, results.build());
 
       allOutputGroups.put(outputGroup, artifacts);
     }
 
-    return new ArtifactsToBuild(allOutputGroups.build());
+    return new ArtifactsToBuild(
+        allOutputGroups.build(), /*allOutputGroupsImportant=*/ allOutputGroupsImportant);
   }
 
   /**

src/main/java/com/google/devtools/build/lib/runtime/NamedArtifactGroup.java

@@ -75,7 +75,7 @@ public Collection<LocalFile> referencedLocalFiles() {
       ExpandedArtifact expandedArtifact = (ExpandedArtifact) elem;
       if (expandedArtifact.relPath == null) {
         LocalFileType outputType =
-            completionContext.isOutputFile(expandedArtifact.artifact)
+            completionContext.isGuaranteedToBeOutputFile(expandedArtifact.artifact)
                 ? LocalFileType.OUTPUT_FILE
                 : LocalFileType.OUTPUT;
         artifacts.add(

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

@@ -13,7 +13,6 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
-
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -55,7 +54,9 @@
 import com.google.devtools.build.skyframe.SkyValue;
 import com.google.devtools.build.skyframe.ValueOrException2;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import javax.annotation.Nullable;
 import net.starlark.java.syntax.Location;
 
@@ -166,15 +167,33 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     ValueT value = valueAndArtifactsToBuild.first;
     ArtifactsToBuild artifactsToBuild = valueAndArtifactsToBuild.second;
 
-    // Avoid iterating over nested set twice.
     ImmutableList<Artifact> allArtifacts = artifactsToBuild.getAllArtifacts().toList();
     Map<SkyKey, ValueOrException2<ActionExecutionException, SourceArtifactException>> inputDeps =
         env.getValuesOrThrow(
             Artifact.keys(allArtifacts),
             ActionExecutionException.class,
             SourceArtifactException.class);
 
+    boolean allArtifactsAreImportant = artifactsToBuild.areAllOutputGroupsImportant();
+
     ActionInputMap inputMap = new ActionInputMap(inputDeps.size());
+    // Prepare an ActionInputMap for important artifacts separately, to be used by BEP events. The
+    // _validation output group can contain orders of magnitude more unimportant artifacts than
+    // there are important artifacts, and BEP events will retain the ActionInputMap until the
+    // event is delivered to transports. If the BEP events reference *all* artifacts it can increase
+    // heap high-watermark by multiple GB.
+    ActionInputMap importantInputMap;
+    Set<Artifact> importantArtifactSet;
+    if (allArtifactsAreImportant) {
+      importantArtifactSet = ImmutableSet.of();
+      importantInputMap = inputMap;
+    } else {
+      ImmutableList<Artifact> importantArtifacts =
+          artifactsToBuild.getImportantArtifacts().toList();
+      importantArtifactSet = new HashSet<>(importantArtifacts);
+      importantInputMap = new ActionInputMap(importantArtifacts.size());
+    }
+
     Map<Artifact, ImmutableCollection<? extends Artifact>> expandedArtifacts = new HashMap<>();
     Map<Artifact, ImmutableList<FilesetOutputSymlink>> expandedFilesets = new HashMap<>();
     Map<SpecialArtifact, ArchivedTreeArtifact> archivedTreeArtifacts = new HashMap<>();
@@ -209,6 +228,22 @@ public SkyValue compute(SkyKey skyKey, Environment env)
                 artifactValue,
                 env,
                 currentConsumer);
+            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,
+              // topLevelFilesets). MetadataConsumerForMetrics.NO_OP is used to avoid
+              // double-counting.
+              ActionInputMapHelper.addToMap(
+                  importantInputMap,
+                  expandedArtifacts,
+                  archivedTreeArtifacts,
+                  expandedFilesets,
+                  topLevelFilesets,
+                  input,
+                  artifactValue,
+                  env,
+                  MetadataConsumerForMetrics.NO_OP);
+            }
           }
         }
       } catch (ActionExecutionException e) {
@@ -238,14 +273,14 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       }
     }
 
-    final CompletionContext ctx;
-    ctx =
+    CompletionContext ctx =
         CompletionContext.create(
             expandedArtifacts,
             expandedFilesets,
             key.topLevelArtifactContext().expandFilesets(),
             key.topLevelArtifactContext().fullyResolveFilesetSymlinks(),
             inputMap,
+            importantInputMap,
             pathResolverFactory,
             skyframeActionExecutor.getExecRoot(),
             workspaceNameValue.getName());

src/test/java/com/google/devtools/build/lib/buildtool/BUILD

@@ -555,6 +555,24 @@ java_library(
     ],
 )
 
+java_test(
+    name = "TargetCompleteEventKeepsMinimalMetadataTest",
+    srcs = ["TargetCompleteEventKeepsMinimalMetadataTest.java"],
+    tags = ["manual"],
+    deps = [
+        "//src/main/java/com/google/devtools/build/lib:runtime",
+        "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
+        "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
+        "//src/main/java/com/google/devtools/build/lib/analysis:configured_target",
+        "//src/main/java/com/google/devtools/build/lib/cmdline",
+        "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
+        "//src/test/java/com/google/devtools/build/lib/buildtool/util",
+        "//third_party:guava",
+        "//third_party:junit4",
+        "//third_party:truth",
+    ],
+)
+
 java_test(
     name = "BuildResultTestCase",
     srcs = ["BuildResultTestCase.java"],