Allow uniquifiers to throw QueryExceptions.

RELNOTES: None.
PiperOrigin-RevId: 202961077

10 files changed, 34 insertions, 38 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/AbstractSkyKeyParallelVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/AbstractSkyKeyParallelVisitor.java
index 972e6703c6..e296e62f1f 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/AbstractSkyKeyParallelVisitor.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/AbstractSkyKeyParallelVisitor.java
@@ -15,6 +15,7 @@ package com.google.devtools.build.lib.query2;
 
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.query2.engine.Callback;
+import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.Uniquifier;
 import com.google.devtools.build.skyframe.SkyKey;
 
@@ -32,7 +33,7 @@ public abstract class AbstractSkyKeyParallelVisitor<T> extends ParallelVisitor<S
   }
 
   @Override
-  protected ImmutableList<SkyKey> getUniqueValues(Iterable<SkyKey> values) {
+  protected ImmutableList<SkyKey> getUniqueValues(Iterable<SkyKey> values) throws QueryException {
     return uniquifier.unique(values);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/DepsUnboundedVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/DepsUnboundedVisitor.java
index cb5e19c196..c38834e423 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/DepsUnboundedVisitor.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/DepsUnboundedVisitor.java
@@ -169,7 +169,8 @@ class DepsUnboundedVisitor extends AbstractEdgeVisitor<SkyKey> {
   }
 
   @Override
-  protected ImmutableList<SkyKey> getUniqueValues(Iterable<SkyKey> keysToVisit) {
+  protected ImmutableList<SkyKey> getUniqueValues(Iterable<SkyKey> keysToVisit)
+      throws QueryException {
     // Legit deps are already filtered using env.getFwdDeps().
     return validDepUniquifier.unique(keysToVisit);
   }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/ParallelVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/ParallelVisitor.java
index 3ea72273a8..45c92f0882 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/ParallelVisitor.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/ParallelVisitor.java
@@ -165,7 +165,8 @@ public abstract class ParallelVisitor<T, V> {
       throws QueryException, InterruptedException;
 
   /** Gets the {@link Visit} representing the local visitation of the given {@code values}. */
-  protected abstract Visit getVisitResult(Iterable<T> values) throws InterruptedException;
+  protected abstract Visit getVisitResult(Iterable<T> values)
+      throws QueryException, InterruptedException;
 
   /** Gets the equivalent of {@link Visit#keysToVisit} for the entry-level SkyKeys. */
   protected abstract Iterable<T> preprocessInitialVisit(Iterable<SkyKey> keys);
@@ -175,7 +176,7 @@ public abstract class ParallelVisitor<T, V> {
    *
    * <p>Used to dedupe visitations before adding them to {@link #processingQueue}.
    */
-  protected abstract ImmutableList<T> getUniqueValues(Iterable<T> values);
+  protected abstract ImmutableList<T> getUniqueValues(Iterable<T> values) throws QueryException;
 
   /** Gets tasks to visit pending keys. */
   protected Iterable<Task> getVisitTasks(Collection<T> pendingKeysToVisit) {
@@ -212,7 +213,7 @@ public abstract class ParallelVisitor<T, V> {
     }
 
     @Override
-    void process() throws InterruptedException {
+    void process() throws QueryException, InterruptedException {
       Visit visit = getVisitResult(keysToVisit);
       for (Iterable<SkyKey> keysToUseForResultBatch :
           Iterables.partition(visit.keysToUseForResult, processResultsBatchSize)) {
diff --git a/src/main/java/com/google/devtools/build/lib/query2/PostAnalysisQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/PostAnalysisQueryEnvironment.java
index 956abf5605..89aae592bc 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/PostAnalysisQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/PostAnalysisQueryEnvironment.java
@@ -403,8 +403,7 @@ abstract class PostAnalysisQueryEnvironment<T> extends AbstractBlazeQueryEnviron
 
   @Override
   public Uniquifier<T> createUniquifier() {
-    return new UniquifierImpl<>(
-        getConfiguredTargetKeyExtractor(), SkyQueryEnvironment.DEFAULT_THREAD_COUNT);
+    return new UniquifierImpl<>(getConfiguredTargetKeyExtractor());
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/query2/RdepsBoundedVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/RdepsBoundedVisitor.java
index d6c59c642c..b0f182b893 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/RdepsBoundedVisitor.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/RdepsBoundedVisitor.java
@@ -29,6 +29,7 @@ import com.google.devtools.build.lib.query2.ParallelSkyQueryUtils.DepAndRdep;
 import com.google.devtools.build.lib.query2.ParallelSkyQueryUtils.DepAndRdepAtDepth;
 import com.google.devtools.build.lib.query2.engine.Callback;
 import com.google.devtools.build.lib.query2.engine.MinDepthUniquifier;
+import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.UniquifierImpl;
 import com.google.devtools.build.lib.query2.engine.Uniquifier;
 import com.google.devtools.build.skyframe.SkyKey;
@@ -211,7 +212,7 @@ class RdepsBoundedVisitor extends AbstractEdgeVisitor<DepAndRdepAtDepth> {
 
   @Override
   protected ImmutableList<DepAndRdepAtDepth> getUniqueValues(
-      Iterable<DepAndRdepAtDepth> depAndRdepAtDepths) {
+      Iterable<DepAndRdepAtDepth> depAndRdepAtDepths) throws QueryException {
     // See the comment in RdepsUnboundedVisitor#getUniqueValues.
     return depAndRdepAtDepthUniquifier.unique(
         Iterables.filter(
diff --git a/src/main/java/com/google/devtools/build/lib/query2/RdepsUnboundedVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/RdepsUnboundedVisitor.java
index e2720490ef..d9f1e0fb70 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/RdepsUnboundedVisitor.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/RdepsUnboundedVisitor.java
@@ -26,6 +26,7 @@ import com.google.devtools.build.lib.concurrent.MultisetSemaphore;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.query2.ParallelSkyQueryUtils.DepAndRdep;
 import com.google.devtools.build.lib.query2.engine.Callback;
+import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryUtil.UniquifierImpl;
 import com.google.devtools.build.lib.query2.engine.Uniquifier;
 import com.google.devtools.build.skyframe.SkyKey;
@@ -109,7 +110,8 @@ class RdepsUnboundedVisitor extends AbstractEdgeVisitor<DepAndRdep> {
   }
 
   @Override
-  protected Visit getVisitResult(Iterable<DepAndRdep> depAndRdeps) throws InterruptedException {
+  protected Visit getVisitResult(Iterable<DepAndRdep> depAndRdeps)
+      throws QueryException, InterruptedException {
     Collection<SkyKey> validRdeps = new ArrayList<>();
 
     // Multimap of dep to all the reverse deps in this visitation. Used to filter out the
@@ -150,11 +152,13 @@ class RdepsUnboundedVisitor extends AbstractEdgeVisitor<DepAndRdep> {
       packageSemaphore.releaseAll(pkgIdsNeededForTargetification);
     }
 
-    ImmutableList<SkyKey> uniqueValidRdeps =
-        validRdeps
-            .stream()
-            .filter(validRdepUniquifier::unique)
-            .collect(ImmutableList.toImmutableList());
+    ImmutableList.Builder<SkyKey> uniqueValidRdepsbuilder = ImmutableList.builder();
+    for (SkyKey rdep : validRdeps) {
+      if (validRdepUniquifier.unique(rdep)) {
+        uniqueValidRdepsbuilder.add(rdep);
+      }
+    }
+    ImmutableList<SkyKey> uniqueValidRdeps = uniqueValidRdepsbuilder.build();
 
     // Retrieve the reverse deps as SkyKeys and defer the targetification and filtering to next
     // recursive visitation.
@@ -188,7 +192,8 @@ class RdepsUnboundedVisitor extends AbstractEdgeVisitor<DepAndRdep> {
   }
 
   @Override
-  protected ImmutableList<DepAndRdep> getUniqueValues(Iterable<DepAndRdep> depAndRdeps) {
+  protected ImmutableList<DepAndRdep> getUniqueValues(Iterable<DepAndRdep> depAndRdeps)
+      throws QueryException {
     // See the javadoc for 'validRdepUniquifier'.
     //
     // N.B. - Except for the visitation roots, 'depAndRdepUniquifier' is actually completely
diff --git a/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
index a53f1d610d..2df8008e2c 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/SkyQueryEnvironment.java
@@ -600,13 +600,13 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
 
   @Override
   public <V> MutableMap<Target, V> createMutableMap() {
-    return new MutableKeyExtractorBackedMapImpl<Target, Label, V>(TargetKeyExtractor.INSTANCE);
+    return new MutableKeyExtractorBackedMapImpl<>(TargetKeyExtractor.INSTANCE);
   }
 
   @ThreadSafe
   @Override
   public Uniquifier<Target> createUniquifier() {
-    return createTargetUniquifier();
+    return new UniquifierImpl<>(TargetKeyExtractor.INSTANCE);
   }
 
   @ThreadSafe
@@ -621,13 +621,8 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   }
 
   @ThreadSafe
-  Uniquifier<Target> createTargetUniquifier() {
-    return new UniquifierImpl<>(TargetKeyExtractor.INSTANCE, DEFAULT_THREAD_COUNT);
-  }
-
-  @ThreadSafe
   public Uniquifier<SkyKey> createSkyKeyUniquifier() {
-    return new UniquifierImpl<>(SkyKeyKeyExtractor.INSTANCE, DEFAULT_THREAD_COUNT);
+    return new UniquifierImpl<>(SkyKeyKeyExtractor.INSTANCE);
   }
 
   private ImmutableSet<PathFragment> getBlacklistedExcludes(TargetPatternKey targetPatternKey)
@@ -1126,8 +1121,8 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     // memory. We should have a threshold for when to invoke the callback with a batch, and also a
     // separate, larger, bound on the number of targets being processed at the same time.
     private final ThreadSafeOutputFormatterCallback<Target> callback;
-    private final Uniquifier<Target> uniquifier =
-        new UniquifierImpl<>(TargetKeyExtractor.INSTANCE, DEFAULT_THREAD_COUNT);
+    private final UniquifierImpl<Target, ?> uniquifier =
+        new UniquifierImpl<>(TargetKeyExtractor.INSTANCE);
     private final Object pendingLock = new Object();
     private List<Target> pending = new ArrayList<>();
     private int batchThreshold;
diff --git a/src/main/java/com/google/devtools/build/lib/query2/TransitiveTraversalValueDTCVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/TransitiveTraversalValueDTCVisitor.java
index 1a5962c7c4..3d887f5ec2 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/TransitiveTraversalValueDTCVisitor.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/TransitiveTraversalValueDTCVisitor.java
@@ -92,7 +92,7 @@ class TransitiveTraversalValueDTCVisitor extends ParallelVisitor<SkyKey, SkyKey>
   }
 
   @Override
-  protected ImmutableList<SkyKey> getUniqueValues(Iterable<SkyKey> values) {
+  protected ImmutableList<SkyKey> getUniqueValues(Iterable<SkyKey> values) throws QueryException {
     return uniquifier.unique(values);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java
index fce662fcef..120d53bee1 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java
@@ -224,15 +224,8 @@ public final class QueryUtil {
     private final Set<K> alreadySeen;
 
     public UniquifierImpl(KeyExtractor<T, K> extractor) {
-      this(extractor, /*concurrencyLevel=*/ 1);
-    }
-
-    public UniquifierImpl(KeyExtractor<T, K> extractor, int concurrencyLevel) {
       this.extractor = extractor;
-      this.alreadySeen =
-          Collections.newSetFromMap(
-              new ConcurrentHashMap<>(
-                  /*initialCapacity=*/ concurrencyLevel, /*loadFactor=*/ 0.75f));
+      this.alreadySeen = Collections.newSetFromMap(new ConcurrentHashMap<>());
     }
 
     @Override
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/Uniquifier.java b/src/main/java/com/google/devtools/build/lib/query2/engine/Uniquifier.java
index 390a14930b..e363c1f7df 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/Uniquifier.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/Uniquifier.java
@@ -31,14 +31,14 @@ public interface Uniquifier<T> {
   boolean uniquePure(T newElement);
 
   /**
-   * Returns whether {@code newElement} has been seen before by {@link #unique(T)} or
-   * {@link #unique(Iterable)}.
+   * Returns whether {@code newElement} has been seen before by {@link #unique(T)} or {@link
+   * #unique(Iterable)}.
    */
-  boolean unique(T newElement);
+  boolean unique(T newElement) throws QueryException;
 
   /**
    * Returns the subset of {@code newElements} that haven't been seen before by {@link #unique(T)}
    * or {@link #unique(Iterable)}.
    */
-  ImmutableList<T> unique(Iterable<T> newElements);
+  ImmutableList<T> unique(Iterable<T> newElements) throws QueryException;
 }