Have SkyQueryEnvironment#getRBuildFiles not visit skyframe nodes more than once. This is helpful when there are multiple skylark reverse 'load' paths from the same popular .bzl file to a single BUILD file. As an implementation detail, introduce a few ThreadSafeUniquifier things (these will also be used in followup CLs).

--
MOS_MIGRATED_REVID=132680777

4 files changed, 101 insertions, 52 deletions

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 6623ca695a..826ccdf843 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
@@ -60,10 +60,12 @@ import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
 import com.google.devtools.build.lib.query2.engine.QueryExpressionEvalListener;
 import com.google.devtools.build.lib.query2.engine.QueryExpressionMapper;
+import com.google.devtools.build.lib.query2.engine.QueryUtil.AbstractThreadSafeUniquifier;
 import com.google.devtools.build.lib.query2.engine.RdepsFunction;
 import com.google.devtools.build.lib.query2.engine.StreamableQueryEnvironment;
 import com.google.devtools.build.lib.query2.engine.TargetLiteral;
 import com.google.devtools.build.lib.query2.engine.ThreadSafeCallback;
+import com.google.devtools.build.lib.query2.engine.ThreadSafeUniquifier;
 import com.google.devtools.build.lib.query2.engine.Uniquifier;
 import com.google.devtools.build.lib.query2.engine.VariableContext;
 import com.google.devtools.build.lib.skyframe.BlacklistedPackagePrefixesValue;
@@ -91,7 +93,6 @@ import com.google.devtools.build.skyframe.WalkableGraph.WalkableGraphFactory;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -102,7 +103,6 @@ import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Queue;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Logger;
@@ -513,8 +513,8 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
 
   @ThreadSafe
   @Override
-  public Uniquifier<Target> createUniquifier() {
-    return new ConcurrentUniquifier();
+  public ThreadSafeUniquifier<Target> createUniquifier() {
+    return new ThreadSafeTargetUniquifier(DEFAULT_THREAD_COUNT);
   }
 
   @ThreadSafe
@@ -880,7 +880,8 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   void getRBuildFiles(Collection<PathFragment> fileIdentifiers, Callback<Target> callback)
       throws QueryException, InterruptedException {
     Collection<SkyKey> files = getSkyKeysForFileFragments(fileIdentifiers);
-    Collection<SkyKey> current = graph.getSuccessfulValues(files).keySet();
+    Uniquifier<SkyKey> keyUniquifier = new ThreadSafeSkyKeyUniquifier(/*concurrencyLevel=*/ 1);
+    Collection<SkyKey> current = keyUniquifier.unique(graph.getSuccessfulValues(files).keySet());
     Set<SkyKey> resultKeys = CompactHashSet.create();
     while (!current.isEmpty()) {
       Collection<Iterable<SkyKey>> reverseDeps = graph.getReverseDeps(current).values();
@@ -890,12 +891,16 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
           resultKeys.add(rdep);
           // Every package has a dep on the external package, so we need to include those edges too.
           if (rdep.equals(PackageValue.key(Label.EXTERNAL_PACKAGE_IDENTIFIER))) {
-            current.add(rdep);
+            if (keyUniquifier.unique(rdep)) {
+              current.add(rdep);
+            }
           }
         } else if (!rdep.functionName().equals(SkyFunctions.PACKAGE_LOOKUP)) {
           // Packages may depend on the existence of subpackages, but these edges aren't relevant to
           // rbuildfiles.
-          current.add(rdep);
+          if (keyUniquifier.unique(rdep)) {
+            current.add(rdep);
+          }
         }
       }
       if (resultKeys.size() >= BATCH_CALLBACK_SIZE) {
@@ -934,23 +939,27 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
     }
   }
 
-  @ThreadSafe
-  private static class ConcurrentUniquifier implements Uniquifier<Target> {
+  private static class ThreadSafeTargetUniquifier
+      extends AbstractThreadSafeUniquifier<Target, Label> {
+    protected ThreadSafeTargetUniquifier(int concurrencyLevel) {
+      super(concurrencyLevel);
+    }
 
-    // Note that setting initialCapacity to BATCH_CALLBACK_SIZE is not especially principled.
-    private final Set<Label> seen =
-        Collections.newSetFromMap(
-            new ConcurrentHashMap<Label, Boolean>(BATCH_CALLBACK_SIZE, .75f, DEFAULT_THREAD_COUNT));
+    @Override
+    protected Label extractKey(Target element) {
+      return element.getLabel();
+    }
+  }
+
+  private static class ThreadSafeSkyKeyUniquifier
+      extends AbstractThreadSafeUniquifier<SkyKey, SkyKey> {
+    protected ThreadSafeSkyKeyUniquifier(int concurrencyLevel) {
+      super(concurrencyLevel);
+    }
 
     @Override
-    public ImmutableList<Target> unique(Iterable<Target> newElements) {
-      ImmutableList.Builder<Target> builder = ImmutableList.builder();
-      for (Target newElement : newElements) {
-        if (seen.add(newElement.getLabel())) {
-          builder.add(newElement);
-        }
-      }
-      return builder.build();
+    protected SkyKey extractKey(SkyKey element) {
+      return element;
     }
   }
 
@@ -971,7 +980,8 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   private static class BatchStreamedCallback implements ThreadSafeCallback<Target> {
 
     private final Callback<Target> callback;
-    private final Uniquifier<Target> uniquifier = new ConcurrentUniquifier();
+    private final ThreadSafeUniquifier<Target> uniquifier =
+        new ThreadSafeTargetUniquifier(DEFAULT_THREAD_COUNT);
     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/engine/QueryUtil.java b/src/main/java/com/google/devtools/build/lib/query2/engine/QueryUtil.java
index 68f0a5a22e..7df7a00c9b 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
@@ -17,8 +17,9 @@ package com.google.devtools.build.lib.query2.engine;
 import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.MapMaker;
 import com.google.devtools.build.lib.collect.CompactHashSet;
-
+import java.util.Collections;
 import java.util.Set;
 
 /** Several query utilities to make easier to work with query callbacks and uniquifiers. */
@@ -83,33 +84,57 @@ public final class QueryUtil {
     };
   }
 
-  /**
-   * An uniquifier that uses a CompactHashSet and a key extractor for making the elements unique.
-   *
-   * <p>Using a key extractor allows to improve memory since we don't have to keep the whole element
-   * in the set but just the key.
-   */
-  public abstract static class AbstractUniquifier<T, K> implements Uniquifier<T> {
-
+  /** A trivial {@link Uniquifier} base class. */
+  public abstract static class AbstractUniquifier<T, K>
+      extends AbstractUniquifierBase<T, K> {
     private final CompactHashSet<K> alreadySeen = CompactHashSet.create();
 
     @Override
+    public final boolean unique(T element) {
+      return alreadySeen.add(extractKey(element));
+    }
+
+    /**
+     * Extracts an unique key that can be used to dedupe the given {@code element}.
+     *
+     * <p>Depending on the choice of {@code K}, this enables potential memory optimizations.
+     */
+    protected abstract K extractKey(T element);
+  }
+
+  /** A trivial {@link ThreadSafeUniquifier} base class. */
+  public abstract static class AbstractThreadSafeUniquifier<T, K>
+      extends AbstractUniquifierBase<T, K> implements ThreadSafeUniquifier<T> {
+    private final Set<K> alreadySeen;
+
+    protected AbstractThreadSafeUniquifier(int concurrencyLevel) {
+      this.alreadySeen = Collections.newSetFromMap(
+          new MapMaker().concurrencyLevel(concurrencyLevel).<K, Boolean>makeMap());
+    }
+
+    @Override
+    public final boolean unique(T element) {
+      return alreadySeen.add(extractKey(element));
+    }
+
+    /**
+     * Extracts an unique key that can be used to dedupe the given {@code element}.
+     *
+     * <p>Depending on the choice of {@code K}, this enables potential memory optimizations.
+     */
+    protected abstract K extractKey(T element);
+  }
+
+  private abstract static class AbstractUniquifierBase<T, K> implements Uniquifier<T> {
+    @Override
     public final ImmutableList<T> unique(Iterable<T> newElements) {
       ImmutableList.Builder<T> result = ImmutableList.builder();
       for (T element : newElements) {
-        if (alreadySeen.add(extractKey(element))) {
+        if (unique(element)) {
           result.add(element);
         }
       }
       return result.build();
     }
-
-    /** Extracts an unique key that represents the target. For example the label. */
-    protected abstract K extractKey(T t);
-
-    @Override
-    public String toString() {
-      return this.getClass().getName() + " uniquifier :" + alreadySeen;
-    }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/ThreadSafeUniquifier.java b/src/main/java/com/google/devtools/build/lib/query2/engine/ThreadSafeUniquifier.java
new file mode 100644
index 0000000000..747185582f
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/ThreadSafeUniquifier.java
@@ -0,0 +1,22 @@
+// Copyright 2016 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.query2.engine;
+
+import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
+
+/** Marker interface for a {@link ThreadSafe} {@link Uniquifier}. */
+@ThreadSafe
+public interface ThreadSafeUniquifier<T> extends Uniquifier<T> {
+}
+
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 d401921176..ed2b2376fa 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
@@ -15,19 +15,11 @@ package com.google.devtools.build.lib.query2.engine;
 
 import com.google.common.collect.ImmutableList;
 
-/**
- * A class used for deduplication of {@code Iterable}s. If called with repeated elements or multiple
- * times with repeated elements between calls, it guarantees to output only those elements exactly once.
- */
+/** A helper for deduping values. */
 public interface Uniquifier<T> {
+  /** Returns whether {@code newElement} has been seen before. */
+  boolean unique(T newElement);
 
-  /**
-   * Receives an iterable and returns the list of elements that were not already seen. The
-   * uniqueness need to be guaranteed for elements of the same iterable and multiple calls to the
-   * {@code unique} method.
-   *
-   * @param newElements The new elements to process.
-   * @return The subset of elements not already seen by this Uniquifier.
-   */
+  /** Returns the subset of {@code newElements} that haven't been seen before. */
   ImmutableList<T> unique(Iterable<T> newElements);
 }