add internal restart mechanism to avoid DEADLINE_EXCEEDED in windmill streams

Author: m-trieu

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

@@ -20,15 +20,11 @@
 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;
@@ -43,11 +39,8 @@
 import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudget;
 import org.apache.beam.runners.dataflow.worker.windmill.work.budget.GetWorkBudgetSpender;
 import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.FixedStreamHeartbeatSender;
-import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender;
 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.
@@ -64,42 +57,28 @@
 @Internal
 @ThreadSafe
 final class WindmillStreamSender implements GetWorkBudgetSpender, StreamSender {
-  private static final Logger LOG = LoggerFactory.getLogger(WindmillStreamSender.class);
-  private static final String STREAM_MANAGER_THREAD_NAME_FORMAT = "WindmillStreamManagerThread";
-
-  // Must be shorter than withLongDeadline() duration in GrpcWindmillStreamFactory to prevent
-  // DEADLINE_EXCEEDED.
-  private static final int GET_WORK_STREAM_TTL_MINUTES = 45;
-
-  private final AtomicBoolean isRunning = new AtomicBoolean(false);
+  private static final String STREAM_STARTER_THREAD_NAME = "StartWindmillStreamThread-%d";
+  private final AtomicBoolean started;
+  private final AtomicReference<GetWorkBudget> getWorkBudget;
+  private final GetWorkStream getWorkStream;
   private final GetDataStream getDataStream;
   private final CommitWorkStream commitWorkStream;
   private final WorkCommitter workCommitter;
   private final StreamingEngineThrottleTimers streamingEngineThrottleTimers;
-  private final ExecutorService streamManagerExecutor;
-  private final String backendWorkerToken;
-  private final Object getWorkStreamLock = new Object();
-
-  @GuardedBy("getWorkStreamLock")
-  private final Supplier<GetWorkStream> getWorkStreamFactory;
-
-  @GuardedBy("getWorkStreamLock")
-  private @Nullable GetWorkStream activeGetWorkStream;
-
-  @GuardedBy("getWorkStreamLock")
-  private GetWorkBudget getWorkBudget;
+  private final ExecutorService streamStarter;
 
   private WindmillStreamSender(
       WindmillConnection connection,
       GetWorkRequest getWorkRequest,
-      GetWorkBudget initialGetWorkBudget,
+      AtomicReference<GetWorkBudget> getWorkBudget,
       GrpcWindmillStreamFactory streamingEngineStreamFactory,
       WorkItemScheduler workItemScheduler,
       Function<GetDataStream, GetDataClient> getDataClientFactory,
       Function<CommitWorkStream, WorkCommitter> workCommitterFactory) {
-    this.backendWorkerToken = connection.backendWorkerToken();
-    this.getWorkBudget = initialGetWorkBudget;
+    this.started = new AtomicBoolean(false);
+    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 =
@@ -109,43 +88,33 @@ private WindmillStreamSender(
         streamingEngineStreamFactory.createDirectCommitWorkStream(
             connection, streamingEngineThrottleTimers.commitWorkThrottleTimer());
     this.workCommitter = workCommitterFactory.apply(commitWorkStream);
-    this.activeGetWorkStream = null;
-
-    HeartbeatSender heartbeatSender = FixedStreamHeartbeatSender.create(getDataStream);
-    GetDataClient getDataClient = getDataClientFactory.apply(getDataStream);
-    this.getWorkStreamFactory =
-        () -> {
-          synchronized (getWorkStreamLock) {
-            return streamingEngineStreamFactory.createDirectGetWorkStream(
-                connection,
-                withRequestBudget(getWorkRequest, getWorkBudget),
-                streamingEngineThrottleTimers.getWorkThrottleTimer(),
-                heartbeatSender,
-                getDataClient,
-                workCommitter,
-                workItemScheduler);
-          }
-        };
+    this.getWorkStream =
+        streamingEngineStreamFactory.createDirectGetWorkStream(
+            connection,
+            withRequestBudget(getWorkRequest, getWorkBudget.get()),
+            streamingEngineThrottleTimers.getWorkThrottleTimer(),
+            FixedStreamHeartbeatSender.create(getDataStream),
+            getDataClientFactory.apply(getDataStream),
+            workCommitter,
+            workItemScheduler);
     // 3 threads, 1 for each stream type (GetWork, GetData, CommitWork).
-    this.streamManagerExecutor =
+    this.streamStarter =
         Executors.newCachedThreadPool(
-            new ThreadFactoryBuilder()
-                .setNameFormat(STREAM_MANAGER_THREAD_NAME_FORMAT + "-" + backendWorkerToken + "-%d")
-                .build());
+            new ThreadFactoryBuilder().setNameFormat(STREAM_STARTER_THREAD_NAME).build());
   }
 
   static WindmillStreamSender create(
       WindmillConnection connection,
       GetWorkRequest getWorkRequest,
-      GetWorkBudget initialGetWorkBudget,
+      GetWorkBudget getWorkBudget,
       GrpcWindmillStreamFactory streamingEngineStreamFactory,
       WorkItemScheduler workItemScheduler,
       Function<GetDataStream, GetDataClient> getDataClientFactory,
       Function<CommitWorkStream, WorkCommitter> workCommitterFactory) {
     return new WindmillStreamSender(
         connection,
         getWorkRequest,
-        initialGetWorkBudget,
+        new AtomicReference<>(getWorkBudget),
         streamingEngineStreamFactory,
         workItemScheduler,
         getDataClientFactory,
@@ -157,31 +126,24 @@ private static GetWorkRequest withRequestBudget(GetWorkRequest request, GetWorkB
   }
 
   synchronized void start() {
-    if (isRunning.compareAndSet(false, true)) {
-      checkState(
-          !streamManagerExecutor.isShutdown(), "WindmillStreamSender has already been shutdown.");
+    if (!started.get()) {
+      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);
-      streamManagerExecutor.execute(() -> getWorkStreamLoop(waitForInitialStream::countDown));
       CompletableFuture.allOf(
-              CompletableFuture.runAsync(getDataStream::start, streamManagerExecutor),
-              CompletableFuture.runAsync(commitWorkStream::start, streamManagerExecutor))
+              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);
-    streamManagerExecutor.shutdownNow();
+    streamStarter.shutdownNow();
+    getWorkStream.shutdown();
     getDataStream.shutdown();
     workCommitter.stop();
     commitWorkStream.shutdown();
@@ -190,16 +152,9 @@ public synchronized void close() {
   @Override
   public void setBudget(long items, long bytes) {
     GetWorkBudget budget = GetWorkBudget.builder().setItems(items).setBytes(bytes).build();
-    synchronized (getWorkStreamLock) {
-      getWorkBudget = budget;
-      if (isRunning.get()) {
-        // activeGetWorkStream could be null if start() was called but activeGetWorkStream was not
-        // populated yet. Populating activeGetWorkStream and setting the budget are guaranteed to
-        // execute serially since both operations synchronize on getWorkStreamLock.
-        if (activeGetWorkStream != null) {
-          activeGetWorkStream.setBudget(budget);
-        }
-      }
+    getWorkBudget.set(budget);
+    if (started.get()) {
+      getWorkStream.setBudget(budget);
     }
   }
 
@@ -210,70 +165,4 @@ 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(Runnable onInitialStream) {
-    boolean shouldCreateNewStream = true;
-    while (isRunning.get() && shouldCreateNewStream) {
-      GetWorkStream newStream;
-      synchronized (getWorkStreamLock) {
-        newStream = getWorkStreamFactory.get();
-        newStream.start();
-        onInitialStream.run();
-        activeGetWorkStream = newStream;
-      }
-
-      shouldCreateNewStream = shouldCreateNewStreamAfterAwaitingTermination(newStream);
-    }
-  }
-
-  /**
-   * Manages stream termination. Returns true if a new stream should be created after the stream is
-   * terminated.
-   *
-   * @implNote This may block for up to {@link #GET_WORK_STREAM_TTL_MINUTES} minutes.
-   */
-  private boolean shouldCreateNewStreamAfterAwaitingTermination(GetWorkStream stream) {
-    try {
-      // Try to gracefully terminate the stream. If awaitTermination() returns before the TTL it
-      // means the server has terminated the connection and we reconnect immediately. If the stream
-      // is alive, terminate and drain the stream from the client to prevent DEADLINE_EXCEEDED
-      // status errors.
-      if (!stream.awaitTermination(GET_WORK_STREAM_TTL_MINUTES, TimeUnit.MINUTES)) {
-        drainStreamAsync(stream);
-      }
-    } catch (InterruptedException e) {
-      assert !isRunning.get();
-      return false;
-    }
-
-    return true;
-  }
-
-  private void drainStreamAsync(GetWorkStream stream) {
-    // Offload the old stream termination to a different thread which will terminate once it
-    // closes the stream. This allows this thread to not wait for the old stream to terminate
-    // before creating a new stream once the old stream has timed out.
-    streamManagerExecutor.execute(
-        () -> {
-          stream.halfClose();
-          try {
-            // Wait a bit for retries/drains then forcefully shutdown if graceful termination is
-            // unsuccessful.
-            if (!stream.awaitTermination(5, TimeUnit.MINUTES)) {
-              stream.shutdown();
-            }
-
-          } catch (InterruptedException e) {
-            assert !isRunning.get();
-          } finally {
-            // Make sure we clean up the stream.
-            stream.shutdown();
-          }
-        });
-  }
 }