Batch all DependencyResolver#getTarget calls. This leads to some duplicate iteration, but that should be cheap, while requesting packages sequentially can hurt...

PiperOrigin-RevId: 208126130

1 files changed, 13 insertions, 3 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/TransitionsOutputFormatterCallback.java b/src/main/java/com/google/devtools/build/lib/query2/TransitionsOutputFormatterCallback.java
index 27685267e8..81298a46f7 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/TransitionsOutputFormatterCallback.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/TransitionsOutputFormatterCallback.java
@@ -19,6 +19,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Streams;
 import com.google.devtools.build.lib.analysis.AspectCollection;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.Dependency;
@@ -58,6 +59,8 @@ import java.io.OutputStream;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
@@ -242,9 +245,16 @@ public class TransitionsOutputFormatterCallback extends CqueryThreadsafeCallback
     }
 
     @Override
-    protected Target getTarget(Target from, Label label, NestedSetBuilder<Cause> rootCauses)
-        throws InterruptedException {
-      return partialResultMap.get(label);
+    protected Map<Label, Target> getTargets(
+        Iterable<Label> labels,
+        Target fromTarget,
+        NestedSetBuilder<Cause> rootCauses,
+        int labelsSizeHint) {
+      return Streams.stream(labels)
+          .distinct()
+          .filter(Objects::nonNull)
+          .filter(partialResultMap::containsKey)
+          .collect(Collectors.toMap(Function.identity(), partialResultMap::get));
     }
 
     @Override