Buffer blaze query

It was previously sending each label individually over gRPC, where each call
has a lot of overhead.

This makes queries with a large amount of output _a lot_ faster. For an example
query where all packages are already loaded, I observe a difference of ~3.5s
before this change to ~1.6s after this change.

PiperOrigin-RevId: 177426957

2 files changed, 21 insertions, 2 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
index 7f6c173b24..d13beac249 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
@@ -153,6 +153,16 @@ public abstract class OutputFormatter implements Serializable {
             : DependencyFilter.NO_IMPLICIT_DEPS);
   }
 
+  /**
+   * Workaround for a bug in {@link java.nio.channels.Channels#newChannel(OutputStream)}, which
+   * attempts to close the output stream on interrupt, which can cause a deadlock if there is an
+   * ongoing write. If this formatter uses Channels.newChannel, then it must return false here, and
+   * perform its own buffering.
+   */
+  public boolean canBeBuffered() {
+    return true;
+  }
+
   public void verifyCompatible(QueryEnvironment<?> env, QueryExpression expr)
       throws QueryException {
   }
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
index ec8ae85db4..8817c71ce8 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
@@ -45,6 +45,7 @@ import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.common.options.OptionsParser;
 import com.google.devtools.common.options.OptionsProvider;
+import java.io.BufferedOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.channels.ClosedByInterruptException;
@@ -156,7 +157,15 @@ public final class QueryCommand implements BlazeCommand {
 
     expr = queryEnv.transformParsedQuery(expr);
 
-    OutputStream out = env.getReporter().getOutErr().getOutputStream();
+    OutputStream out;
+    if (formatter.canBeBuffered()) {
+      // There is no particular reason for the 16384 constant here, except its a multiple of the
+      // gRPC buffer size. We mainly don't want to send each label individually because the output
+      // stream is connected to gRPC, and every write gets converted to one gRPC call.
+      out = new BufferedOutputStream(env.getReporter().getOutErr().getOutputStream(), 16384);
+    } else {
+      out = env.getReporter().getOutErr().getOutputStream();
+    }
     ThreadSafeOutputFormatterCallback<Target> callback;
     if (streamResults) {
       disableAnsiCharactersFiltering(env);
@@ -217,7 +226,7 @@ public final class QueryCommand implements BlazeCommand {
             result,
             targets,
             formatter,
-            env.getReporter().getOutErr().getOutputStream(),
+            out,
             queryOptions.aspectDeps.createResolver(env.getPackageManager(), env.getReporter()));
       } catch (ClosedByInterruptException | InterruptedException e) {
         env.getReporter().handle(Event.error("query interrupted"));