Improve artifact->owner label accounting in two ways. First, don't do the full mapping unless requested. This gets rid of any performance issue for the vast majority of builds. Second, if requested, use a custom data structure so that we don't have to create a full HashSet for artifacts whose only owning labels are their own owner labels.

PiperOrigin-RevId: 206610370

4 files changed, 157 insertions, 47 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/AnalysisResult.java b/src/main/java/com/google/devtools/build/lib/analysis/AnalysisResult.java
index ff8dbe435b..1fc7048aa5 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/AnalysisResult.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/AnalysisResult.java
@@ -16,12 +16,9 @@ package com.google.devtools.build.lib.analysis;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.SetMultimap;
 import com.google.devtools.build.lib.actions.ActionGraph;
-import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.PackageRoots;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationCollection;
-import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.skyframe.AspectValue;
 import java.util.Collection;
 import java.util.List;
@@ -37,7 +34,7 @@ public final class AnalysisResult {
   private final ImmutableSet<ConfiguredTarget> targetsToSkip;
   @Nullable private final String error;
   private final ActionGraph actionGraph;
-  private final SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels;
+  private final ArtifactsToOwnerLabels topLevelArtifactsToOwnerLabels;
   private final ImmutableSet<ConfiguredTarget> parallelTests;
   private final ImmutableSet<ConfiguredTarget> exclusiveTests;
   @Nullable private final TopLevelArtifactContext topLevelContext;
@@ -54,7 +51,7 @@ public final class AnalysisResult {
       Collection<ConfiguredTarget> targetsToSkip,
       @Nullable String error,
       ActionGraph actionGraph,
-      SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels,
+      ArtifactsToOwnerLabels topLevelArtifactsToOwnerLabels,
       Collection<ConfiguredTarget> parallelTests,
       Collection<ConfiguredTarget> exclusiveTests,
       TopLevelArtifactContext topLevelContext,
@@ -122,7 +119,7 @@ public final class AnalysisResult {
     return targetsToSkip;
   }
 
-  public SetMultimap<Artifact, Label> getTopLevelArtifactsToOwnerLabels() {
+  public ArtifactsToOwnerLabels getTopLevelArtifactsToOwnerLabels() {
     return topLevelArtifactsToOwnerLabels;
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/ArtifactsToOwnerLabels.java b/src/main/java/com/google/devtools/build/lib/analysis/ArtifactsToOwnerLabels.java
new file mode 100644
index 0000000000..e8cc545a3a
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/analysis/ArtifactsToOwnerLabels.java
@@ -0,0 +1,116 @@
+// Copyright 2018 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.analysis;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.SetMultimap;
+import com.google.common.collect.Sets;
+import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.cmdline.Label;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * Morally a {@link SetMultimap} from artifacts to the labels that own them, which may not be just
+ * {@link Artifact#getOwnerLabel} depending on the build. Optimizes for artifacts that are owned by
+ * their {@link Artifact#getOwnerLabel} by storing them separately.
+ */
+public class ArtifactsToOwnerLabels {
+  private final SetMultimap<Artifact, Label> artifactToMultipleOrDifferentOwnerLabels;
+  private final Set<Artifact> artifactsOwnedOnlyByTheirLabels;
+
+  private ArtifactsToOwnerLabels(
+      SetMultimap<Artifact, Label> artifactToMultipleOrDifferentOwnerLabels,
+      Set<Artifact> artifactsOwnedOnlyByTheirLabels) {
+    this.artifactToMultipleOrDifferentOwnerLabels = artifactToMultipleOrDifferentOwnerLabels;
+    this.artifactsOwnedOnlyByTheirLabels = artifactsOwnedOnlyByTheirLabels;
+  }
+
+  public Set<Label> getOwners(Artifact artifact) {
+    if (artifactsOwnedOnlyByTheirLabels.contains(artifact)) {
+      Preconditions.checkState(
+          !artifactToMultipleOrDifferentOwnerLabels.containsKey(artifact),
+          "Artifact %s incorrectly in multiple categories",
+          artifact);
+      Label ownerLabel = artifact.getOwnerLabel();
+      if (ownerLabel == null) {
+        return ImmutableSet.of();
+      }
+      return ImmutableSet.of(ownerLabel);
+    }
+    return Preconditions.checkNotNull(
+        artifactToMultipleOrDifferentOwnerLabels.get(artifact), artifact);
+  }
+
+  public Set<Artifact> getArtifacts() {
+    return Sets.union(
+        artifactsOwnedOnlyByTheirLabels, artifactToMultipleOrDifferentOwnerLabels.keySet());
+  }
+
+  public Builder toBuilder() {
+    return new Builder(
+        HashMultimap.create(artifactToMultipleOrDifferentOwnerLabels),
+        new HashSet<>(artifactsOwnedOnlyByTheirLabels));
+  }
+
+  static class Builder {
+    private final SetMultimap<Artifact, Label> artifactToMultipleOrDifferentOwnerLabels;
+    private final Set<Artifact> artifactsOwnedOnlyByTheirLabels;
+
+    Builder() {
+      this(HashMultimap.create(), new HashSet<>());
+    }
+
+    private Builder(
+        SetMultimap<Artifact, Label> artifactToMultipleOrDifferentOwnerLabels,
+        Set<Artifact> artifactsOwnedOnlyByTheirLabels) {
+      this.artifactToMultipleOrDifferentOwnerLabels = artifactToMultipleOrDifferentOwnerLabels;
+      this.artifactsOwnedOnlyByTheirLabels = artifactsOwnedOnlyByTheirLabels;
+    }
+
+    public Builder addArtifact(Artifact artifact) {
+      if (artifactToMultipleOrDifferentOwnerLabels.containsKey(artifact)) {
+        Label ownerLabel = artifact.getOwnerLabel();
+        if (ownerLabel != null) {
+          artifactToMultipleOrDifferentOwnerLabels.put(artifact, artifact.getOwnerLabel());
+        }
+      } else {
+        artifactsOwnedOnlyByTheirLabels.add(artifact);
+      }
+      return this;
+    }
+
+    public Builder addArtifact(Artifact artifact, Label label) {
+      Preconditions.checkNotNull(label, artifact);
+      if (label.equals(artifact.getOwnerLabel())) {
+        addArtifact(artifact);
+      } else {
+        artifactToMultipleOrDifferentOwnerLabels.put(artifact, label);
+        if (artifactsOwnedOnlyByTheirLabels.remove(artifact)) {
+          // Redoing this call now that we have a mismatched label will force addition into the
+          // multimap.
+          addArtifact(artifact);
+        }
+      }
+      return this;
+    }
+
+    ArtifactsToOwnerLabels build() {
+      return new ArtifactsToOwnerLabels(
+          artifactToMultipleOrDifferentOwnerLabels, artifactsOwnedOnlyByTheirLabels);
+    }
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/BuildView.java b/src/main/java/com/google/devtools/build/lib/analysis/BuildView.java
index ec049cd5fd..9995de0c73 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/BuildView.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/BuildView.java
@@ -18,13 +18,11 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ArrayListMultimap;
-import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Multimap;
-import com.google.common.collect.SetMultimap;
 import com.google.common.collect.Sets;
 import com.google.common.eventbus.EventBus;
 import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
@@ -377,7 +375,8 @@ public class BuildView {
       allTargetsToTest = filterTestsByTargets(configuredTargets, testsToRun);
     }
 
-    SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels = HashMultimap.create();
+    ArtifactsToOwnerLabels.Builder artifactsToOwnerLabelsBuilder =
+        new ArtifactsToOwnerLabels.Builder();
     Set<ConfiguredTarget> parallelTests = new HashSet<>();
     Set<ConfiguredTarget> exclusiveTests = new HashSet<>();
 
@@ -385,15 +384,15 @@ public class BuildView {
     Collection<Artifact> buildInfoArtifacts =
         skyframeExecutor.getWorkspaceStatusArtifacts(eventHandler);
     Preconditions.checkState(buildInfoArtifacts.size() == 2, buildInfoArtifacts);
-    addArtifactsWithNoOwner(buildInfoArtifacts, topLevelArtifactsToOwnerLabels);
+    buildInfoArtifacts.forEach(artifactsToOwnerLabelsBuilder::addArtifact);
 
     // Extra actions
     addExtraActionsIfRequested(
-        viewOptions, configuredTargets, aspects, topLevelArtifactsToOwnerLabels, eventHandler);
+        viewOptions, configuredTargets, aspects, artifactsToOwnerLabelsBuilder, eventHandler);
 
     // Coverage
     NestedSet<Artifact> baselineCoverageArtifacts =
-        getBaselineCoverageArtifacts(configuredTargets, topLevelArtifactsToOwnerLabels);
+        getBaselineCoverageArtifacts(configuredTargets, artifactsToOwnerLabelsBuilder);
     if (coverageReportActionFactory != null) {
       CoverageReportActionsWrapper actionsWrapper;
       actionsWrapper =
@@ -408,8 +407,7 @@ public class BuildView {
       if (actionsWrapper != null) {
         ImmutableList<ActionAnalysisMetadata> actions = actionsWrapper.getActions();
         skyframeExecutor.injectCoverageReportData(actions);
-        addArtifactsWithNoOwner(
-            actionsWrapper.getCoverageOutputs(), topLevelArtifactsToOwnerLabels);
+        actionsWrapper.getCoverageOutputs().forEach(artifactsToOwnerLabelsBuilder::addArtifact);
       }
     }
 
@@ -453,7 +451,7 @@ public class BuildView {
         targetsToSkip,
         error,
         actionGraph,
-        topLevelArtifactsToOwnerLabels,
+        artifactsToOwnerLabelsBuilder.build(),
         parallelTests,
         exclusiveTests,
         topLevelOptions,
@@ -462,11 +460,6 @@ public class BuildView {
         topLevelTargetsWithConfigs);
   }
 
-  private static void addArtifactsWithNoOwner(
-      Collection<Artifact> artifacts, SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels) {
-    artifacts.forEach((a) -> topLevelArtifactsToOwnerLabels.put(a, a.getOwnerLabel()));
-  }
-
   @Nullable
   public static String createErrorMessage(
       TargetPatternPhaseValue loadingResult,
@@ -483,7 +476,7 @@ public class BuildView {
 
   private static NestedSet<Artifact> getBaselineCoverageArtifacts(
       Collection<ConfiguredTarget> configuredTargets,
-      SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels) {
+      ArtifactsToOwnerLabels.Builder topLevelArtifactsToOwnerLabels) {
     NestedSetBuilder<Artifact> baselineCoverageArtifacts = NestedSetBuilder.stableOrder();
     for (ConfiguredTarget target : configuredTargets) {
       InstrumentedFilesProvider provider = target.getProvider(InstrumentedFilesProvider.class);
@@ -503,7 +496,7 @@ public class BuildView {
       AnalysisOptions viewOptions,
       Collection<ConfiguredTarget> configuredTargets,
       Collection<AspectValue> aspects,
-      SetMultimap<Artifact, Label> artifactsToTopLevelLabelsMap,
+      ArtifactsToOwnerLabels.Builder artifactsToTopLevelLabelsMap,
       ExtendedEventHandler eventHandler) {
     RegexFilter filter = viewOptions.extraActionFilter;
     for (ConfiguredTarget target : configuredTargets) {
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/TopLevelArtifactHelper.java b/src/main/java/com/google/devtools/build/lib/analysis/TopLevelArtifactHelper.java
index 853ecb93a8..ff153030d6 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/TopLevelArtifactHelper.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/TopLevelArtifactHelper.java
@@ -17,16 +17,16 @@ package com.google.devtools.build.lib.analysis;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.HashMultimap;
-import com.google.common.collect.SetMultimap;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.test.TestProvider;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
+import com.google.devtools.build.lib.profiler.AutoProfiler;
 import com.google.devtools.build.lib.skyframe.AspectValue;
 import com.google.devtools.build.lib.util.RegexFilter;
+import java.util.logging.Logger;
 import javax.annotation.Nullable;
 
 /**
@@ -34,6 +34,7 @@ import javax.annotation.Nullable;
  * extra top-level artifacts into the build.
  */
 public final class TopLevelArtifactHelper {
+  private static Logger logger = Logger.getLogger(TopLevelArtifactHelper.class.getName());
 
   /** Set of {@link Artifact}s in an output group. */
   @Immutable
@@ -119,46 +120,49 @@ public final class TopLevelArtifactHelper {
   }
 
   @VisibleForTesting
-  public static SetMultimap<Artifact, Label> makeTopLevelArtifactsToOwnerLabels(
+  public static ArtifactsToOwnerLabels makeTopLevelArtifactsToOwnerLabels(
       AnalysisResult analysisResult, Iterable<AspectValue> aspects) {
-    SetMultimap<Artifact, Label> topLevelArtifactsToOwnerLabels =
-        HashMultimap.create(analysisResult.getTopLevelArtifactsToOwnerLabels());
+    try (AutoProfiler ignored = AutoProfiler.logged("assigning owner labels", logger, 10)) {
+
+      ArtifactsToOwnerLabels.Builder artifactsToOwnerLabelsBuilder =
+          analysisResult.getTopLevelArtifactsToOwnerLabels().toBuilder();
     TopLevelArtifactContext artifactContext = analysisResult.getTopLevelContext();
     for (ConfiguredTarget target : analysisResult.getTargetsToBuild()) {
-      addArtifactsWithOwnerLabel(
-          getAllArtifactsToBuild(target, artifactContext).getAllArtifacts(),
-          null,
-          target.getLabel(),
-          topLevelArtifactsToOwnerLabels);
+        addArtifactsWithOwnerLabel(
+            getAllArtifactsToBuild(target, artifactContext).getAllArtifacts(),
+            null,
+            target.getLabel(),
+            artifactsToOwnerLabelsBuilder);
     }
     for (AspectValue aspect : aspects) {
-      addArtifactsWithOwnerLabel(
-          getAllArtifactsToBuild(aspect, artifactContext).getAllArtifacts(),
-          null,
-          aspect.getLabel(),
-          topLevelArtifactsToOwnerLabels);
+        addArtifactsWithOwnerLabel(
+            getAllArtifactsToBuild(aspect, artifactContext).getAllArtifacts(),
+            null,
+            aspect.getLabel(),
+            artifactsToOwnerLabelsBuilder);
     }
     if (analysisResult.getTargetsToTest() != null) {
       for (ConfiguredTarget target : analysisResult.getTargetsToTest()) {
-        addArtifactsWithOwnerLabel(
-            TestProvider.getTestStatusArtifacts(target),
-            null,
-            target.getLabel(),
-            topLevelArtifactsToOwnerLabels);
+          addArtifactsWithOwnerLabel(
+              TestProvider.getTestStatusArtifacts(target),
+              null,
+              target.getLabel(),
+              artifactsToOwnerLabelsBuilder);
       }
     }
-    // TODO(dslomov): Artifacts to test from aspects?
-    return topLevelArtifactsToOwnerLabels;
+      // TODO(dslomov): Artifacts to test from aspects?
+      return artifactsToOwnerLabelsBuilder.build();
+    }
   }
 
-  public static void addArtifactsWithOwnerLabel(
+  static void addArtifactsWithOwnerLabel(
       Iterable<Artifact> artifacts,
       @Nullable RegexFilter filter,
       Label ownerLabel,
-      SetMultimap<Artifact, Label> artifactToTopLevelLabels) {
+      ArtifactsToOwnerLabels.Builder artifactsToOwnerLabelsBuilder) {
     for (Artifact artifact : artifacts) {
       if (filter == null || filter.isIncluded(artifact.getOwnerLabel().toString())) {
-        artifactToTopLevelLabels.put(artifact, ownerLabel);
+        artifactsToOwnerLabelsBuilder.addArtifact(artifact, ownerLabel);
       }
     }
   }