From 6504a251b54e8aaab3b12aa7c89c1036c6c790bb Mon Sep 17 00:00:00 2001
From: Googler <leba@google.com>
Date: Tue, 8 Aug 2023 01:04:58 -0700
Subject: [PATCH] [aquery] only flush the relevant output channel.

Depending on the output type.

PiperOrigin-RevId: 554730946
Change-Id: I62c2fbc62dcf1e9022cc0cccad9d1bd244bea982
---
 .../v2/StreamedConsumingOutputHandler.java       | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

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
index d7e1d0955d1036..ca19d80e20b5bd 100644
--- 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.