add graceful restart mechanism for GetWorkStream to prevent DEADLINE_EXCEEDED; add long deadlines of 1hr to direct streams

Author: m-trieu

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSender.java

@@ -20,11 +20,16 @@
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+import javax.annotation.concurrent.GuardedBy;
 import javax.annotation.concurrent.ThreadSafe;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GetWorkRequest;
 import org.apache.beam.runners.dataflow.worker.windmill.WindmillConnection;
@@ -41,6 +46,8 @@
 import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.FixedStreamHeartbeatSender;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Owns and maintains a set of streams used to communicate with a specific Windmill worker.
@@ -57,15 +64,26 @@
 @Internal
 @ThreadSafe
 final class WindmillStreamSender implements GetWorkBudgetSpender, StreamSender {
-  private static final String STREAM_STARTER_THREAD_NAME = "StartWindmillStreamThread-%d";
-  private final AtomicBoolean started;
-  private final AtomicReference<GetWorkBudget> getWorkBudget;
-  private final GetWorkStream getWorkStream;
+  private static final Logger LOG = LoggerFactory.getLogger(WindmillStreamSender.class);
+  private static final String STREAM_MANAGER_THREAD_NAME_FORMAT = "WindmillStreamManagerThread-%d";
+  private static final int GET_WORK_STREAM_TTL_MINUTES = 45;
+
+  private final AtomicBoolean isRunning = new AtomicBoolean(false);
   private final GetDataStream getDataStream;
   private final CommitWorkStream commitWorkStream;
   private final WorkCommitter workCommitter;
   private final StreamingEngineThrottleTimers streamingEngineThrottleTimers;
   private final ExecutorService streamStarter;
+  private final String backendWorkerToken;
+
+  @GuardedBy("activeGetWorkStream")
+  private final AtomicReference<GetWorkStream> activeGetWorkStream;
+
+  @GuardedBy("activeGetWorkStream")
+  private final AtomicReference<GetWorkBudget> getWorkBudget;
+
+  @GuardedBy("activeGetWorkStream")
+  private final Supplier<GetWorkStream> getWorkStreamFactory;
 
   private WindmillStreamSender(
       WindmillConnection connection,
@@ -75,10 +93,9 @@ private WindmillStreamSender(
       WorkItemScheduler workItemScheduler,
       Function<GetDataStream, GetDataClient> getDataClientFactory,
       Function<CommitWorkStream, WorkCommitter> workCommitterFactory) {
-    this.started = new AtomicBoolean(false);
+    this.backendWorkerToken = connection.backendWorkerToken();
     this.getWorkBudget = getWorkBudget;
     this.streamingEngineThrottleTimers = StreamingEngineThrottleTimers.create();
-
     // Stream instances connect/reconnect internally, so we can reuse the same instance through the
     // entire lifecycle of WindmillStreamSender.
     this.getDataStream =
@@ -88,19 +105,21 @@ private WindmillStreamSender(
         streamingEngineStreamFactory.createDirectCommitWorkStream(
             connection, streamingEngineThrottleTimers.commitWorkThrottleTimer());
     this.workCommitter = workCommitterFactory.apply(commitWorkStream);
-    this.getWorkStream =
-        streamingEngineStreamFactory.createDirectGetWorkStream(
-            connection,
-            withRequestBudget(getWorkRequest, getWorkBudget.get()),
-            streamingEngineThrottleTimers.getWorkThrottleTimer(),
-            FixedStreamHeartbeatSender.create(getDataStream),
-            getDataClientFactory.apply(getDataStream),
-            workCommitter,
-            workItemScheduler);
+    this.getWorkStreamFactory =
+        () ->
+            streamingEngineStreamFactory.createDirectGetWorkStream(
+                connection,
+                withRequestBudget(getWorkRequest, getWorkBudget.get()),
+                streamingEngineThrottleTimers.getWorkThrottleTimer(),
+                FixedStreamHeartbeatSender.create(getDataStream),
+                getDataClientFactory.apply(getDataStream),
+                workCommitter,
+                workItemScheduler);
+    this.activeGetWorkStream = new AtomicReference<>();
     // 3 threads, 1 for each stream type (GetWork, GetData, CommitWork).
     this.streamStarter =
         Executors.newFixedThreadPool(
-            3, new ThreadFactoryBuilder().setNameFormat(STREAM_STARTER_THREAD_NAME).build());
+            3, new ThreadFactoryBuilder().setNameFormat(STREAM_MANAGER_THREAD_NAME_FORMAT).build());
   }
 
   static WindmillStreamSender create(
@@ -126,35 +145,45 @@ private static GetWorkRequest withRequestBudget(GetWorkRequest request, GetWorkB
   }
 
   synchronized void start() {
-    if (!started.get()) {
+    if (isRunning.compareAndSet(false, true)) {
       checkState(!streamStarter.isShutdown(), "WindmillStreamSender has already been shutdown.");
-
       // Start these 3 streams in parallel since they each may perform blocking IO.
+      CountDownLatch waitForInitialStream = new CountDownLatch(1);
+      streamStarter.execute(() -> getWorkStreamLoop(waitForInitialStream));
       CompletableFuture.allOf(
-              CompletableFuture.runAsync(getWorkStream::start, streamStarter),
               CompletableFuture.runAsync(getDataStream::start, streamStarter),
               CompletableFuture.runAsync(commitWorkStream::start, streamStarter))
           .join();
+      try {
+        waitForInitialStream.await();
+      } catch (InterruptedException e) {
+        close();
+        LOG.error("GetWorkStream to {} was never able to start.", backendWorkerToken);
+        throw new IllegalStateException("GetWorkStream unable to start aborting.", e);
+      }
       workCommitter.start();
-      started.set(true);
     }
   }
 
   @Override
   public synchronized void close() {
+    isRunning.set(false);
     streamStarter.shutdownNow();
-    getWorkStream.shutdown();
     getDataStream.shutdown();
     workCommitter.stop();
     commitWorkStream.shutdown();
   }
 
   @Override
   public void setBudget(long items, long bytes) {
-    GetWorkBudget budget = GetWorkBudget.builder().setItems(items).setBytes(bytes).build();
-    getWorkBudget.set(budget);
-    if (started.get()) {
-      getWorkStream.setBudget(budget);
+    synchronized (activeGetWorkStream) {
+      GetWorkBudget budget = GetWorkBudget.builder().setItems(items).setBytes(bytes).build();
+      getWorkBudget.set(budget);
+      if (isRunning.get()) {
+        GetWorkStream stream = activeGetWorkStream.get();
+        assert stream != null;
+        stream.setBudget(budget);
+      }
     }
   }
 
@@ -165,4 +194,39 @@ long getAndResetThrottleTime() {
   long getCurrentActiveCommitBytes() {
     return workCommitter.currentActiveCommitBytes();
   }
+
+  /**
+   * Creates, starts, and gracefully terminates {@link GetWorkStream} before the clientside deadline
+   * to prevent {@link org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Status#DEADLINE_EXCEEDED} errors.
+   * If at any point the server closes the stream, reconnects immediately.
+   */
+  private void getWorkStreamLoop(CountDownLatch waitForInitialStream) {
+    @Nullable GetWorkStream newStream = null;
+    while (isRunning.get()) {
+      synchronized (activeGetWorkStream) {
+        newStream = getWorkStreamFactory.get();
+        newStream.start();
+        waitForInitialStream.countDown();
+        activeGetWorkStream.set(newStream);
+      }
+      try {
+        // Try to gracefully terminate the stream.
+        if (!newStream.awaitTermination(GET_WORK_STREAM_TTL_MINUTES, TimeUnit.MINUTES)) {
+          newStream.halfClose();
+        }
+
+        // If graceful termination is unsuccessful, forcefully shutdown.
+        if (!newStream.awaitTermination(30, TimeUnit.SECONDS)) {
+          newStream.shutdown();
+        }
+
+      } catch (InterruptedException e) {
+        // continue until !isRunning.
+      }
+    }
+
+    if (newStream != null) {
+      newStream.shutdown();
+    }
+  }
 }

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

@@ -189,6 +189,12 @@ private static <T extends AbstractStub<T>> T withDefaultDeadline(T stub) {
     return stub.withDeadlineAfter(DEFAULT_STREAM_RPC_DEADLINE_SECONDS, TimeUnit.SECONDS);
   }
 
+  private static <T extends AbstractStub<T>> T withLongDeadline(T stub) {
+    // Deadlines are absolute points in time, so generate a new one everytime this function is
+    // called.
+    return stub.withDeadlineAfter(1, TimeUnit.HOURS);
+  }
+
   private static void printSummaryHtmlForWorker(
       String workerToken, Collection<AbstractWindmillStream<?, ?>> streams, PrintWriter writer) {
     writer.write(
@@ -230,7 +236,7 @@ public GetWorkStream createDirectGetWorkStream(
       WorkItemScheduler workItemScheduler) {
     return GrpcDirectGetWorkStream.create(
         connection.backendWorkerToken(),
-        responseObserver -> connection.stub().getWorkStream(responseObserver),
+        responseObserver -> withLongDeadline(connection.stub()).getWorkStream(responseObserver),
         request,
         grpcBackOff.get(),
         newStreamObserverFactory(),
@@ -265,7 +271,7 @@ public GetDataStream createDirectGetDataStream(
       WindmillConnection connection, ThrottleTimer getDataThrottleTimer) {
     return GrpcGetDataStream.create(
         connection.backendWorkerToken(),
-        responseObserver -> connection.stub().getDataStream(responseObserver),
+        responseObserver -> withLongDeadline(connection.stub()).getDataStream(responseObserver),
         grpcBackOff.get(),
         newStreamObserverFactory(),
         streamRegistry,
@@ -297,7 +303,7 @@ public CommitWorkStream createDirectCommitWorkStream(
       WindmillConnection connection, ThrottleTimer commitWorkThrottleTimer) {
     return GrpcCommitWorkStream.create(
         connection.backendWorkerToken(),
-        responseObserver -> connection.stub().commitWorkStream(responseObserver),
+        responseObserver -> withLongDeadline(connection.stub()).commitWorkStream(responseObserver),
         grpcBackOff.get(),
         newStreamObserverFactory(),
         streamRegistry,

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

@@ -44,7 +44,7 @@ final class DirectStreamObserver<T> implements TerminatingStreamObserver<T> {
   private static final long OUTPUT_CHANNEL_CONSIDERED_STALLED_SECONDS = 30;
 
   private final Phaser isReadyNotifier;
-  private final long deadlineSeconds;
+  private final long inactivityTimeout;
   private final int messagesBetweenIsReadyChecks;
   private final Object lock = new Object();
 
@@ -63,11 +63,11 @@ final class DirectStreamObserver<T> implements TerminatingStreamObserver<T> {
   DirectStreamObserver(
       Phaser isReadyNotifier,
       CallStreamObserver<T> outboundObserver,
-      long deadlineSeconds,
+      long inactivityTimeout,
       int messagesBetweenIsReadyChecks) {
     this.isReadyNotifier = isReadyNotifier;
     this.outboundObserver = outboundObserver;
-    this.deadlineSeconds = deadlineSeconds;
+    this.inactivityTimeout = inactivityTimeout;
     // We always let the first message pass through without blocking because it is performed under
     // the StreamPool synchronized block and single header message isn't going to cause memory
     // issues due to excessive buffering within grpc.
@@ -153,7 +153,7 @@ public void onNext(T value) {
         }
       } catch (TimeoutException e) {
         totalSecondsWaited += waitSeconds;
-        if (totalSecondsWaited > deadlineSeconds) {
+        if (totalSecondsWaited > inactivityTimeout) {
           String errorMessage = constructStreamCancelledErrorMessage(totalSecondsWaited);
           LOG.error(errorMessage);
           throw new WindmillRpcException(errorMessage, e);
@@ -217,11 +217,11 @@ public void terminate(Throwable terminationException) {
   }
 
   private String constructStreamCancelledErrorMessage(long totalSecondsWaited) {
-    return deadlineSeconds > 0
+    return inactivityTimeout > 0
         ? "Waited "
             + totalSecondsWaited
             + "s which exceeds given deadline of "
-            + deadlineSeconds
+            + inactivityTimeout
             + "s for the outboundObserver to become ready meaning "
             + "that the stream deadline was not respected."
         : "Output channel has been blocked for "

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

@@ -29,20 +29,20 @@
  */
 public abstract class StreamObserverFactory {
   public static StreamObserverFactory direct(
-      long deadlineSeconds, int messagesBetweenIsReadyChecks) {
-    return new Direct(deadlineSeconds, messagesBetweenIsReadyChecks);
+      long inactivityTimeout, int messagesBetweenIsReadyChecks) {
+    return new Direct(inactivityTimeout, messagesBetweenIsReadyChecks);
   }
 
   public abstract <ResponseT, RequestT> TerminatingStreamObserver<RequestT> from(
       Function<StreamObserver<ResponseT>, StreamObserver<RequestT>> clientFactory,
       StreamObserver<ResponseT> responseObserver);
 
   private static class Direct extends StreamObserverFactory {
-    private final long deadlineSeconds;
+    private final long inactivityTimeout;
     private final int messagesBetweenIsReadyChecks;
 
-    Direct(long deadlineSeconds, int messagesBetweenIsReadyChecks) {
-      this.deadlineSeconds = deadlineSeconds;
+    Direct(long inactivityTimeout, int messagesBetweenIsReadyChecks) {
+      this.inactivityTimeout = inactivityTimeout;
       this.messagesBetweenIsReadyChecks = messagesBetweenIsReadyChecks;
     }
 
@@ -57,7 +57,7 @@ public <ResponseT, RequestT> TerminatingStreamObserver<RequestT> from(
                   new ForwardingClientResponseObserver<ResponseT, RequestT>(
                       inboundObserver, phaser::arrive, phaser::forceTermination));
       return new DirectStreamObserver<>(
-          phaser, outboundObserver, deadlineSeconds, messagesBetweenIsReadyChecks);
+          phaser, outboundObserver, inactivityTimeout, messagesBetweenIsReadyChecks);
     }
   }
 }

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java

@@ -20,6 +20,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
@@ -226,7 +227,7 @@ public void testCloseAllStreams_closesAllStreams() {
     windmillStreamSender.start();
     windmillStreamSender.close();
 
-    verify(mockGetWorkStream).shutdown();
+    verify(mockGetWorkStream, atLeastOnce()).shutdown();
     verify(mockGetDataStream).shutdown();
     verify(mockCommitWorkStream).shutdown();
   }
@@ -268,7 +269,6 @@ public void testCloseAllStreams_doesNotStartStreamsAfterClose() {
     verify(mockGetDataStream, times(0)).start();
     verify(mockCommitWorkStream, times(0)).start();
 
-    verify(mockGetWorkStream).shutdown();
     verify(mockGetDataStream).shutdown();
     verify(mockCommitWorkStream).shutdown();
   }