Don't force PrintStream on streaming output formatters

Usually an OutputStream will do. Forgo the extra layer of indirection and
stream directly to the output.

--
MOS_MIGRATED_REVID=134682243

5 files changed, 89 insertions, 58 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/GraphOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/GraphOutputFormatter.java
index 0d7c9406ff..869901f894 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/GraphOutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/GraphOutputFormatter.java
@@ -23,7 +23,7 @@ import com.google.devtools.build.lib.graph.LabelSerializer;
 import com.google.devtools.build.lib.graph.Node;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.query2.output.QueryOptions.OrderOutput;
-import java.io.PrintStream;
+import java.io.OutputStream;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -47,7 +47,7 @@ class GraphOutputFormatter extends OutputFormatter {
   }
 
   @Override
-  public void output(QueryOptions options, Digraph<Target> result, PrintStream out,
+  public void output(QueryOptions options, Digraph<Target> result, OutputStream out,
       AspectResolver aspectProvider) {
     this.graphNodeStringLimit = options.graphNodeStringLimit;
 
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 d8de116f88..38ae978894 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
@@ -38,6 +38,7 @@ import com.google.devtools.build.lib.syntax.Printer;
 import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.common.options.EnumConverter;
 import java.io.IOException;
+import java.io.OutputStream;
 import java.io.PrintStream;
 import java.io.Serializable;
 import java.util.ArrayList;
@@ -154,7 +155,7 @@ public abstract class OutputFormatter implements Serializable {
    * Format the result (a set of target nodes implicitly ordered according to
    * the graph maintained by the QueryEnvironment), and print it to "out".
    */
-  public abstract void output(QueryOptions options, Digraph<Target> result, PrintStream out,
+  public abstract void output(QueryOptions options, Digraph<Target> result, OutputStream out,
       AspectResolver aspectProvider)
           throws IOException, InterruptedException;
 
@@ -183,14 +184,14 @@ public abstract class OutputFormatter implements Serializable {
      * <p>Intended to be use for streaming out during evaluation of a query.
      */
     OutputFormatterCallback<Target> createStreamCallback(
-        PrintStream out, QueryOptions options, QueryEnvironment<?> env);
+        OutputStream out, QueryOptions options, QueryEnvironment<?> env);
 
     /**
      * Same as {@link #createStreamCallback}, but intended to be used for outputting the
      * already-computed result of a query.
      */
     OutputFormatterCallback<Target> createPostFactoStreamCallback(
-        PrintStream out, QueryOptions options);
+        OutputStream out, QueryOptions options);
   }
 
   /**
@@ -224,7 +225,7 @@ public abstract class OutputFormatter implements Serializable {
     public void output(
         QueryOptions options,
         Digraph<Target> result,
-        PrintStream out,
+        OutputStream out,
         AspectResolver aspectResolver) throws IOException, InterruptedException {
       setOptions(options, aspectResolver);
       OutputFormatterCallback.processAllTargets(
@@ -232,6 +233,20 @@ public abstract class OutputFormatter implements Serializable {
     }
   }
 
+  /** Abstract class supplying a {@link PrintStream} to implementations, flushing it on close. */
+  private abstract static class TextOutputFormatterCallback<T> extends OutputFormatterCallback<T> {
+    protected PrintStream printStream;
+
+    private TextOutputFormatterCallback(OutputStream out) {
+      this.printStream = new PrintStream(out);
+    }
+
+    @Override
+    public void close() throws IOException {
+      flushAndCheckError(printStream);
+    }
+  }
+
   /**
    * An output formatter that prints the labels of the resulting target set in
    * topological order, optionally with the target's kind.
@@ -251,18 +266,16 @@ public abstract class OutputFormatter implements Serializable {
 
     @Override
     public OutputFormatterCallback<Target> createPostFactoStreamCallback(
-        final PrintStream out, final QueryOptions options) {
-      return new OutputFormatterCallback<Target>() {
-
+        OutputStream out, final QueryOptions options) {
+      return new TextOutputFormatterCallback<Target>(out) {
         @Override
-        protected void processOutput(Iterable<Target> partialResult)
-            throws IOException, InterruptedException {
+        protected void processOutput(Iterable<Target> partialResult) {
           for (Target target : partialResult) {
             if (showKind) {
-              out.print(target.getTargetKind());
-              out.print(' ');
+              printStream.print(target.getTargetKind());
+              printStream.print(' ');
             }
-            out.printf(
+            printStream.printf(
                 "%s%s", target.getLabel().getDefaultCanonicalForm(), options.getLineTerminator());
           }
         }
@@ -271,7 +284,7 @@ public abstract class OutputFormatter implements Serializable {
 
     @Override
     public OutputFormatterCallback<Target> createStreamCallback(
-        PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+        OutputStream out, QueryOptions options, QueryEnvironment<?> env) {
       return createPostFactoStreamCallback(out, options);
     }
   }
@@ -301,13 +314,12 @@ public abstract class OutputFormatter implements Serializable {
 
     @Override
     public OutputFormatterCallback<Target> createPostFactoStreamCallback(
-        final PrintStream out, final QueryOptions options) {
-      return new OutputFormatterCallback<Target>() {
+        OutputStream out, final QueryOptions options) {
+      return new TextOutputFormatterCallback<Target>(out) {
         private final Set<String> packageNames = Sets.newTreeSet();
 
         @Override
-        protected void processOutput(Iterable<Target> partialResult)
-            throws IOException, InterruptedException {
+        protected void processOutput(Iterable<Target> partialResult) {
 
           for (Target target : partialResult) {
             packageNames.add(target.getLabel().getPackageName());
@@ -318,15 +330,16 @@ public abstract class OutputFormatter implements Serializable {
         public void close() throws IOException {
           final String lineTerm = options.getLineTerminator();
           for (String packageName : packageNames) {
-            out.printf("%s%s", packageName, lineTerm);
+            printStream.printf("%s%s", packageName, lineTerm);
           }
+          super.close();
         }
       };
     }
 
     @Override
     public OutputFormatterCallback<Target> createStreamCallback(
-        PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+        OutputStream out, QueryOptions options, QueryEnvironment<?> env) {
       return createPostFactoStreamCallback(out, options);
     }
   }
@@ -346,16 +359,15 @@ public abstract class OutputFormatter implements Serializable {
 
     @Override
     public OutputFormatterCallback<Target> createPostFactoStreamCallback(
-        final PrintStream out, final QueryOptions options) {
-      return new OutputFormatterCallback<Target>() {
+        OutputStream out, final QueryOptions options) {
+      return new TextOutputFormatterCallback<Target>(out) {
 
         @Override
-        protected void processOutput(Iterable<Target> partialResult)
-            throws IOException, InterruptedException {
+        protected void processOutput(Iterable<Target> partialResult) {
           final String lineTerm = options.getLineTerminator();
           for (Target target : partialResult) {
             Location location = target.getLocation();
-            out.print(
+            printStream.print(
                 location.print()
                     + ": "
                     + target.getTargetKind()
@@ -369,7 +381,7 @@ public abstract class OutputFormatter implements Serializable {
 
     @Override
     public OutputFormatterCallback<Target> createStreamCallback(
-        PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+        OutputStream out, QueryOptions options, QueryEnvironment<?> env) {
       return createPostFactoStreamCallback(out, options);
     }
   }
@@ -388,15 +400,15 @@ public abstract class OutputFormatter implements Serializable {
 
     @Override
     public OutputFormatterCallback<Target> createPostFactoStreamCallback(
-        final PrintStream out, final QueryOptions options) {
-      return new OutputFormatterCallback<Target>() {
+        OutputStream out, final QueryOptions options) {
+      return new TextOutputFormatterCallback<Target>(out) {
         private final Set<Label> printed = CompactHashSet.create();
 
-        private void outputRule(Rule rule, PrintStream out) {
+        private void outputRule(Rule rule, PrintStream printStream) {
           final String lineTerm = options.getLineTerminator();
-          out.printf("# %s%s", rule.getLocation(), lineTerm);
-          out.printf("%s(%s", rule.getRuleClass(), lineTerm);
-          out.printf("  name = \"%s\",%s", rule.getName(), lineTerm);
+          printStream.printf("# %s%s", rule.getLocation(), lineTerm);
+          printStream.printf("%s(%s", rule.getRuleClass(), lineTerm);
+          printStream.printf("  name = \"%s\",%s", rule.getName(), lineTerm);
 
           for (Attribute attr : rule.getAttributes()) {
             Pair<Iterable<Object>, AttributeValueSource> values =
@@ -408,11 +420,11 @@ public abstract class OutputFormatter implements Serializable {
               continue; // Don't print default values.
             }
             Object value = Iterables.getOnlyElement(values.first);
-            out.printf("  %s = ", attr.getPublicName());
+            printStream.printf("  %s = ", attr.getPublicName());
             if (value instanceof Label) {
               value = ((Label) value).getDefaultCanonicalForm();
             } else if (value instanceof License) {
-              List<String> licenseTypes = new ArrayList<String>();
+              List<String> licenseTypes = new ArrayList<>();
               for (License.LicenseType licenseType : ((License) value).getLicenseTypes()) {
                 licenseTypes.add(licenseType.toString().toLowerCase());
               }
@@ -424,22 +436,21 @@ public abstract class OutputFormatter implements Serializable {
             // It is *much* faster to write to a StringBuilder compared to the PrintStream object.
             StringBuilder builder = new StringBuilder();
             Printer.write(builder, value);
-            out.print(builder);
-            out.printf(",%s", lineTerm);
+            printStream.print(builder);
+            printStream.printf(",%s", lineTerm);
           }
-          out.printf(")\n%s", lineTerm);
+          printStream.printf(")\n%s", lineTerm);
         }
 
         @Override
-        protected void processOutput(Iterable<Target> partialResult)
-            throws IOException, InterruptedException {
+        protected void processOutput(Iterable<Target> partialResult) {
 
           for (Target target : partialResult) {
             Rule rule = target.getAssociatedRule();
             if (rule == null || printed.contains(rule.getLabel())) {
               continue;
             }
-            outputRule(rule, out);
+            outputRule(rule, printStream);
             printed.add(rule.getLabel());
           }
         }
@@ -448,7 +459,7 @@ public abstract class OutputFormatter implements Serializable {
 
     @Override
     public OutputFormatterCallback<Target> createStreamCallback(
-        PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+        OutputStream out, QueryOptions options, QueryEnvironment<?> env) {
       return createPostFactoStreamCallback(out, options);
     }
   }
@@ -507,15 +518,20 @@ public abstract class OutputFormatter implements Serializable {
     }
 
     @Override
-    public void output(QueryOptions options, Digraph<Target> result, PrintStream out,
-        AspectResolver aspectResolver) {
+    public void output(
+        QueryOptions options,
+        Digraph<Target> result,
+        OutputStream out,
+        AspectResolver aspectResolver)
+            throws IOException {
+      PrintStream printStream = new PrintStream(out);
       // getRoots() isn't defined for cyclic graphs, so in order to handle
       // cycles correctly, we need work on the strong component graph, as
       // cycles should be treated a "clump" of nodes all on the same rank.
       // Graphs may contain cycles because there are errors in BUILD files.
 
       List<RankAndLabel> outputToOrder =
-          options.orderOutput == OrderOutput.FULL ? new ArrayList<RankAndLabel>() : null;
+          options.orderOutput == OrderOutput.FULL ? new ArrayList<>() : null;
       Digraph<Set<Node<Target>>> scGraph = result.getStrongComponentGraph();
       Set<Node<Set<Node<Target>>>> rankNodes = scGraph.getRoots();
       Set<Node<Set<Node<Target>>>> seen = new HashSet<>();
@@ -525,7 +541,8 @@ public abstract class OutputFormatter implements Serializable {
         // Print out this rank:
         for (Node<Set<Node<Target>>> xScc : rankNodes) {
           for (Node<Target> x : xScc.getLabel()) {
-            outputToStreamOrSave(rank, x.getLabel().getLabel(), out, outputToOrder, lineTerm);
+            outputToStreamOrSave(
+                rank, x.getLabel().getLabel(), printStream, outputToOrder, lineTerm);
           }
         }
 
@@ -543,9 +560,11 @@ public abstract class OutputFormatter implements Serializable {
       if (outputToOrder != null) {
         Collections.sort(outputToOrder);
         for (RankAndLabel item : outputToOrder) {
-          out.printf("%s%s", item, lineTerm);
+          printStream.printf("%s%s", item, lineTerm);
         }
       }
+
+      flushAndCheckError(printStream);
     }
   }
 
@@ -568,8 +587,12 @@ public abstract class OutputFormatter implements Serializable {
     }
 
     @Override
-    public void output(QueryOptions options, Digraph<Target> result, PrintStream out,
-        AspectResolver aspectResolver) {
+    public void output(
+        QueryOptions options,
+        Digraph<Target> result,
+        OutputStream out,
+        AspectResolver aspectResolver)
+            throws IOException {
       // In order to handle cycles correctly, we need work on the strong
       // component graph, as cycles should be treated a "clump" of nodes all on
       // the same rank. Graphs may contain cycles because there are errors in BUILD files.
@@ -617,9 +640,11 @@ public abstract class OutputFormatter implements Serializable {
             });
       }
       final String lineTerm = options.getLineTerminator();
+      PrintStream printStream = new PrintStream(out);
       for (RankAndLabel item : output) {
-        out.printf("%s%s", item, lineTerm);
+        printStream.printf("%s%s", item, lineTerm);
       }
+      flushAndCheckError(printStream);
     }
   }
 
@@ -656,6 +681,12 @@ public abstract class OutputFormatter implements Serializable {
     return Pair.of(possibleAttributeValues, source);
   }
 
+  private static void flushAndCheckError(PrintStream printStream) throws IOException {
+    if (printStream.checkError()) {
+      throw new IOException("PrintStream encountered an error");
+    }
+  }
+
   /**
    * Returns the target location, eventually stripping out the workspace path to obtain a relative
    * target (stable across machines / workspaces).
diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java
index 2cc49f0c5b..8b388bbd61 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java
@@ -48,7 +48,7 @@ import com.google.devtools.build.lib.query2.proto.proto2api.Build.QueryResult.Bu
 import com.google.devtools.build.lib.query2.proto.proto2api.Build.SourceFile;
 import com.google.devtools.build.lib.syntax.Environment;
 import java.io.IOException;
-import java.io.PrintStream;
+import java.io.OutputStream;
 import java.util.Collection;
 import java.util.LinkedHashSet;
 import java.util.Map;
@@ -89,7 +89,7 @@ public class ProtoOutputFormatter extends AbstractUnorderedFormatter {
 
   @Override
   public OutputFormatterCallback<Target> createPostFactoStreamCallback(
-      final PrintStream out, final QueryOptions options) {
+      final OutputStream out, final QueryOptions options) {
     return new OutputFormatterCallback<Target>() {
 
       private Builder queryResult;
@@ -117,7 +117,7 @@ public class ProtoOutputFormatter extends AbstractUnorderedFormatter {
 
   @Override
   public OutputFormatterCallback<Target> createStreamCallback(
-      PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+      OutputStream out, QueryOptions options, QueryEnvironment<?> env) {
     return createPostFactoStreamCallback(out, options);
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/QueryOutputUtils.java b/src/main/java/com/google/devtools/build/lib/query2/output/QueryOutputUtils.java
index 9d4696b938..d28b96a8ec 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/QueryOutputUtils.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/QueryOutputUtils.java
@@ -20,7 +20,7 @@ import com.google.devtools.build.lib.query2.engine.QueryEvalResult;
 import com.google.devtools.build.lib.query2.output.OutputFormatter.StreamedFormatter;
 import com.google.devtools.build.lib.query2.output.QueryOptions.OrderOutput;
 import java.io.IOException;
-import java.io.PrintStream;
+import java.io.OutputStream;
 import java.util.Set;
 
 /** Static utility methods for outputting a query. */
@@ -34,7 +34,7 @@ public class QueryOutputUtils {
   }
 
   public static void output(QueryOptions queryOptions, QueryEvalResult result,
-      Set<Target> targetsResult, OutputFormatter formatter, PrintStream outputStream,
+      Set<Target> targetsResult, OutputFormatter formatter, OutputStream outputStream,
       AspectResolver aspectResolver)
       throws IOException, InterruptedException {
     /*
diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java
index b60250ebfe..5727723604 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java
@@ -33,7 +33,7 @@ import com.google.devtools.build.lib.query2.output.AspectResolver.BuildFileDepen
 import com.google.devtools.build.lib.query2.output.OutputFormatter.AbstractUnorderedFormatter;
 import com.google.devtools.build.lib.util.Pair;
 import java.io.IOException;
-import java.io.PrintStream;
+import java.io.OutputStream;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.Map;
@@ -62,13 +62,13 @@ class XmlOutputFormatter extends AbstractUnorderedFormatter {
 
   @Override
   public OutputFormatterCallback<Target> createStreamCallback(
-      PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+      OutputStream out, QueryOptions options, QueryEnvironment<?> env) {
     return createPostFactoStreamCallback(out, options);
   }
 
   @Override
   public OutputFormatterCallback<Target> createPostFactoStreamCallback(
-      final PrintStream out, final QueryOptions options) {
+      final OutputStream out, final QueryOptions options) {
     return new OutputFormatterCallback<Target>() {
 
       private Document doc;