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 6d6dcd569e85..044792148c94 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 @@ -91,7 +91,9 @@ final class GrpcGetDataStream @GuardedBy("this") private final Deque batches; - private final Supplier 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 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 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; diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamTest.java index 4f584022c8a5..849b2612cecf 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStreamTest.java +++ b/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)