[Dataflow Streaming] Fix race in GetDataStream that could leave get data requests orphaned on stream errors after half-closing. (#36401)

This caused test flakiness of GetDataStreamTest.testRequestKeyedData_reconnectOnStreamErrorAfterHalfClose

Fixes #36347

Author: scwhittle

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java

@@ -91,7 +91,9 @@ final class GrpcGetDataStream
   @GuardedBy("this")
   private final Deque<QueuedBatch> batches;
 
-  private final Supplier<Integer> batchesDebugSizeSupplier;
+  // Size of the batches that may be read without synchronization.  If it is under synchronized
+  // block it is guaranteed to be correct.
+  private final Supplier<Integer> batchesSizeSupplier;
 
   private final AtomicLong idGenerator;
   private final JobHeader jobHeader;
@@ -133,7 +135,7 @@ private GrpcGetDataStream(
     // Otherwise the deque is accessed via batches which has a guardedby annotation.
     ConcurrentLinkedDeque<QueuedBatch> batches = new ConcurrentLinkedDeque<>();
     this.batches = batches;
-    this.batchesDebugSizeSupplier = batches::size;
+    this.batchesSizeSupplier = batches::size;
     this.sendKeyedGetDataRequests = sendKeyedGetDataRequests;
     this.processHeartbeatResponses = processHeartbeatResponses;
   }
@@ -224,7 +226,7 @@ public void onResponse(StreamingGetDataResponse chunk) {
 
     @Override
     public boolean hasPendingRequests() {
-      return !pending.isEmpty();
+      return !pending.isEmpty() || batchesSizeSupplier.get() > 0;
     }
 
     @Override
@@ -276,7 +278,9 @@ protected synchronized void onFlushPending(boolean isNewStream)
     while (!batches.isEmpty()) {
       QueuedBatch batch = checkNotNull(batches.peekFirst());
       verify(!batch.isEmpty());
-      if (!batch.isFinalized()) break;
+      if (!batch.isFinalized()) {
+        break;
+      }
       try {
         verify(
             batch == batches.pollFirst(),
@@ -419,7 +423,7 @@ protected synchronized void shutdownInternal() {
 
   @Override
   public void appendSpecificHtml(PrintWriter writer) {
-    int batches = batchesDebugSizeSupplier.get();
+    int batches = batchesSizeSupplier.get();
     if (batches > 0) {
       writer.format("GetDataStream: %d queued batches ", batches);
     } else {
@@ -516,7 +520,7 @@ private synchronized void trySendBatch(QueuedBatch batch) throws WindmillStreamS
     }
     final @Nullable GetDataPhysicalStreamHandler currentGetDataPhysicalStream =
         (GetDataPhysicalStreamHandler) currentPhysicalStream;
-    if (currentGetDataPhysicalStream == null) {
+    if (currentGetDataPhysicalStream == null || clientClosed) {
       // Leave the batch finalized but in the batches queue.  Finalized batches will be sent on a
       // new stream in onFlushPending.
       return;

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamTest.java

@@ -316,16 +316,26 @@ public void testRequestKeyedData_reconnectOnStreamErrorAfterHalfClose()
     assertNull(streamInfo.onDone.get());
 
     // Simulate an error on the grpc stream, this should trigger retrying the requests on a new
-    // stream
-    // which is half-closed.
+    // stream which is half-closed.
     streamInfo.responseObserver.onError(new IOException("test error"));
 
-    FakeWindmillGrpcService.GetDataStreamInfo streamInfo2 = waitForConnectionAndConsumeHeader();
-    Windmill.StreamingGetDataRequest request2 = streamInfo2.requests.take();
-    assertThat(request2.getRequestIdList()).containsExactly(1L);
-    assertEquals(keyedGetDataRequest, request2.getStateRequest(0).getRequests(0));
-    assertNull(streamInfo2.onDone.get());
     Windmill.KeyedGetDataResponse keyedGetDataResponse = createTestResponse(1);
+    FakeWindmillGrpcService.GetDataStreamInfo streamInfo2;
+    while (true) {
+      streamInfo2 = waitForConnectionAndConsumeHeader();
+      streamInfo2.onDone.get();
+      Windmill.StreamingGetDataRequest request2 = streamInfo2.requests.poll(5, TimeUnit.SECONDS);
+      if (request2 == null) {
+        // Client half-closed but didn't send the request, this can happen due to race but
+        // should recover by resending stream with requests.
+        streamInfo2.responseObserver.onCompleted();
+        continue;
+      }
+      assertThat(request2.getRequestIdList()).containsExactly(1L);
+      assertEquals(keyedGetDataRequest, request2.getStateRequest(0).getRequests(0));
+      break;
+    }
+
     streamInfo2.responseObserver.onNext(
         Windmill.StreamingGetDataResponse.newBuilder()
             .addRequestId(1)