[aquery] only flush the relevant output channel.

joeleba · copybara-github · commit 6504a251b54e · 2023-08-08T01:06:27.000-07:00
Depending on the output type.

PiperOrigin-RevId: 554730946
Change-Id: I62c2fbc62dcf1e9022cc0cccad9d1bd244bea982
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/StreamedConsumingOutputHandler.java b/src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/StreamedConsumingOutputHandler.java
@@ -158,9 +158,19 @@ void addTaskToQueue(Message message, int fieldNumber, String messageLabel) {
 
   @Override
   public void close() throws IOException {
-    outputStream.flush();
-    codedOutputStream.flush();
-    printStream.flush();
+    switch (outputType) {
+      case BINARY:
+        codedOutputStream.flush();
+        break;
+      case DELIMITED_BINARY:
+        outputStream.flush();
+        break;
+      case TEXT:
+        printStream.flush();
+        break;
+      default:
+        throw new IllegalStateException("Unknown outputType: " + outputType);
+    }
   }
 
   // Only runs on 1 single thread.
