1 files changed, 23 insertions, 23 deletions
diff --git a/src/main/java/com/google/devtools/build/lib/query2/ParallelSkyQueryUtils.java b/src/main/java/com/google/devtools/build/lib/query2/ParallelSkyQueryUtils.java
index 95486c6ef9..0052892130 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/ParallelSkyQueryUtils.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/ParallelSkyQueryUtils.java
@@ -36,10 +36,10 @@ import com.google.devtools.build.lib.concurrent.QuiescingExecutor;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.query2.engine.Callback;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment.QueryTaskFuture;
 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.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.PackageValue;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
@@ -77,14 +77,13 @@ class ParallelSkyQueryUtils {
    * Specialized parallel variant of {@link SkyQueryEnvironment#getAllRdeps} that is appropriate
    * when there is no depth-bound.
    */
-  static void getAllRdepsUnboundedParallel(
+  static QueryTaskFuture<Void> getAllRdepsUnboundedParallel(
       SkyQueryEnvironment env,
       QueryExpression expression,
       VariableContext<Target> context,
-      ThreadSafeCallback<Target> callback,
-      MultisetSemaphore<PackageIdentifier> packageSemaphore)
-          throws QueryException, InterruptedException {
-    env.eval(
+      Callback<Target> callback,
+      MultisetSemaphore<PackageIdentifier> packageSemaphore) {
+    return env.eval(
         expression,
         context,
         new SkyKeyBFSVisitorCallback(
@@ -95,10 +94,10 @@ class ParallelSkyQueryUtils {
   static void getRBuildFilesParallel(
       SkyQueryEnvironment env,
       Collection<PathFragment> fileIdentifiers,
-      ThreadSafeCallback<Target> callback,
+      Callback<Target> callback,
       MultisetSemaphore<PackageIdentifier> packageSemaphore)
           throws QueryException, InterruptedException {
-    ThreadSafeUniquifier<SkyKey> keyUniquifier = env.createSkyKeyUniquifier();
+    Uniquifier<SkyKey> keyUniquifier = env.createSkyKeyUniquifier();
     RBuildFilesVisitor visitor =
         new RBuildFilesVisitor(env, keyUniquifier, callback, packageSemaphore);
     visitor.visitAndWaitForCompletion(env.getSkyKeysForFileFragments(fileIdentifiers));
@@ -110,7 +109,7 @@ class ParallelSkyQueryUtils {
 
     private RBuildFilesVisitor(
         SkyQueryEnvironment env,
-        ThreadSafeUniquifier<SkyKey> uniquifier,
+        Uniquifier<SkyKey> uniquifier,
         Callback<Target> callback,
         MultisetSemaphore<PackageIdentifier> packageSemaphore) {
       super(env, uniquifier, callback);
@@ -180,8 +179,8 @@ class ParallelSkyQueryUtils {
 
     private AllRdepsUnboundedVisitor(
         SkyQueryEnvironment env,
-        ThreadSafeUniquifier<Pair<SkyKey, SkyKey>> uniquifier,
-        ThreadSafeCallback<Target> callback,
+        Uniquifier<Pair<SkyKey, SkyKey>> uniquifier,
+        Callback<Target> callback,
         MultisetSemaphore<PackageIdentifier> packageSemaphore) {
       super(env, uniquifier, callback);
       this.packageSemaphore = packageSemaphore;
@@ -190,19 +189,18 @@ class ParallelSkyQueryUtils {
     /**
      * A {@link Factory} for {@link AllRdepsUnboundedVisitor} instances, each of which will be used
      * to perform visitation of the reverse transitive closure of the {@link Target}s passed in a
-     * single {@link ThreadSafeCallback#process} call. Note that all the created
-     * instances share the same {@code ThreadSafeUniquifier<SkyKey>} so that we don't visit the
-     * same Skyframe node more than once.
+     * single {@link Callback#process} call. Note that all the created instances share the same
+     * {@link Uniquifier} so that we don't visit the same Skyframe node more than once.
      */
     private static class Factory implements AbstractSkyKeyBFSVisitor.Factory {
       private final SkyQueryEnvironment env;
-      private final ThreadSafeUniquifier<Pair<SkyKey, SkyKey>> uniquifier;
-      private final ThreadSafeCallback<Target> callback;
+      private final Uniquifier<Pair<SkyKey, SkyKey>> uniquifier;
+      private final Callback<Target> callback;
       private final MultisetSemaphore<PackageIdentifier> packageSemaphore;
 
       private Factory(
         SkyQueryEnvironment env,
-        ThreadSafeCallback<Target> callback,
+        Callback<Target> callback,
         MultisetSemaphore<PackageIdentifier> packageSemaphore) {
         this.env = env;
         this.uniquifier = env.createReverseDepSkyKeyUniquifier();
@@ -341,10 +339,10 @@ class ParallelSkyQueryUtils {
   }
 
   /**
-   * A {@link ThreadSafeCallback} whose {@link ThreadSafeCallback#process} method kicks off a BFS
-   * visitation via a fresh {@link AbstractSkyKeyBFSVisitor} instance.
+   * A {@link Callback} whose {@link Callback#process} method kicks off a BFS visitation via a fresh
+   * {@link AbstractSkyKeyBFSVisitor} instance.
    */
-  private static class SkyKeyBFSVisitorCallback implements ThreadSafeCallback<Target> {
+  private static class SkyKeyBFSVisitorCallback implements Callback<Target> {
     private final AbstractSkyKeyBFSVisitor.Factory visitorFactory;
 
     private SkyKeyBFSVisitorCallback(AbstractSkyKeyBFSVisitor.Factory visitorFactory) {
@@ -355,6 +353,8 @@ class ParallelSkyQueryUtils {
     public void process(Iterable<Target> partialResult)
         throws QueryException, InterruptedException {
       AbstractSkyKeyBFSVisitor<?> visitor = visitorFactory.create();
+      // TODO(nharmata): It's not ideal to have an operation like this in #process that blocks on
+      // another, potentially expensive computation. Refactor to something like "processAsync".
       visitor.visitAndWaitForCompletion(
           SkyQueryEnvironment.makeTransitiveTraversalKeysStrict(partialResult));
     }
@@ -370,7 +370,7 @@ class ParallelSkyQueryUtils {
   @ThreadSafe
   private abstract static class AbstractSkyKeyBFSVisitor<T> {
     protected final SkyQueryEnvironment env;
-    private final ThreadSafeUniquifier<T> uniquifier;
+    private final Uniquifier<T> uniquifier;
     private final Callback<Target> callback;
 
     private final BFSVisitingTaskExecutor executor;
@@ -434,7 +434,7 @@ class ParallelSkyQueryUtils {
             new ThreadFactoryBuilder().setNameFormat("skykey-bfs-visitor %d").build());
 
     private AbstractSkyKeyBFSVisitor(
-        SkyQueryEnvironment env, ThreadSafeUniquifier<T> uniquifier, Callback<Target> callback) {
+        SkyQueryEnvironment env, Uniquifier<T> uniquifier, Callback<Target> callback) {
       this.env = env;
       this.uniquifier = uniquifier;
       this.callback = callback;