From 93c5b31c65fb6fd84682654098a108a675ece9bc Mon Sep 17 00:00:00 2001
From: Martin Trieu <martinkt@google.com>
Date: Fri, 25 Oct 2024 10:30:41 -0700
Subject: [PATCH] address PR comments

---
 .../client/grpc/GrpcGetDataStream.java        | 42 ++++++++-----------
 1 file changed, 18 insertions(+), 24 deletions(-)

diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java
index cda246065ab9..a5d5b1882fd7 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java
@@ -144,6 +144,16 @@ static GrpcGetDataStream create(
         processHeartbeatResponses);
   }
 
+  private static WindmillStreamShutdownException shutdownException(QueuedBatch batch) {
+    return new WindmillStreamShutdownException(
+        "Stream was closed when attempting to send " + batch.requestsCount() + " requests.");
+  }
+
+  private static WindmillStreamShutdownException shutdownException(QueuedRequest request) {
+    return new WindmillStreamShutdownException(
+        "Cannot send request=[" + request + "] on closed stream.");
+  }
+
   @Override
   protected synchronized void onNewStream() {
     if (isShutdown()) {
@@ -349,34 +359,21 @@ private <ResponseT> ResponseT issueRequest(QueuedRequest request, ParseFn<Respon
         "Cannot send request=[" + request + "] on closed stream.");
   }
 
-  private void handleShutdown(QueuedRequest request, Throwable... causes) {
+  private void handleShutdown(QueuedRequest request, Throwable cause) {
     if (isShutdown()) {
-      WindmillStreamShutdownException shutdownException =
-          new WindmillStreamShutdownException(
-              "Cannot send request=[" + request + "] on closed stream.");
-
-      for (Throwable cause : causes) {
-        shutdownException.addSuppressed(cause);
-      }
-
+      WindmillStreamShutdownException shutdownException = shutdownException(request);
+      shutdownException.addSuppressed(cause);
       throw shutdownException;
     }
   }
 
-  private void handleShutdown(QueuedBatch batch) {
-    if (isShutdown()) {
-      throw new WindmillStreamShutdownException(
-          "Stream was closed when attempting to send " + batch.requestsCount() + " requests.");
-    }
-  }
-
   private void queueRequestAndWait(QueuedRequest request) throws InterruptedException {
     QueuedBatch batch;
     boolean responsibleForSend = false;
     @Nullable QueuedBatch prevBatch = null;
     synchronized (shutdownLock) {
       if (isShutdown()) {
-        handleShutdown(request);
+        throw shutdownException(request);
       }
 
       batch = batches.isEmpty() ? null : batches.getLast();
@@ -405,7 +402,7 @@ private void queueRequestAndWait(QueuedRequest request) throws InterruptedExcept
       // queue so that a subsequent batch will wait for its completion.
       synchronized (shutdownLock) {
         if (isShutdown()) {
-          handleShutdown(batch);
+          throw shutdownException(batch);
         }
 
         verify(batch == batches.peekFirst(), "GetDataStream request batch removed before send().");
@@ -423,7 +420,7 @@ void trySendBatch(QueuedBatch batch) {
       sendBatch(batch);
       synchronized (shutdownLock) {
         if (isShutdown()) {
-          handleShutdown(batch);
+          throw shutdownException(batch);
         }
 
         verify(
@@ -434,7 +431,6 @@ void trySendBatch(QueuedBatch batch) {
       // of the next batch (if one exists).
       batch.notifySent();
     } catch (Exception e) {
-      LOG.error("Error occurred sending batch.", e);
       // Free waiters if the send() failed.
       batch.notifyFailed();
       // Propagate the exception to the calling thread.
@@ -447,7 +443,6 @@ private void sendBatch(QueuedBatch batch) {
       return;
     }
 
-    StreamingGetDataRequest batchedRequest = batch.asGetDataRequest();
     synchronized (shutdownLock) {
       // Synchronization of pending inserts is necessary with send to ensure duplicates are not
       // sent on stream reconnect.
@@ -455,8 +450,7 @@ private void sendBatch(QueuedBatch batch) {
         // shutdown() clears pending, once the stream is shutdown, prevent values from being added
         // to it.
         if (isShutdown()) {
-          throw new WindmillStreamShutdownException(
-              "Stream was closed when attempting to send " + batch.requestsCount() + " requests.");
+          throw shutdownException(batch);
         }
 
         for (QueuedRequest request : batch.requestsReadOnly()) {
@@ -470,7 +464,7 @@ private void sendBatch(QueuedBatch batch) {
     }
 
     try {
-      send(batchedRequest);
+      send(batch.asGetDataRequest());
     } catch (IllegalStateException e) {
       // The stream broke before this call went through; onNewStream will retry the fetch.
       LOG.warn("GetData stream broke before call started.", e);