[Dataflow Java Streaming] Add a timeout to how long commits will retry to the service. (#39085)

* [Dataflow Java Streaming] Add a timeout to how long commits will retry to the service. In cases where the service is unavailable this prevents build-up of commits on workers that are no longer relevant. This defaults to 30 minutes but can be disabled via an experiment or by setting the option.

Author: scwhittle

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowStreamingPipelineOptions.java

@@ -196,6 +196,13 @@ public interface DataflowStreamingPipelineOptions extends PipelineOptions {
 
   void setStuckCommitDurationMillis(int value);
 
+  @Description(
+      "Retry commits on stream errors until this much time has elapsed since the commit was scheduled. If zero, retry forever.")
+  @Default.InstanceFactory(CommitWorkStreamRetryTimeoutMillisFactory.class)
+  long getCommitWorkStreamRetryTimeoutMillis();
+
+  void setCommitWorkStreamRetryTimeoutMillis(long value);
+
   @Description(
       "Period for sending 'global get config' requests to the service. The duration is "
           + "specified as seconds in 'PTx.yS' format, e.g. 'PT5.125S'."
@@ -343,4 +350,14 @@ public Boolean create(PipelineOptions options) {
       return ExperimentalOptions.hasExperiment(options, "enable_windmill_service_direct_path");
     }
   }
+
+  class CommitWorkStreamRetryTimeoutMillisFactory implements DefaultValueFactory<Long> {
+    @Override
+    public Long create(PipelineOptions options) {
+      if (ExperimentalOptions.hasExperiment(options, "disable_commit_retry_timeout")) {
+        return 0L;
+      }
+      return Duration.standardMinutes(30).getMillis();
+    }
+  }
 }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

@@ -756,6 +756,8 @@ public static StreamingDataflowWorker fromOptions(DataflowWorkerHarnessOptions o
                   new WorkHeartbeatResponseProcessor(computationStateCache::get))
               .setHealthCheckIntervalMillis(
                   options.getWindmillServiceStreamingRpcHealthCheckPeriodMs())
+              .setCommitWorkStreamRetryTimeout(
+                  java.time.Duration.ofMillis(options.getCommitWorkStreamRetryTimeoutMillis()))
               .build();
       return ConfigFetcherComputationStateCacheAndWindmillClient.builder()
           .setWindmillDispatcherClient(dispatcherClient)

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

@@ -20,7 +20,6 @@
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
-import com.google.auto.value.AutoValue;
 import java.io.PrintWriter;
 import java.time.Duration;
 import java.util.HashMap;
@@ -77,6 +76,7 @@ private static class StreamAndRequest {
   private final JobHeader jobHeader;
   private final int streamingRpcBatchLimit;
   private volatile boolean logMissingResponse = true;
+  private final Duration maxRetryDuration;
 
   private GrpcCommitWorkStream(
       String backendWorkerToken,
@@ -90,6 +90,7 @@ private GrpcCommitWorkStream(
       AtomicLong idGenerator,
       int streamingRpcBatchLimit,
       Duration halfClosePhysicalStreamAfter,
+      Duration maxRetryDuration,
       ScheduledExecutorService executor) {
     super(
         LOG,
@@ -104,6 +105,7 @@ private GrpcCommitWorkStream(
     this.idGenerator = idGenerator;
     this.jobHeader = jobHeader;
     this.streamingRpcBatchLimit = streamingRpcBatchLimit;
+    this.maxRetryDuration = maxRetryDuration;
   }
 
   static GrpcCommitWorkStream create(
@@ -118,6 +120,7 @@ static GrpcCommitWorkStream create(
       AtomicLong idGenerator,
       int streamingRpcBatchLimit,
       Duration halfClosePhysicalStreamAfter,
+      Duration maxRetryDuration,
       ScheduledExecutorService executor) {
     return new GrpcCommitWorkStream(
         backendWorkerToken,
@@ -130,6 +133,7 @@ static GrpcCommitWorkStream create(
         idGenerator,
         streamingRpcBatchLimit,
         halfClosePhysicalStreamAfter,
+        maxRetryDuration,
         executor);
   }
 
@@ -224,14 +228,48 @@ public void onResponse(StreamingCommitResponse response) {
       failureHandler.throwIfNonEmpty();
     }
 
-    @Override
     @SuppressWarnings("ReferenceEquality")
+    private boolean belongsToThisHandler(StreamAndRequest streamAndRequest) {
+      return streamAndRequest.handler == this;
+    }
+
+    @Override
     public boolean hasPendingRequests() {
-      return pending.entrySet().stream().anyMatch(e -> e.getValue().handler == this);
+      return pending.entrySet().stream().anyMatch(e -> belongsToThisHandler(e.getValue()));
     }
 
     @Override
+    @SuppressWarnings("ReferenceEquality")
     public void onDone(Status status) {
+      if (maxRetryDuration.compareTo(Duration.ZERO) > 0) {
+        // Remove the requests that have exceeded the retry time so they are not retried.
+        long nowNanos = System.nanoTime();
+        long maxRetryDurationNanos = maxRetryDuration.toNanos();
+        Iterator<Map.Entry<Long, StreamAndRequest>> iterator = pending.entrySet().iterator();
+        int keptRequests = 0, removedRequests = 0;
+        while (iterator.hasNext()) {
+          StreamAndRequest streamAndRequest = checkNotNull(iterator.next().getValue());
+          PendingRequest pendingRequest = streamAndRequest.request;
+          if (!belongsToThisHandler(streamAndRequest)
+              || nowNanos - pendingRequest.getStartTimeNanos() < maxRetryDurationNanos) {
+            ++keptRequests;
+            continue;
+          }
+          ++removedRequests;
+          iterator.remove();
+          try {
+            pendingRequest.completeWithStatus(CommitStatus.ABORTED);
+          } catch (RuntimeException e) {
+            LOG.warn("Exception while aborting commit due to retry timeout.", e);
+          }
+        }
+        if (removedRequests > 0) {
+          LOG.info(
+              "Aborting {} commits which have exceeded retry deadline, kept {}. Work will be retried as needed by service.",
+              removedRequests,
+              keptRequests);
+        }
+      }
       if (status.isOk() && hasPendingRequests()) {
         LOG.warn("Unexpected requests without responses on drained physical stream, retrying.");
       }
@@ -270,7 +308,7 @@ private void flushInternal(Map<Long, PendingRequest> requests)
 
     if (requests.size() == 1) {
       Map.Entry<Long, PendingRequest> elem = requests.entrySet().iterator().next();
-      if (elem.getValue().request().getSerializedSize()
+      if (elem.getValue().getRequest().getSerializedSize()
           > AbstractWindmillStream.RPC_STREAM_CHUNK_SIZE) {
         issueMultiChunkRequest(elem.getKey(), elem.getValue());
       } else {
@@ -286,7 +324,7 @@ private void issueSingleRequest(long id, PendingRequest pendingRequest)
     StreamingCommitWorkRequest.Builder requestBuilder = StreamingCommitWorkRequest.newBuilder();
     requestBuilder
         .addCommitChunkBuilder()
-        .setComputationId(pendingRequest.computationId())
+        .setComputationId(pendingRequest.getComputationId())
         .setRequestId(id)
         .setShardingKey(pendingRequest.shardingKey())
         .setSerializedWorkItemCommit(pendingRequest.serializedCommit());
@@ -311,9 +349,9 @@ private void issueBatchedRequest(Map<Long, PendingRequest> requests)
     for (Map.Entry<Long, PendingRequest> entry : requests.entrySet()) {
       PendingRequest request = entry.getValue();
       StreamingCommitRequestChunk.Builder chunkBuilder = requestBuilder.addCommitChunkBuilder();
-      if (lastComputation == null || !lastComputation.equals(request.computationId())) {
-        chunkBuilder.setComputationId(request.computationId());
-        lastComputation = request.computationId();
+      if (lastComputation == null || !lastComputation.equals(request.getComputationId())) {
+        chunkBuilder.setComputationId(request.getComputationId());
+        lastComputation = request.getComputationId();
       }
       chunkBuilder
           .setRequestId(entry.getKey())
@@ -338,7 +376,7 @@ private void issueBatchedRequest(Map<Long, PendingRequest> requests)
 
   private void issueMultiChunkRequest(long id, PendingRequest pendingRequest)
       throws WindmillStreamShutdownException {
-    checkNotNull(pendingRequest.computationId(), "Cannot commit WorkItem w/o a computationId.");
+    checkNotNull(pendingRequest.getComputationId(), "Cannot commit WorkItem w/o a computationId.");
     ByteString serializedCommit = pendingRequest.serializedCommit();
     synchronized (this) {
       if (isShutdown) {
@@ -359,7 +397,7 @@ private void issueMultiChunkRequest(long id, PendingRequest pendingRequest)
             StreamingCommitRequestChunk.newBuilder()
                 .setRequestId(id)
                 .setSerializedWorkItemCommit(chunk)
-                .setComputationId(pendingRequest.computationId())
+                .setComputationId(pendingRequest.getComputationId())
                 .setShardingKey(pendingRequest.shardingKey());
         int remaining = serializedCommit.size() - end;
         if (remaining > 0) {
@@ -376,34 +414,46 @@ private void issueMultiChunkRequest(long id, PendingRequest pendingRequest)
     }
   }
 
-  @AutoValue
-  abstract static class PendingRequest {
+  private static class PendingRequest {
+    private final String computationId;
+    private final WorkItemCommitRequest request;
+    private final Consumer<CommitStatus> onDone;
+    private final long startTimeNanos; // System.nanoTime() of when request began.
 
-    private static PendingRequest create(
+    private PendingRequest(
         String computationId, WorkItemCommitRequest request, Consumer<CommitStatus> onDone) {
-      return new AutoValue_GrpcCommitWorkStream_PendingRequest(computationId, request, onDone);
+      this.computationId = computationId;
+      this.request = request;
+      this.onDone = onDone;
+      this.startTimeNanos = System.nanoTime();
     }
 
-    abstract String computationId();
+    String getComputationId() {
+      return computationId;
+    }
 
-    abstract WorkItemCommitRequest request();
+    WorkItemCommitRequest getRequest() {
+      return request;
+    }
 
-    abstract Consumer<CommitStatus> onDone();
+    long getStartTimeNanos() {
+      return startTimeNanos;
+    }
 
     private long getBytes() {
-      return (long) request().getSerializedSize() + computationId().length();
+      return (long) request.getSerializedSize() + computationId.length();
     }
 
     private ByteString serializedCommit() {
-      return request().toByteString();
+      return request.toByteString();
     }
 
     private void completeWithStatus(CommitStatus commitStatus) {
-      onDone().accept(commitStatus);
+      onDone.accept(commitStatus);
     }
 
     private long shardingKey() {
-      return request().getShardingKey();
+      return request.getShardingKey();
     }
 
     private void abort() {
@@ -462,7 +512,7 @@ public boolean commitWorkItem(
         return false;
       }
 
-      PendingRequest request = PendingRequest.create(computation, commitRequest, onDone);
+      PendingRequest request = new PendingRequest(computation, commitRequest, onDone);
       add(idGenerator.incrementAndGet(), request);
       return true;
     }

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

@@ -100,6 +100,7 @@ public class GrpcWindmillStreamFactory implements StatusDataProvider {
   private final Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses;
   private final java.time.Duration directStreamingRpcPhysicalStreamHalfCloseAfter;
   private final Supplier<ScheduledExecutorService> executorServiceSupplier;
+  private final java.time.Duration commitWorkStreamRetryTimeout;
 
   private GrpcWindmillStreamFactory(
       JobHeader jobHeader,
@@ -111,6 +112,7 @@ private GrpcWindmillStreamFactory(
       Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses,
       Supplier<Duration> maxBackOffSupplier,
       java.time.Duration directStreamingRpcPhysicalStreamHalfCloseAfter,
+      java.time.Duration commitWorkStreamRetryTimeout,
       Supplier<ScheduledExecutorService> executorServiceSupplier) {
     this.jobHeader = jobHeader;
     this.logEveryNStreamFailures = logEveryNStreamFailures;
@@ -132,6 +134,7 @@ private GrpcWindmillStreamFactory(
     this.directStreamingRpcPhysicalStreamHalfCloseAfter =
         directStreamingRpcPhysicalStreamHalfCloseAfter;
     this.executorServiceSupplier = executorServiceSupplier;
+    this.commitWorkStreamRetryTimeout = commitWorkStreamRetryTimeout;
   }
 
   /** @implNote Used for {@link AutoBuilder} {@link Builder} class, do not call directly. */
@@ -146,6 +149,7 @@ static GrpcWindmillStreamFactory create(
       Supplier<Duration> maxBackOffSupplier,
       int healthCheckIntervalMillis,
       java.time.Duration directStreamingRpcPhysicalStreamHalfCloseAfter,
+      java.time.Duration commitWorkStreamRetryTimeout,
       Supplier<ScheduledExecutorService> scheduledExecutorServiceSupplier) {
     GrpcWindmillStreamFactory streamFactory =
         new GrpcWindmillStreamFactory(
@@ -158,6 +162,7 @@ static GrpcWindmillStreamFactory create(
             processHeartbeatResponses,
             maxBackOffSupplier,
             directStreamingRpcPhysicalStreamHalfCloseAfter,
+            commitWorkStreamRetryTimeout,
             scheduledExecutorServiceSupplier);
 
     if (healthCheckIntervalMillis >= 0) {
@@ -200,6 +205,7 @@ public static GrpcWindmillStreamFactory.Builder of(JobHeader jobHeader) {
         .setProcessHeartbeatResponses(ignored -> {})
         .setDirectStreamingRpcPhysicalStreamHalfCloseAfter(
             DEFAULT_DIRECT_STREAMING_RPC_PHYSICAL_STREAM_HALF_CLOSE_AFTER)
+        .setCommitWorkStreamRetryTimeout(java.time.Duration.ZERO)
         .setScheduledExecutorServiceSupplier(() -> null);
   }
 
@@ -347,6 +353,7 @@ public CommitWorkStream createCommitWorkStream(CloudWindmillServiceV1Alpha1Stub
         streamIdGenerator,
         streamingRpcBatchLimit,
         java.time.Duration.ZERO,
+        commitWorkStreamRetryTimeout,
         executorForDispatchedStreams("CommitWork"));
   }
 
@@ -363,6 +370,7 @@ public CommitWorkStream createDirectCommitWorkStream(WindmillConnection connecti
         streamIdGenerator,
         streamingRpcBatchLimit,
         directStreamingRpcPhysicalStreamHalfCloseAfter,
+        java.time.Duration.ZERO,
         executorForDirectStreams(connection.backendWorkerToken(), "CommitWork"));
   }
 
@@ -426,6 +434,8 @@ Builder setProcessHeartbeatResponses(
 
     Builder setDirectStreamingRpcPhysicalStreamHalfCloseAfter(java.time.Duration timeout);
 
+    Builder setCommitWorkStreamRetryTimeout(java.time.Duration timeout);
+
     Builder setScheduledExecutorServiceSupplier(
         Supplier<ScheduledExecutorService> scheduledExecutorServiceSupplier);