Some code cleanups and improvements around the various DTC visitor implementations in SkyQuery:

(1) Get rid of the 'errorReporter' param used in the two unbounded 'deps' implementations. In practice, this callback is always used right alongside the normal query callback. So we can re-implement it as a single, composite callback.

(2) Revert the recent change to processResultsAndReturnTargets (and restore the old method name). It general, it's a bad idea to return Targets outside that method, because it means we aren't respecting the MultisetSemaphore<PackageIdentifier>. Along with (1), this fixes a real issue where DepsUnboundedVisitor#processPartialResults was using Targets (and thus keeping their Packages live in memory) outside of the package semaphore.

RELNOTES: None
PiperOrigin-RevId: 206606747

7 files changed, 76 insertions, 60 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/AbstractUnfilteredTTVDTCVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/AbstractUnfilteredTTVDTCVisitor.java
new file mode 100644
index 0000000000..9f86d60c49
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/query2/AbstractUnfilteredTTVDTCVisitor.java
@@ -0,0 +1,58 @@
+// 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.query2;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Multimap;
+import com.google.devtools.build.lib.query2.engine.Callback;
+import com.google.devtools.build.lib.query2.engine.Uniquifier;
+import com.google.devtools.build.skyframe.SkyKey;
+
+/**
+ * Helper class for visiting the TTV-only DTC of some given TTV keys, via BFS following all
+ * TTV->TTV dep edges. Disallowed edge filtering is *not* performed.
+ */
+public abstract class AbstractUnfilteredTTVDTCVisitor<T> extends AbstractSkyKeyParallelVisitor<T> {
+  protected final SkyQueryEnvironment env;
+
+  protected AbstractUnfilteredTTVDTCVisitor(
+      SkyQueryEnvironment env,
+      Uniquifier<SkyKey> uniquifier,
+      int visitBatchSize,
+      int processResultsBatchSize,
+      Callback<T> callback) {
+    super(uniquifier, callback, visitBatchSize, processResultsBatchSize);
+    this.env = env;
+  }
+
+  @Override
+  protected Visit getVisitResult(Iterable<SkyKey> ttvKeys) throws InterruptedException {
+    Multimap<SkyKey, SkyKey> deps = env.getUnfilteredDirectDepsOfSkyKeys(ttvKeys);
+    return new Visit(
+        /*keysToUseForResult=*/ deps.keySet(),
+        /*keysToVisit=*/ deps.values()
+        .stream()
+        .filter(SkyQueryEnvironment.IS_TTV)
+        .collect(ImmutableList.toImmutableList()));
+  }
+
+  @Override
+  protected Iterable<SkyKey> preprocessInitialVisit(Iterable<SkyKey> keys) {
+    // ParallelVisitorCallback passes in TTV keys.
+    Preconditions.checkState(Iterables.all(keys, SkyQueryEnvironment.IS_TTV), keys);
+    return keys;
+  }
+}
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 c38834e423..4fb020f534 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
@@ -47,7 +47,6 @@ class DepsUnboundedVisitor extends AbstractEdgeVisitor<SkyKey> {
   private final Uniquifier<SkyKey> validDepUniquifier;
 
   private final boolean depsNeedFiltering;
-  private final Callback<Target> errorReporter;
   private final QueryExpressionContext<Target> context;
 
   DepsUnboundedVisitor(
@@ -56,12 +55,10 @@ class DepsUnboundedVisitor extends AbstractEdgeVisitor<SkyKey> {
       Callback<Target> callback,
       MultisetSemaphore<PackageIdentifier> packageSemaphore,
       boolean depsNeedFiltering,
-      Callback<Target> errorReporter,
       QueryExpressionContext<Target> context) {
     super(env, callback, packageSemaphore);
     this.validDepUniquifier = validDepUniquifier;
     this.depsNeedFiltering = depsNeedFiltering;
-    this.errorReporter = errorReporter;
     this.context = context;
   }
 
@@ -77,7 +74,6 @@ class DepsUnboundedVisitor extends AbstractEdgeVisitor<SkyKey> {
     private final Callback<Target> callback;
     private final MultisetSemaphore<PackageIdentifier> packageSemaphore;
     private final boolean depsNeedFiltering;
-    private final Callback<Target> errorReporter;
     private final QueryExpressionContext<Target> context;
 
     Factory(
@@ -85,14 +81,12 @@ class DepsUnboundedVisitor extends AbstractEdgeVisitor<SkyKey> {
         Callback<Target> callback,
         MultisetSemaphore<PackageIdentifier> packageSemaphore,
         boolean depsNeedFiltering,
-        Callback<Target> errorReporter,
         QueryExpressionContext<Target> context) {
       this.env = env;
       this.validDepUniquifier = env.createSkyKeyUniquifier();
       this.callback = callback;
       this.packageSemaphore = packageSemaphore;
       this.depsNeedFiltering = depsNeedFiltering;
-      this.errorReporter = errorReporter;
       this.context = context;
     }
 
@@ -104,7 +98,6 @@ class DepsUnboundedVisitor extends AbstractEdgeVisitor<SkyKey> {
           callback,
           packageSemaphore,
           depsNeedFiltering,
-          errorReporter,
           context);
     }
   }
@@ -157,13 +150,6 @@ class DepsUnboundedVisitor extends AbstractEdgeVisitor<SkyKey> {
   }
 
   @Override
-  protected void processPartialResults(
-      Iterable<SkyKey> keysToUseForResult, Callback<Target> callback)
-      throws QueryException, InterruptedException {
-    errorReporter.process(processResultsAndReturnTargets(keysToUseForResult, callback));
-  }
-
-  @Override
   protected SkyKey getNewNodeFromEdge(SkyKey visit) {
     return visit;
   }
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 5ad164c565..030ec3e2de 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
@@ -127,7 +127,7 @@ public class ParallelSkyQueryUtils {
               () -> {
                 Callback<Target> visitorCallback =
                     ParallelVisitor.createParallelVisitorCallback(
-                        new UnfilteredTransitiveTraversalValueDTCVisitor.Factory(
+                        new UnfilteredSkyKeyTTVDTCVisitor.Factory(
                             env,
                             env.createSkyKeyUniquifier(),
                             processResultsBatchSize,
@@ -177,14 +177,13 @@ public class ParallelSkyQueryUtils {
       QueryExpressionContext<Target> context,
       Callback<Target> callback,
       MultisetSemaphore<PackageIdentifier> packageSemaphore,
-      boolean depsNeedFiltering,
-      Callback<Target> errorReporter) {
+      boolean depsNeedFiltering) {
     return env.eval(
         expression,
         context,
         ParallelVisitor.createParallelVisitorCallback(
             new DepsUnboundedVisitor.Factory(
-                env, callback, packageSemaphore, depsNeedFiltering, errorReporter, context)));
+                env, callback, packageSemaphore, depsNeedFiltering, context)));
   }
 
   static class DepAndRdep {
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 cba76ba8c6..0c759a8444 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
@@ -1217,16 +1217,14 @@ public class SkyQueryEnvironment extends AbstractBlazeQueryEnvironment<Target>
   public QueryTaskFuture<Void> getDepsUnboundedParallel(
       QueryExpression expression,
       QueryExpressionContext<Target> context,
-      Callback<Target> callback,
-      Callback<Target> errorReporter) {
+      Callback<Target> callback) {
     return ParallelSkyQueryUtils.getDepsUnboundedParallel(
         SkyQueryEnvironment.this,
         expression,
         context,
         callback,
         packageSemaphore,
-        /*depsNeedFiltering=*/ !dependencyFilter.equals(DependencyFilter.ALL_DEPS),
-        errorReporter);
+        /*depsNeedFiltering=*/ !dependencyFilter.equals(DependencyFilter.ALL_DEPS));
   }
 
   @ThreadSafe
diff --git a/src/main/java/com/google/devtools/build/lib/query2/UnfilteredTransitiveTraversalValueDTCVisitor.java b/src/main/java/com/google/devtools/build/lib/query2/UnfilteredSkyKeyTTVDTCVisitor.java
index e1d10eb82e..4f73b150f4 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/UnfilteredTransitiveTraversalValueDTCVisitor.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/UnfilteredSkyKeyTTVDTCVisitor.java
@@ -13,11 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.query2;
 
-import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Multimap;
 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.AggregateAllCallback;
@@ -25,24 +21,21 @@ import com.google.devtools.build.lib.query2.engine.Uniquifier;
 import com.google.devtools.build.skyframe.SkyKey;
 
 /**
- * Helper class that computes the TTV-only DTC of some given TTV keys, via BFS following all
- * TTV->TTV dep edges. Disallowed edge filtering is *not* performed.
+ * Helper class for visiting the TTV-only DTC of some given TTV keys, and feeding those TTVs to a
+ * callback.
  */
-// TODO(nharmata): Unify with UnfilteredUnboundedDepsSkyKeyVisitor.
-class UnfilteredTransitiveTraversalValueDTCVisitor extends AbstractSkyKeyParallelVisitor<SkyKey> {
-  private final SkyQueryEnvironment env;
-
-  private UnfilteredTransitiveTraversalValueDTCVisitor(
+class UnfilteredSkyKeyTTVDTCVisitor extends AbstractUnfilteredTTVDTCVisitor<SkyKey> {
+  private UnfilteredSkyKeyTTVDTCVisitor(
       SkyQueryEnvironment env,
       Uniquifier<SkyKey> uniquifier,
       int processResultsBatchSize,
       AggregateAllCallback<SkyKey, ImmutableSet<SkyKey>> aggregateAllCallback) {
     super(
+        env,
         uniquifier,
-        aggregateAllCallback,
         ParallelSkyQueryUtils.VISIT_BATCH_SIZE,
-        processResultsBatchSize);
-    this.env = env;
+        processResultsBatchSize,
+        aggregateAllCallback);
   }
 
   static class Factory implements ParallelVisitor.Factory {
@@ -64,7 +57,7 @@ class UnfilteredTransitiveTraversalValueDTCVisitor extends AbstractSkyKeyParalle
 
     @Override
     public ParallelVisitor<SkyKey, SkyKey> create() {
-      return new UnfilteredTransitiveTraversalValueDTCVisitor(
+      return new UnfilteredSkyKeyTTVDTCVisitor(
           env, uniquifier, processResultsBatchSize, aggregateAllCallback);
     }
   }
@@ -75,22 +68,4 @@ class UnfilteredTransitiveTraversalValueDTCVisitor extends AbstractSkyKeyParalle
       throws QueryException, InterruptedException {
     callback.process(keysToUseForResult);
   }
-
-  @Override
-  protected Visit getVisitResult(Iterable<SkyKey> ttvKeys) throws InterruptedException {
-    Multimap<SkyKey, SkyKey> deps = env.getUnfilteredDirectDepsOfSkyKeys(ttvKeys);
-    return new Visit(
-        /*keysToUseForResult=*/ deps.keySet(),
-        /*keysToVisit=*/ deps.values()
-            .stream()
-            .filter(SkyQueryEnvironment.IS_TTV)
-            .collect(ImmutableList.toImmutableList()));
-  }
-
-  @Override
-  protected Iterable<SkyKey> preprocessInitialVisit(Iterable<SkyKey> keys) {
-    // ParallelVisitorCallback passes in TTV keys.
-    Preconditions.checkState(Iterables.all(keys, SkyQueryEnvironment.IS_TTV), keys);
-    return keys;
-  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/DepsFunction.java b/src/main/java/com/google/devtools/build/lib/query2/engine/DepsFunction.java
index 190e29137e..1a430e3686 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/DepsFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/DepsFunction.java
@@ -62,10 +62,11 @@ final class DepsFunction implements QueryFunction {
       return streamableEnv.getDepsUnboundedParallel(
           queryExpression,
           context,
-          callback,
-          targets -> {
+          /*callback=*/ partialResult -> {
+            callback.process(partialResult);
             ThreadSafeMutableSet<T> set = env.createThreadSafeMutableSet();
-            Iterables.addAll(set, targets);
+            Iterables.addAll(set, partialResult);
+            // Ensure the proper error messages are reported.
             env.buildTransitiveClosure(expression, set, /*maxDepth=*/ 1);
           });
     }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/StreamableQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/engine/StreamableQueryEnvironment.java
index 07cf9a226c..a2939b688c 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/engine/StreamableQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/engine/StreamableQueryEnvironment.java
@@ -44,6 +44,5 @@ public interface StreamableQueryEnvironment<T> extends QueryEnvironment<T> {
   QueryTaskFuture<Void> getDepsUnboundedParallel(
       QueryExpression expression,
       QueryExpressionContext<T> context,
-      Callback<T> callback,
-      Callback<T> errorReporter);
+      Callback<T> callback);
 }