fix retries of GetDataStream

Author: scwhittle

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

@@ -261,26 +261,8 @@ protected PhysicalStreamHandler newResponseHandler() {
   @Override
   protected synchronized void onNewStream() throws WindmillStreamShutdownException {
     trySend(StreamingGetDataRequest.newBuilder().setHeader(jobHeader).build());
-    while (!batches.isEmpty()) {
-      QueuedBatch batch = batches.peekFirst();
-      verify(!batch.isEmpty());
-      if (batch.isFinalized()) {
-        try {
-          ((GetDataPhysicalStreamHandler) currentPhysicalStream).sendBatch(batch);
-          verify(
-              batch == batches.pollFirst(),
-              "Sent GetDataStream request batch removed before send() was complete.");
-          // Notify all waiters with requests in this batch as well as the sender
-          // of the next batch (if one exists).
-          batch.notifySent();
-        } catch (Exception e) {
-          LOG.debug("Batch failed to send on new stream", e);
-          // Free waiters if the send() failed.
-          batch.notifyFailed();
-          throw e;
-        }
-      }
-    }
+    // If there are any requests in batches, they will either complete being sent by the queuing
+    // thread or will encounter an error due to their input stream cancelling and will retry.
   }
 
   private long uniqueId() {
@@ -500,11 +482,9 @@ private synchronized void trySendBatch(QueuedBatch batch) throws WindmillStreamS
 
     try {
       verify(batch == batches.peekFirst(), "GetDataStream request batch removed before send().");
+      verify(batch == batches.pollFirst());
       verify(!batch.isEmpty());
       ((GetDataPhysicalStreamHandler) currentPhysicalStream).sendBatch(batch);
-      verify(
-          batch == batches.pollFirst(),
-          "Sent GetDataStream request batch removed before send() was complete.");
       // Notify all waiters with requests in this batch as well as the sender
       // of the next batch (if one exists).
       batch.notifySent();

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

@@ -43,7 +43,11 @@ public class WindmillStreamPoolTest {
   private final ConcurrentHashMap<
           TestWindmillStream, WindmillStreamPool.StreamData<TestWindmillStream>>
       holds = new ConcurrentHashMap<>();
-  @Rule public transient Timeout globalTimeout = Timeout.seconds(600);
+
+  @Rule
+  public transient Timeout globalTimeout =
+      Timeout.builder().withTimeout(10, TimeUnit.MINUTES).withLookingForStuckThread(true).build();
+
   private List<WindmillStreamPool.@Nullable StreamData<TestWindmillStream>> streams;
 
   @Before

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

@@ -28,6 +28,7 @@
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
@@ -63,7 +64,11 @@ public class GrpcCommitWorkStreamTest {
 
   @Rule public final ErrorCollector errorCollector = new ErrorCollector();
   @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
-  @Rule public transient Timeout globalTimeout = Timeout.seconds(60);
+
+  @Rule
+  public transient Timeout globalTimeout =
+      Timeout.builder().withTimeout(10, TimeUnit.MINUTES).withLookingForStuckThread(true).build();
+
   private final FakeWindmillGrpcService fakeService = new FakeWindmillGrpcService(errorCollector);
   private ManagedChannel inProcessChannel;
   private Server inProcessServer;

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

@@ -81,7 +81,11 @@ public class GrpcDirectGetWorkStreamTest {
   private static final String FAKE_SERVER_NAME = "Fake server for GrpcDirectGetWorkStreamTest";
   @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
   private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
-  @Rule public transient Timeout globalTimeout = Timeout.seconds(600);
+
+  @Rule
+  public transient Timeout globalTimeout =
+      Timeout.builder().withTimeout(10, TimeUnit.MINUTES).withLookingForStuckThread(true).build();
+
   private ManagedChannel inProcessChannel;
   private GrpcDirectGetWorkStream stream;
 

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

@@ -30,6 +30,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
@@ -63,7 +64,11 @@ public class GrpcGetDataStreamTest {
 
   @Rule public final ErrorCollector errorCollector = new ErrorCollector();
   @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
-  @Rule public transient Timeout globalTimeout = Timeout.seconds(600);
+
+  @Rule
+  public transient Timeout globalTimeout =
+      Timeout.builder().withTimeout(10, TimeUnit.MINUTES).withLookingForStuckThread(true).build();
+
   private final FakeWindmillGrpcService fakeService = new FakeWindmillGrpcService(errorCollector);
   private ManagedChannel inProcessChannel;
   private Server inProcessServer;

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

@@ -117,7 +117,11 @@ public class GrpcWindmillServerTest {
   private final long clientId = 10L;
   private final Set<ManagedChannel> openedChannels = new HashSet<>();
   private final MutableHandlerRegistry serviceRegistry = new MutableHandlerRegistry();
-  @Rule public transient Timeout globalTimeout = Timeout.seconds(600);
+
+  @Rule
+  public transient Timeout globalTimeout =
+      Timeout.builder().withTimeout(10, TimeUnit.MINUTES).withLookingForStuckThread(true).build();
+
   @Rule public GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
   @Rule public ErrorCollector errorCollector = new ErrorCollector();
   private Server server;