Provide parallel implementations of 'allrdeps' and 'rbuildfiles', using a shared naive BFS implementation. Also implement RegexFilterExpression#parEval. -- MOS_MIGRATED_REVID=134598046
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 new file mode 100644 index 0000000..4d9212f --- /dev/null +++ b/src/main/java/com/google/devtools/build/lib/query2/ParallelSkyQueryUtils.java
@@ -0,0 +1,360 @@ +// 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; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; +import com.google.devtools.build.lib.cmdline.Label; +import com.google.devtools.build.lib.collect.CompactHashSet; +import com.google.devtools.build.lib.concurrent.ForkJoinQuiescingExecutor; +import com.google.devtools.build.lib.concurrent.MoreFutures; +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.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.VariableContext; +import com.google.devtools.build.lib.skyframe.PackageValue; +import com.google.devtools.build.lib.skyframe.SkyFunctions; +import com.google.devtools.build.lib.vfs.PathFragment; +import com.google.devtools.build.skyframe.SkyKey; +import java.util.Collection; +import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.ForkJoinTask; +import java.util.concurrent.RecursiveAction; + + +/** + * Parallel implementations of various functionality in {@link SkyQueryEnvironment}. + * + * <p>Special attention is given to memory usage. Naive parallel implementations of query + * functionality would lead to memory blowup. Instead of dealing with {@link Target}s, we try to + * deal with {@link SkyKey}s as much as possible to reduce the number of {@link Package}s forcibly + * in memory at any given time. + */ +// TODO(bazel-team): Be more deliberate about bounding memory usage here. +class ParallelSkyQueryUtils { + private ParallelSkyQueryUtils() { + } + + /** + * Specialized parallel variant of {@link SkyQueryEnvironment#getAllRdeps} that is appropriate + * when there is no depth-bound. + */ + static void getAllRdepsUnboundedParallel( + SkyQueryEnvironment env, + QueryExpression expression, + VariableContext<Target> context, + ThreadSafeCallback<Target> callback, + ForkJoinPool forkJoinPool) + throws QueryException, InterruptedException { + env.eval( + expression, + context, + new SkyKeyBFSVisitorCallback( + new AllRdepsUnboundedVisitor.Factory(env, callback, forkJoinPool))); + } + + /** Specialized parallel variant of {@link SkyQueryEnvironment#getRBuildFiles}. */ + static void getRBuildFilesParallel( + SkyQueryEnvironment env, + Collection<PathFragment> fileIdentifiers, + ThreadSafeCallback<Target> callback, + ForkJoinPool forkJoinPool) + throws QueryException, InterruptedException { + ThreadSafeUniquifier<SkyKey> keyUniquifier = env.createSkyKeyUniquifier(); + RBuildFilesVisitor visitor = new RBuildFilesVisitor(env, forkJoinPool, keyUniquifier, callback); + visitor.visitAndWaitForCompletion(env.getSkyKeysForFileFragments(fileIdentifiers)); + } + + /** A helper class that computes 'rbuildfiles(<blah>)' via BFS. */ + private static class RBuildFilesVisitor extends AbstractSkyKeyBFSVisitor { + private final SkyQueryEnvironment env; + + private RBuildFilesVisitor( + SkyQueryEnvironment env, + ForkJoinPool forkJoinPool, + ThreadSafeUniquifier<SkyKey> uniquifier, + Callback<Target> callback) { + super(forkJoinPool, uniquifier, callback); + this.env = env; + } + + @Override + protected Visit getVisitResult(Iterable<SkyKey> values) throws InterruptedException { + Collection<Iterable<SkyKey>> reverseDeps = env.graph.getReverseDeps(values).values(); + Set<SkyKey> keysToUseForResult = CompactHashSet.create(); + Set<SkyKey> keysToVisitNext = CompactHashSet.create(); + for (SkyKey rdep : Iterables.concat(reverseDeps)) { + if (rdep.functionName().equals(SkyFunctions.PACKAGE)) { + keysToUseForResult.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))) { + keysToVisitNext.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. + keysToVisitNext.add(rdep); + } + } + return new Visit(keysToUseForResult, keysToVisitNext); + } + + @Override + protected Iterable<Target> getTargetsToAddToResult(Iterable<SkyKey> keysToUseForResult) + throws InterruptedException { + return SkyQueryEnvironment.getBuildFilesForPackageValues( + env.graph.getSuccessfulValues(keysToUseForResult).values()); + } + } + + /** A helper class that computes 'allrdeps(<blah>)' via BFS. */ + private static class AllRdepsUnboundedVisitor extends AbstractSkyKeyBFSVisitor { + private final SkyQueryEnvironment env; + + private AllRdepsUnboundedVisitor( + SkyQueryEnvironment env, + ForkJoinPool forkJoinPool, + ThreadSafeUniquifier<SkyKey> uniquifier, + ThreadSafeCallback<Target> callback) { + super(forkJoinPool, uniquifier, callback); + this.env = env; + } + + /** + * 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. + */ + private static class Factory implements AbstractSkyKeyBFSVisitor.Factory { + private final SkyQueryEnvironment env; + private final ForkJoinPool forkJoinPool; + private final ThreadSafeUniquifier<SkyKey> uniquifier; + private final ThreadSafeCallback<Target> callback; + + private Factory( + SkyQueryEnvironment env, + ThreadSafeCallback<Target> callback, + ForkJoinPool forkJoinPool) { + this.env = env; + this.forkJoinPool = forkJoinPool; + this.uniquifier = env.createSkyKeyUniquifier(); + this.callback = callback; + } + + @Override + public AbstractSkyKeyBFSVisitor create() { + return new AllRdepsUnboundedVisitor(env, forkJoinPool, uniquifier, callback); + } + } + + @Override + protected Visit getVisitResult(Iterable<SkyKey> keys) throws InterruptedException { + // TODO(bazel-team): Defer some of this work to the next recursive visitation. Instead, have + // this visitation merely get the Skyframe-land rdeps. + Iterable<SkyKey> keysToVisit = SkyQueryEnvironment.makeTransitiveTraversalKeysStrict( + env.getReverseDeps(env.makeTargetsFromSkyKeys(keys).values())); + return new Visit( + /*keysToUseForResult=*/ keys, + /*keysToVisit=*/ keysToVisit); + } + + @Override + protected Iterable<Target> getTargetsToAddToResult(Iterable<SkyKey> keysToUseForResult) + throws InterruptedException { + return env.makeTargetsFromSkyKeys(keysToUseForResult).values(); + } + } + + /** + * A {@link ThreadSafeCallback} whose {@link ThreadSafeCallback#process} method kicks off a BFS + * visitation via a fresh {@link AbstractSkyKeyBFSVisitor} instance. + */ + private static class SkyKeyBFSVisitorCallback implements ThreadSafeCallback<Target> { + private final AbstractSkyKeyBFSVisitor.Factory visitorFactory; + + private SkyKeyBFSVisitorCallback(AbstractSkyKeyBFSVisitor.Factory visitorFactory) { + this.visitorFactory = visitorFactory; + } + + @Override + public void process(Iterable<Target> partialResult) + throws QueryException, InterruptedException { + AbstractSkyKeyBFSVisitor visitor = visitorFactory.create(); + visitor.visitAndWaitForCompletion( + SkyQueryEnvironment.makeTransitiveTraversalKeysStrict(partialResult)); + } + } + + /** + * A helper class for performing a custom BFS visitation on the Skyframe graph, using + * {@link ForkJoinQuiescingExecutor}. + * + * <p>The choice of {@link ForkJoinPool} over, say, AbstractQueueVisitor backed by a + * ThreadPoolExecutor, is very deliberate. {@link SkyKeyBFSVisitorCallback#process} kicks off + * a visitation and blocks on completion of it. But this visitation may never complete if there + * are a bounded number of threads in the global thread pool used for query evaluation! + */ + @ThreadSafe + private abstract static class AbstractSkyKeyBFSVisitor { + private final ForkJoinPool forkJoinPool; + private final ThreadSafeUniquifier<SkyKey> uniquifier; + private final Callback<Target> callback; + /** The maximum number of keys to visit at once. */ + private static final int VISIT_BATCH_SIZE = 10000; + + private AbstractSkyKeyBFSVisitor( + ForkJoinPool forkJoinPool, + ThreadSafeUniquifier<SkyKey> uniquifier, + Callback<Target> callback) { + this.forkJoinPool = forkJoinPool; + this.uniquifier = uniquifier; + this.callback = callback; + } + + /** Factory for {@link AbstractSkyKeyBFSVisitor} instances. */ + private static interface Factory { + AbstractSkyKeyBFSVisitor create(); + } + + protected static final class Visit { + private final Iterable<SkyKey> keysToUseForResult; + private final Iterable<SkyKey> keysToVisit; + + private Visit(Iterable<SkyKey> keysToUseForResult, Iterable<SkyKey> keysToVisit) { + this.keysToUseForResult = keysToUseForResult; + this.keysToVisit = keysToVisit; + } + } + + void visitAndWaitForCompletion(Iterable<SkyKey> keys) + throws QueryException, InterruptedException { + Iterable<ForkJoinTask<?>> tasks = getTasks(new Visit( + /*keysToUseForResult=*/ ImmutableList.<SkyKey>of(), + /*keysToVisit=*/ keys)); + for (ForkJoinTask<?> task : tasks) { + forkJoinPool.execute(task); + } + try { + MoreFutures.waitForAllInterruptiblyFailFast(tasks); + } catch (ExecutionException ee) { + Throwable cause = ee.getCause(); + if (cause instanceof RuntimeQueryException) { + throw (QueryException) cause.getCause(); + } else if (cause instanceof RuntimeInterruptedException) { + throw (InterruptedException) cause.getCause(); + } else { + throw new IllegalStateException(cause); + } + } + } + + private abstract static class AbstractInternalRecursiveAction extends RecursiveAction { + protected abstract void computeImpl() throws QueryException, InterruptedException; + + @Override + public final void compute() { + try { + computeImpl(); + } catch (QueryException queryException) { + throw new RuntimeQueryException(queryException); + } catch (InterruptedException interruptedException) { + throw new RuntimeInterruptedException(interruptedException); + } + } + } + + private class VisitTask extends AbstractInternalRecursiveAction { + private final Iterable<SkyKey> keysToVisit; + + private VisitTask(Iterable<SkyKey> keysToVisit) { + this.keysToVisit = keysToVisit; + } + + @Override + protected void computeImpl() throws InterruptedException { + ImmutableList<SkyKey> uniqueKeys = uniquifier.unique(keysToVisit); + if (uniqueKeys.isEmpty()) { + return; + } + Iterable<ForkJoinTask<?>> tasks = getTasks(getVisitResult(uniqueKeys)); + for (ForkJoinTask<?> task : tasks) { + task.fork(); + } + for (ForkJoinTask<?> task : tasks) { + task.join(); + } + } + } + + private class GetAndProcessResultsTask extends AbstractInternalRecursiveAction { + private final Iterable<SkyKey> keysToUseForResult; + + private GetAndProcessResultsTask(Iterable<SkyKey> keysToUseForResult) { + this.keysToUseForResult = keysToUseForResult; + } + + @Override + protected void computeImpl() throws QueryException, InterruptedException { + callback.process(getTargetsToAddToResult(keysToUseForResult)); + } + } + + private Iterable<ForkJoinTask<?>> getTasks(Visit visit) { + // Split the given visit request into ForkJoinTasks for visiting keys and ForkJoinTasks for + // getting and outputting results, each of which obeys the separate batch limits. + // TODO(bazel-team): Attempt to group work on targets within the same package. + ImmutableList.Builder<ForkJoinTask<?>> tasksBuilder = ImmutableList.builder(); + for (Iterable<SkyKey> keysToVisitBatch + : Iterables.partition(visit.keysToVisit, VISIT_BATCH_SIZE)) { + tasksBuilder.add(new VisitTask(keysToVisitBatch)); + } + for (Iterable<SkyKey> keysToUseForResultBatch : Iterables.partition( + visit.keysToUseForResult, SkyQueryEnvironment.BATCH_CALLBACK_SIZE)) { + tasksBuilder.add(new GetAndProcessResultsTask(keysToUseForResultBatch)); + } + return tasksBuilder.build(); + } + + /** + * Gets the given {@code keysToUseForResult}'s contribution to the set of {@link Target}s in the + * full visitation. + */ + protected abstract Iterable<Target> getTargetsToAddToResult( + Iterable<SkyKey> keysToUseForResult) throws InterruptedException; + + /** Gets the {@link Visit} representing the local visitation of the given {@code values}. */ + protected abstract Visit getVisitResult(Iterable<SkyKey> values) throws InterruptedException; + } + + private static class RuntimeQueryException extends RuntimeException { + private RuntimeQueryException(QueryException queryException) { + super(queryException); + } + } + + private static class RuntimeInterruptedException extends RuntimeException { + private RuntimeInterruptedException(InterruptedException interruptedException) { + super(interruptedException); + } + } +} +
diff --git a/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java b/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java index d25717b..79614c2 100644 --- a/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java +++ b/src/main/java/com/google/devtools/build/lib/query2/RBuildFilesFunction.java
@@ -83,6 +83,7 @@ Collections2.transform(args, ARGUMENT_TO_PATH_FRAGMENT), (Callback<Target>) callback); } + @SuppressWarnings("unchecked") @Override public <T> void parEval( QueryEnvironment<T> env, @@ -91,6 +92,13 @@ List<Argument> args, ThreadSafeCallback<T> callback, ForkJoinPool forkJoinPool) throws QueryException, InterruptedException { - eval(env, context, expression, args, callback); + if (!(env instanceof SkyQueryEnvironment)) { + throw new QueryException("rbuildfiles can only be used with SkyQueryEnvironment"); + } + ((SkyQueryEnvironment) env) + .getRBuildFilesParallel( + Collections2.transform(args, ARGUMENT_TO_PATH_FRAGMENT), + (ThreadSafeCallback<Target>) callback, + forkJoinPool); } }
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 51fb19c..2913d1e 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
@@ -116,8 +116,8 @@ implements StreamableQueryEnvironment<Target> { // 10k is likely a good balance between using batch efficiently and not blowing up memory. // TODO(janakr): Unify with RecursivePackageProviderBackedTargetPatternResolver's constant. - private static final int BATCH_CALLBACK_SIZE = 10000; - private static final int DEFAULT_THREAD_COUNT = Runtime.getRuntime().availableProcessors(); + static final int BATCH_CALLBACK_SIZE = 10000; + protected static final int DEFAULT_THREAD_COUNT = Runtime.getRuntime().availableProcessors(); private static final int MAX_QUERY_EXPRESSION_LOG_CHARS = 1000; private static final Logger LOG = Logger.getLogger(SkyQueryEnvironment.class.getName()); private static final Function<Target, Label> TARGET_LABEL_FUNCTION = @@ -538,10 +538,20 @@ @ThreadSafe @Override public ThreadSafeUniquifier<Target> createUniquifier() { + return createTargetUniquifier(); + } + + @ThreadSafe + ThreadSafeUniquifier<Target> createTargetUniquifier() { return new ThreadSafeTargetUniquifier(DEFAULT_THREAD_COUNT); } @ThreadSafe + ThreadSafeUniquifier<SkyKey> createSkyKeyUniquifier() { + return new ThreadSafeSkyKeyUniquifier(DEFAULT_THREAD_COUNT); + } + + @ThreadSafe @Override public void getTargetsMatchingPattern( QueryExpression owner, String pattern, Callback<Target> callback) @@ -735,6 +745,7 @@ } }; + @ThreadSafe public Map<SkyKey, Target> makeTargetsFromSkyKeys(Iterable<SkyKey> keys) throws InterruptedException { Multimap<SkyKey, SkyKey> packageKeyToTargetKeyMap = ArrayListMultimap.create(); @@ -822,7 +833,7 @@ * * <p>Note that there may not be nodes in the graph corresponding to the returned SkyKeys. */ - private Collection<SkyKey> getSkyKeysForFileFragments(Iterable<PathFragment> pathFragments) + Collection<SkyKey> getSkyKeysForFileFragments(Iterable<PathFragment> pathFragments) throws InterruptedException { Set<SkyKey> result = new HashSet<>(); Multimap<PathFragment, PathFragment> currentToOriginal = ArrayListMultimap.create(); @@ -894,12 +905,21 @@ } }; - private static Iterable<Target> getBuildFilesForPackageValues(Iterable<SkyValue> packageValues) { + static Iterable<Target> getBuildFilesForPackageValues(Iterable<SkyValue> packageValues) { return Iterables.transform( Iterables.filter(Iterables.transform(packageValues, EXTRACT_PACKAGE), ERROR_FREE_PACKAGE), GET_BUILD_FILE); } + @ThreadSafe + void getRBuildFilesParallel( + Collection<PathFragment> fileIdentifiers, + ThreadSafeCallback<Target> callback, + ForkJoinPool forkJoinPool) + throws QueryException, InterruptedException { + ParallelSkyQueryUtils.getRBuildFilesParallel(this, fileIdentifiers, callback, forkJoinPool); + } + /** * Calculates the set of {@link Package} objects, represented as source file targets, that depend * on the given list of BUILD files and subincludes (other files are filtered out). @@ -1045,6 +1065,18 @@ @ThreadSafe @Override + public void getAllRdepsUnboundedParallel( + QueryExpression expression, + VariableContext<Target> context, + ThreadSafeCallback<Target> callback, + ForkJoinPool forkJoinPool) + throws QueryException, InterruptedException { + ParallelSkyQueryUtils.getAllRdepsUnboundedParallel( + this, expression, context, callback, forkJoinPool); + } + + @ThreadSafe + @Override public void getAllRdeps( QueryExpression expression, Predicate<Target> universe,
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/AllRdepsFunction.java b/src/main/java/com/google/devtools/build/lib/query2/engine/AllRdepsFunction.java index faa9977..518b674 100644 --- a/src/main/java/com/google/devtools/build/lib/query2/engine/AllRdepsFunction.java +++ b/src/main/java/com/google/devtools/build/lib/query2/engine/AllRdepsFunction.java
@@ -112,6 +112,12 @@ List<Argument> args, ThreadSafeCallback<T> callback, ForkJoinPool forkJoinPool) throws QueryException, InterruptedException { - eval(env, context, expression, args, callback); + boolean unbounded = args.size() == 1; + if (unbounded && env instanceof StreamableQueryEnvironment<?>) { + ((StreamableQueryEnvironment<T>) env).getAllRdepsUnboundedParallel( + args.get(0).getExpression(), context, callback, forkJoinPool); + } else { + eval(env, context, expression, args, callback); + } } }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/Callback.java b/src/main/java/com/google/devtools/build/lib/query2/engine/Callback.java index bee90ad..0f43211 100644 --- a/src/main/java/com/google/devtools/build/lib/query2/engine/Callback.java +++ b/src/main/java/com/google/devtools/build/lib/query2/engine/Callback.java
@@ -27,5 +27,6 @@ * However, {@code QueryExpression}s calling the callback do not need to maintain this property, * as the {@code QueryEnvironment} should filter out duplicates. */ + @Override void process(Iterable<T> partialResult) throws QueryException, InterruptedException; }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java index 5f2b1ab..a31196ab 100644 --- a/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java +++ b/src/main/java/com/google/devtools/build/lib/query2/engine/FunctionExpression.java
@@ -23,6 +23,7 @@ import java.util.Collection; import java.util.List; +import java.util.concurrent.ForkJoinPool; /** * A query expression for user-defined query functions. @@ -52,6 +53,15 @@ } @Override + protected <T> void parEvalImpl( + QueryEnvironment<T> env, + VariableContext<T> context, + ThreadSafeCallback<T> callback, + ForkJoinPool forkJoinPool) throws QueryException, InterruptedException { + function.parEval(env, context, this, args, callback, forkJoinPool); + } + + @Override public void collectTargetPatterns(Collection<String> literals) { for (Argument arg : args) { if (arg.getType() == ArgumentType.EXPRESSION) {
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 7df7a00..4447488 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
@@ -14,7 +14,6 @@ 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; @@ -60,30 +59,6 @@ return callback.result; } - /** - * Notify {@code parentCallback} only about the events that match {@code retainIfTrue} predicate. - * - * @param parentCallback The parent callback to notify with the matching elements - * @param retainIfTrue A predicate that defines what elements to notify to the parent callback. - */ - public static <T> Callback<T> filteredCallback(final Callback<T> parentCallback, - final Predicate<T> retainIfTrue) { - return new Callback<T>() { - @Override - public void process(Iterable<T> partialResult) throws QueryException, InterruptedException { - Iterable<T> filter = Iterables.filter(partialResult, retainIfTrue); - if (!Iterables.isEmpty(filter)) { - parentCallback.process(filter); - } - } - - @Override - public String toString() { - return "filtered parentCallback of : " + retainIfTrue; - } - }; - } - /** A trivial {@link Uniquifier} base class. */ public abstract static class AbstractUniquifier<T, K> extends AbstractUniquifierBase<T, K> {
diff --git a/src/main/java/com/google/devtools/build/lib/query2/engine/RegexFilterExpression.java b/src/main/java/com/google/devtools/build/lib/query2/engine/RegexFilterExpression.java index 9b24014..fd76686 100644 --- a/src/main/java/com/google/devtools/build/lib/query2/engine/RegexFilterExpression.java +++ b/src/main/java/com/google/devtools/build/lib/query2/engine/RegexFilterExpression.java
@@ -47,6 +47,8 @@ + e.getMessage()); } + // Note that Patttern#matcher is thread-safe and so this Predicate can safely be used + // concurrently. final Predicate<T> matchFilter = new Predicate<T>() { @Override public boolean apply(T target) { @@ -62,7 +64,7 @@ env.eval( Iterables.getLast(args).getExpression(), context, - QueryUtil.filteredCallback(callback, matchFilter)); + filteredCallback(callback, matchFilter)); } @Override @@ -101,4 +103,50 @@ } protected abstract String getPattern(List<Argument> args); + + /** + * Returns a new {@link Callback} that forwards values that satisfies the given {@link Predicate} + * to the given {@code parentCallback}. + * + * <p>The returned {@link Callback} will be a {@link ThreadSafeCallback} iff + * {@code parentCallback} is as well. + */ + private static <T> Callback<T> filteredCallback( + final Callback<T> parentCallback, + final Predicate<T> retainIfTrue) { + return (parentCallback instanceof ThreadSafeCallback) + ? new ThreadSafeFilteredCallback<>((ThreadSafeCallback<T>) parentCallback, retainIfTrue) + : new FilteredCallback<>(parentCallback, retainIfTrue); + } + + private static class FilteredCallback<T> implements Callback<T> { + private final Callback<T> parentCallback; + private final Predicate<T> retainIfTrue; + + private FilteredCallback(Callback<T> parentCallback, Predicate<T> retainIfTrue) { + this.parentCallback = parentCallback; + this.retainIfTrue = retainIfTrue; + } + + @Override + public void process(Iterable<T> partialResult) throws QueryException, InterruptedException { + Iterable<T> filter = Iterables.filter(partialResult, retainIfTrue); + if (!Iterables.isEmpty(filter)) { + parentCallback.process(filter); + } + } + + @Override + public String toString() { + return "filtered parentCallback of : " + retainIfTrue; + } + } + + private static class ThreadSafeFilteredCallback<T> + extends FilteredCallback<T> implements ThreadSafeCallback<T> { + private ThreadSafeFilteredCallback( + ThreadSafeCallback<T> parentCallback, Predicate<T> retainIfTrue) { + super(parentCallback, retainIfTrue); + } + } }
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 db52c32..eda505a 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
@@ -14,6 +14,7 @@ package com.google.devtools.build.lib.query2.engine; import com.google.common.base.Predicate; +import java.util.concurrent.ForkJoinPool; /** * The environment of a Blaze query which supports predefined streaming operations. @@ -30,4 +31,15 @@ Callback<T> callback, int depth) throws QueryException, InterruptedException; + + /** + * Similar to {@link #getAllRdeps} but finds all rdeps without a depth bound, making use of the + * provided {@code forkJoinPool}. + */ + void getAllRdepsUnboundedParallel( + QueryExpression expression, + VariableContext<T> context, + ThreadSafeCallback<T> callback, + ForkJoinPool forkJoinPool) + throws QueryException, InterruptedException; }