Render streamed proto output in parallel

Author: keithl-stripe

src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java

@@ -16,8 +16,12 @@
 import com.google.devtools.build.lib.packages.LabelPrinter;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
+import com.google.devtools.build.lib.query2.proto.proto2api.Build;
+
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.util.stream.StreamSupport;
 
 /**
  * An output formatter that outputs a protocol buffer representation of a query result and outputs
@@ -34,13 +38,70 @@ public String getName() {
   public OutputFormatterCallback<Target> createPostFactoStreamCallback(
       final OutputStream out, final QueryOptions options, LabelPrinter labelPrinter) {
     return new OutputFormatterCallback<Target>() {
+      private final LabelPrinter ourLabelPrinter = labelPrinter;
+
       @Override
       public void processOutput(Iterable<Target> partialResult)
           throws IOException, InterruptedException {
-        for (Target target : partialResult) {
-          toTargetProtoBuffer(target, labelPrinter).writeDelimitedTo(out);
+        try {
+          StreamSupport.stream(partialResult.spliterator(), /* parallel= */true)
+              .map(this::toProto)
+              .map(StreamedProtoOutputFormatter::writeDelimited)
+              .forEach(this::writeToOutputStreamThreadSafe);
+        } catch (WrappedIOException e) {
+          throw e.getCause();
+        } catch (WrappedInterruptedException e) {
+          throw e.getCause();
+        }
+      }
+
+      private Build.Target toProto(Target target) {
+        try {
+          return toTargetProtoBuffer(target, ourLabelPrinter);
+        } catch (InterruptedException e) {
+          throw new WrappedInterruptedException(e);
+        }
+      }
+
+      private synchronized void writeToOutputStreamThreadSafe(ByteArrayOutputStream bout) {
+        try {
+          bout.writeTo(out);
+        } catch (IOException e) {
+          throw new RuntimeException(e);
         }
       }
     };
   }
+
+  private static ByteArrayOutputStream writeDelimited(Build.Target targetProtoBuffer) {
+    try {
+      var bout = new ByteArrayOutputStream(targetProtoBuffer.getSerializedSize() + 10);
+      targetProtoBuffer.writeDelimitedTo(bout);
+      return bout;
+    } catch (IOException e) {
+      throw new WrappedIOException(e);
+    }
+  }
+
+  private static class WrappedIOException extends RuntimeException {
+    private WrappedIOException(IOException cause) {
+      super(cause);
+    }
+
+    @Override
+    public synchronized IOException getCause() {
+      return (IOException) super.getCause();
+    }
+  }
+
+  private static class WrappedInterruptedException extends RuntimeException {
+    private WrappedInterruptedException(InterruptedException cause) {
+      super(cause);
+    }
+
+    @Override
+    public synchronized InterruptedException getCause() {
+      return (InterruptedException) super.getCause();
+    }
+  }
 }