add internal restart mechanism to avoid DEADLINE_EXCEEDED in windmill streams

m-trieu · m-trieu · commit 57ad0f161218 · 2025-04-08T12:29:21.000-07:00
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/AbstractWindmillStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/AbstractWindmillStream.java
@@ -25,6 +25,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 import javax.annotation.concurrent.GuardedBy;
@@ -76,6 +77,8 @@ public abstract class AbstractWindmillStream<RequestT, ResponseT> implements Win
 
   private final Logger logger;
   private final ExecutorService executor;
+  private final WindmillStreamTimeout streamTimeout;
+  private final ScheduledExecutorService restartExecutor;
   private final BackOff backoff;
   private final CountDownLatch finishLatch;
   private final Set<AbstractWindmillStream<?, ?>> streamRegistry;
@@ -101,13 +104,30 @@ protected AbstractWindmillStream(
       StreamObserverFactory streamObserverFactory,
       Set<AbstractWindmillStream<?, ?>> streamRegistry,
       int logEveryNStreamFailures,
-      String backendWorkerToken) {
+      String backendWorkerToken,
+      WindmillStreamTimeout streamTimeout) {
     this.backendWorkerToken = backendWorkerToken;
     this.executor =
         Executors.newSingleThreadExecutor(
             new ThreadFactoryBuilder()
                 .setDaemon(true)
-                .setNameFormat(createThreadName(debugStreamType, backendWorkerToken))
+                .setNameFormat(
+                    createThreadName(debugStreamType, backendWorkerToken, "WindmillStream"))
+                .build());
+    this.restartExecutor =
+        // Exceptions thrown on this thread will not crash the user worker process.
+        Executors.newSingleThreadScheduledExecutor(
+            new ThreadFactoryBuilder()
+                .setNameFormat(
+                    createThreadName(
+                        debugStreamType, backendWorkerToken, "WindmillStreamRestarter"))
+                // We need to explicitly log the error since it may be swallowed.
+                .setUncaughtExceptionHandler(
+                    (t, e) ->
+                        logger.error(
+                            "{} failed due to uncaught exception during execution. ",
+                            t.getName(),
+                            e))
                 .build());
     this.backoff = backoff;
     this.streamRegistry = streamRegistry;
@@ -126,12 +146,14 @@ protected AbstractWindmillStream(
             logger);
     this.sleeper = Sleeper.DEFAULT;
     this.debugMetrics = StreamDebugMetrics.create();
+    this.streamTimeout = streamTimeout;
   }
 
-  private static String createThreadName(String streamType, String backendWorkerToken) {
+  private static String createThreadName(
+      String threadType, String streamType, String backendWorkerToken) {
     return !backendWorkerToken.isEmpty()
-        ? String.format("%s-%s-WindmillStream-thread", streamType, backendWorkerToken)
-        : String.format("%s-WindmillStream-thread", streamType);
+        ? String.format("%s-%s-%s-thread", streamType, backendWorkerToken, threadType)
+        : String.format("%s-%s-thread", streamType, threadType);
   }
 
   /** Called on each response from the server. */
@@ -177,6 +199,9 @@ public final void start() {
 
     if (shouldStartStream) {
       startStream();
+      // Restart the stream at every streamTimeout interval.
+      restartExecutor.scheduleAtFixedRate(
+          this::restart, streamTimeout.time(), streamTimeout.time(), streamTimeout.unit());
     }
   }
 
@@ -296,6 +321,21 @@ public final synchronized void halfClose() {
     // Synchronization of close and onCompleted necessary for correct retry logic in onNewStream.
     debugMetrics.recordHalfClose();
     clientClosed = true;
+    resetCurrentObserver();
+  }
+
+  /**
+   * Internally restart the stream to avoid {@link Status#DEADLINE_EXCEEDED} errors.
+   *
+   * @implNote Similar behavior to {@link #halfClose()}, except we allow callers to interact with
+   *     the stream after restarts.
+   */
+  private synchronized void restart() {
+    debugMetrics.recordRestartReason("Internal Timeout");
+    resetCurrentObserver();
+  }
+
+  private synchronized void resetCurrentObserver() {
     try {
       requestObserver.onCompleted();
     } catch (ResettableThrowingStreamObserver.StreamClosedException e) {
@@ -350,6 +390,8 @@ private synchronized boolean maybeTearDownStream() {
     return false;
   }
 
+  private void scheduleRestarts() {}
+
   private class ResponseObserver implements StreamObserver<ResponseT> {
 
     @Override
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStreamTimeout.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStreamTimeout.java
@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.client;
+
+import com.google.auto.value.AutoValue;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.annotations.Internal;
+
+/**
+ * Timeout for a {@link WindmillStream}. This should be less than a stream's deadline since we use
+ * this to internally restart the stream before we will get a {@link
+ * org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Status#DEADLINE_EXCEEDED} errors.
+ */
+@Internal
+@AutoValue
+public abstract class WindmillStreamTimeout {
+  public static WindmillStreamTimeout create(int time, TimeUnit unit) {
+    return new AutoValue_WindmillStreamTimeout(time, unit);
+  }
+
+  public abstract int time();
+
+  public abstract TimeUnit unit();
+
+  long asMillis() {
+    return unit().toMillis(time());
+  }
+}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java
@@ -41,6 +41,7 @@
 import org.apache.beam.runners.dataflow.worker.windmill.client.AbstractWindmillStream;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.CommitWorkStream;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamShutdownException;
+import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamTimeout;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.observers.StreamObserverFactory;
 import org.apache.beam.runners.dataflow.worker.windmill.client.throttling.ThrottleTimer;
 import org.apache.beam.sdk.util.BackOff;
@@ -75,7 +76,8 @@ private GrpcCommitWorkStream(
       ThrottleTimer commitWorkThrottleTimer,
       JobHeader jobHeader,
       AtomicLong idGenerator,
-      int streamingRpcBatchLimit) {
+      int streamingRpcBatchLimit,
+      WindmillStreamTimeout streamTimeout) {
     super(
         LOG,
         "CommitWorkStream",
@@ -84,7 +86,8 @@ private GrpcCommitWorkStream(
         streamObserverFactory,
         streamRegistry,
         logEveryNStreamFailures,
-        backendWorkerToken);
+        backendWorkerToken,
+        streamTimeout);
     pending = new ConcurrentHashMap<>();
     this.idGenerator = idGenerator;
     this.jobHeader = jobHeader;
@@ -103,7 +106,8 @@ static GrpcCommitWorkStream create(
       ThrottleTimer commitWorkThrottleTimer,
       JobHeader jobHeader,
       AtomicLong idGenerator,
-      int streamingRpcBatchLimit) {
+      int streamingRpcBatchLimit,
+      WindmillStreamTimeout streamTimeout) {
     return new GrpcCommitWorkStream(
         backendWorkerToken,
         startCommitWorkRpcFn,
@@ -114,7 +118,8 @@ static GrpcCommitWorkStream create(
         commitWorkThrottleTimer,
         jobHeader,
         idGenerator,
-        streamingRpcBatchLimit);
+        streamingRpcBatchLimit,
+        streamTimeout);
   }
 
   @Override
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java
@@ -35,6 +35,7 @@
 import org.apache.beam.runners.dataflow.worker.windmill.client.AbstractWindmillStream;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetWorkStream;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamShutdownException;
+import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamTimeout;
 import org.apache.beam.runners.dataflow.worker.windmill.client.commits.WorkCommitter;
 import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.GetDataClient;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GetWorkResponseChunkAssembler.AssembledWorkItem;
@@ -107,7 +108,8 @@ private GrpcDirectGetWorkStream(
       HeartbeatSender heartbeatSender,
       GetDataClient getDataClient,
       WorkCommitter workCommitter,
-      WorkItemScheduler workItemScheduler) {
+      WorkItemScheduler workItemScheduler,
+      WindmillStreamTimeout streamTimeout) {
     super(
         LOG,
         "GetWorkStream",
@@ -116,7 +118,8 @@ private GrpcDirectGetWorkStream(
         streamObserverFactory,
         streamRegistry,
         logEveryNStreamFailures,
-        backendWorkerToken);
+        backendWorkerToken,
+        streamTimeout);
     this.requestHeader = requestHeader;
     this.getWorkThrottleTimer = getWorkThrottleTimer;
     this.workItemScheduler = workItemScheduler;
@@ -150,7 +153,8 @@ static GrpcDirectGetWorkStream create(
       HeartbeatSender heartbeatSender,
       GetDataClient getDataClient,
       WorkCommitter workCommitter,
-      WorkItemScheduler workItemScheduler) {
+      WorkItemScheduler workItemScheduler,
+      WindmillStreamTimeout streamTimeout) {
     return new GrpcDirectGetWorkStream(
         backendWorkerToken,
         startGetWorkRpcFn,
@@ -164,7 +168,8 @@ static GrpcDirectGetWorkStream create(
         heartbeatSender,
         getDataClient,
         workCommitter,
-        workItemScheduler);
+        workItemScheduler,
+        streamTimeout);
   }
 
   private static Watermarks createWatermarks(
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
@@ -51,6 +51,7 @@
 import org.apache.beam.runners.dataflow.worker.windmill.client.AbstractWindmillStream;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetDataStream;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamShutdownException;
+import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamTimeout;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcGetDataStreamRequests.QueuedBatch;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcGetDataStreamRequests.QueuedRequest;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.observers.StreamObserverFactory;
@@ -95,7 +96,8 @@ private GrpcGetDataStream(
       AtomicLong idGenerator,
       int streamingRpcBatchLimit,
       boolean sendKeyedGetDataRequests,
-      Consumer<List<Windmill.ComputationHeartbeatResponse>> processHeartbeatResponses) {
+      Consumer<List<Windmill.ComputationHeartbeatResponse>> processHeartbeatResponses,
+      WindmillStreamTimeout streamTimeout) {
     super(
         LOG,
         "GetDataStream",
@@ -104,7 +106,8 @@ private GrpcGetDataStream(
         streamObserverFactory,
         streamRegistry,
         logEveryNStreamFailures,
-        backendWorkerToken);
+        backendWorkerToken,
+        streamTimeout);
     this.idGenerator = idGenerator;
     this.getDataThrottleTimer = getDataThrottleTimer;
     this.jobHeader = jobHeader;
@@ -128,7 +131,8 @@ static GrpcGetDataStream create(
       AtomicLong idGenerator,
       int streamingRpcBatchLimit,
       boolean sendKeyedGetDataRequests,
-      Consumer<List<Windmill.ComputationHeartbeatResponse>> processHeartbeatResponses) {
+      Consumer<List<Windmill.ComputationHeartbeatResponse>> processHeartbeatResponses,
+      WindmillStreamTimeout streamTimeout) {
     return new GrpcGetDataStream(
         backendWorkerToken,
         startGetDataRpcFn,
@@ -141,7 +145,8 @@ static GrpcGetDataStream create(
         idGenerator,
         streamingRpcBatchLimit,
         sendKeyedGetDataRequests,
-        processHeartbeatResponses);
+        processHeartbeatResponses,
+        streamTimeout);
   }
 
   private static WindmillStreamShutdownException shutdownExceptionFor(QueuedBatch batch) {
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java
@@ -30,6 +30,7 @@
 import org.apache.beam.runners.dataflow.worker.windmill.client.AbstractWindmillStream;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetWorkStream;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamShutdownException;
+import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamTimeout;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GetWorkResponseChunkAssembler.AssembledWorkItem;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.observers.StreamObserverFactory;
 import org.apache.beam.runners.dataflow.worker.windmill.client.throttling.ThrottleTimer;
@@ -72,7 +73,8 @@ private GrpcGetWorkStream(
       int logEveryNStreamFailures,
       boolean requestBatchedGetWorkResponse,
       ThrottleTimer getWorkThrottleTimer,
-      WorkItemReceiver receiver) {
+      WorkItemReceiver receiver,
+      WindmillStreamTimeout streamTimeout) {
     super(
         LOG,
         "GetWorkStream",
@@ -81,7 +83,8 @@ private GrpcGetWorkStream(
         streamObserverFactory,
         streamRegistry,
         logEveryNStreamFailures,
-        backendWorkerToken);
+        backendWorkerToken,
+        streamTimeout);
     this.request = request;
     this.getWorkThrottleTimer = getWorkThrottleTimer;
     this.receiver = receiver;
@@ -104,7 +107,8 @@ public static GrpcGetWorkStream create(
       int logEveryNStreamFailures,
       boolean requestBatchedGetWorkResponse,
       ThrottleTimer getWorkThrottleTimer,
-      WorkItemReceiver receiver) {
+      WorkItemReceiver receiver,
+      WindmillStreamTimeout streamTimeout) {
     return new GrpcGetWorkStream(
         backendWorkerToken,
         startGetWorkRpcFn,
@@ -115,7 +119,8 @@ public static GrpcGetWorkStream create(
         logEveryNStreamFailures,
         requestBatchedGetWorkResponse,
         getWorkThrottleTimer,
-        receiver);
+        receiver,
+        streamTimeout);
   }
 
   private void sendRequestExtension(long moreItems, long moreBytes) {
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkerMetadataStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkerMetadataStream.java
@@ -30,6 +30,7 @@
 import org.apache.beam.runners.dataflow.worker.windmill.client.AbstractWindmillStream;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetWorkerMetadataStream;
 import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamShutdownException;
+import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamTimeout;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.observers.StreamObserverFactory;
 import org.apache.beam.runners.dataflow.worker.windmill.client.throttling.ThrottleTimer;
 import org.apache.beam.sdk.util.BackOff;
@@ -60,7 +61,8 @@ private GrpcGetWorkerMetadataStream(
       int logEveryNStreamFailures,
       JobHeader jobHeader,
       ThrottleTimer getWorkerMetadataThrottleTimer,
-      Consumer<WindmillEndpoints> serverMappingConsumer) {
+      Consumer<WindmillEndpoints> serverMappingConsumer,
+      WindmillStreamTimeout streamTimeout) {
     super(
         LOG,
         "GetWorkerMetadataStream",
@@ -69,7 +71,8 @@ private GrpcGetWorkerMetadataStream(
         streamObserverFactory,
         streamRegistry,
         logEveryNStreamFailures,
-        "");
+        "",
+        streamTimeout);
     this.workerMetadataRequest = WorkerMetadataRequest.newBuilder().setHeader(jobHeader).build();
     this.getWorkerMetadataThrottleTimer = getWorkerMetadataThrottleTimer;
     this.serverMappingConsumer = serverMappingConsumer;
@@ -86,7 +89,8 @@ public static GrpcGetWorkerMetadataStream create(
       int logEveryNStreamFailures,
       JobHeader jobHeader,
       ThrottleTimer getWorkerMetadataThrottleTimer,
-      Consumer<WindmillEndpoints> serverMappingUpdater) {
+      Consumer<WindmillEndpoints> serverMappingUpdater,
+      WindmillStreamTimeout streamTimeout) {
     return new GrpcGetWorkerMetadataStream(
         startGetWorkerMetadataRpcFn,
         backoff,
@@ -95,7 +99,8 @@ public static GrpcGetWorkerMetadataStream create(
         logEveryNStreamFailures,
         jobHeader,
         getWorkerMetadataThrottleTimer,
-        serverMappingUpdater);
+        serverMappingUpdater,
+        streamTimeout);
   }
 
   /**
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillStreamFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillStreamFactory.java
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/AbstractWindmillStreamTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/AbstractWindmillStreamTest.java
