This is an automated email from the ASF dual-hosted git repository.
scwhittle pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new f33feafb2b2 [Dataflow Java Streaming] Add a timeout to how long
commits will retry to the service. (#39085)
f33feafb2b2 is described below
commit f33feafb2b2ab46dd73b198a29ea182a1c3c2a00
Author: Sam Whittle <[email protected]>
AuthorDate: Wed Jun 24 15:30:10 2026 +0200
[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.
---
.../options/DataflowStreamingPipelineOptions.java | 17 ++++
.../dataflow/worker/StreamingDataflowWorker.java | 2 +
.../windmill/client/grpc/GrpcCommitWorkStream.java | 94 +++++++++++++++++-----
.../client/grpc/GrpcWindmillStreamFactory.java | 10 +++
.../client/grpc/GrpcCommitWorkStreamTest.java | 82 +++++++++++++++++++
5 files changed, 183 insertions(+), 22 deletions(-)
diff --git
a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowStreamingPipelineOptions.java
b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowStreamingPipelineOptions.java
index 0897e526382..cca261fff62 100644
---
a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowStreamingPipelineOptions.java
+++
b/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 interface DataflowStreamingPipelineOptions extends
PipelineOptions {
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();
+ }
+ }
}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
index 9e82343474c..9f063d39370 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
@@ -756,6 +756,8 @@ public final class StreamingDataflowWorker {
new
WorkHeartbeatResponseProcessor(computationStateCache::get))
.setHealthCheckIntervalMillis(
options.getWindmillServiceStreamingRpcHealthCheckPeriodMs())
+ .setCommitWorkStreamRetryTimeout(
+
java.time.Duration.ofMillis(options.getCommitWorkStreamRetryTimeoutMillis()))
.build();
return ConfigFetcherComputationStateCacheAndWindmillClient.builder()
.setWindmillDispatcherClient(dispatcherClient)
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
index d24676652fd..160b0cce013 100644
---
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
@@ -20,7 +20,6 @@ package
org.apache.beam.runners.dataflow.worker.windmill.client.grpc;
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 @@ final class GrpcCommitWorkStream
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 @@ final class GrpcCommitWorkStream
AtomicLong idGenerator,
int streamingRpcBatchLimit,
Duration halfClosePhysicalStreamAfter,
+ Duration maxRetryDuration,
ScheduledExecutorService executor) {
super(
LOG,
@@ -104,6 +105,7 @@ final class GrpcCommitWorkStream
this.idGenerator = idGenerator;
this.jobHeader = jobHeader;
this.streamingRpcBatchLimit = streamingRpcBatchLimit;
+ this.maxRetryDuration = maxRetryDuration;
}
static GrpcCommitWorkStream create(
@@ -118,6 +120,7 @@ final class GrpcCommitWorkStream
AtomicLong idGenerator,
int streamingRpcBatchLimit,
Duration halfClosePhysicalStreamAfter,
+ Duration maxRetryDuration,
ScheduledExecutorService executor) {
return new GrpcCommitWorkStream(
backendWorkerToken,
@@ -130,6 +133,7 @@ final class GrpcCommitWorkStream
idGenerator,
streamingRpcBatchLimit,
halfClosePhysicalStreamAfter,
+ maxRetryDuration,
executor);
}
@@ -224,14 +228,48 @@ final class GrpcCommitWorkStream
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 @@ final class GrpcCommitWorkStream
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 @@ final class GrpcCommitWorkStream
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 @@ final class GrpcCommitWorkStream
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 @@ final class GrpcCommitWorkStream
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 @@ final class GrpcCommitWorkStream
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 @@ final class GrpcCommitWorkStream
}
}
- @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 @@ final class GrpcCommitWorkStream
return false;
}
- PendingRequest request = PendingRequest.create(computation,
commitRequest, onDone);
+ PendingRequest request = new PendingRequest(computation, commitRequest,
onDone);
add(idGenerator.incrementAndGet(), request);
return true;
}
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
index 0184b88d53c..97ca3c4e83d 100644
---
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
@@ -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 @@ public class GrpcWindmillStreamFactory implements
StatusDataProvider {
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 @@ public class GrpcWindmillStreamFactory implements
StatusDataProvider {
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 @@ public class GrpcWindmillStreamFactory implements
StatusDataProvider {
Supplier<Duration> maxBackOffSupplier,
int healthCheckIntervalMillis,
java.time.Duration directStreamingRpcPhysicalStreamHalfCloseAfter,
+ java.time.Duration commitWorkStreamRetryTimeout,
Supplier<ScheduledExecutorService> scheduledExecutorServiceSupplier) {
GrpcWindmillStreamFactory streamFactory =
new GrpcWindmillStreamFactory(
@@ -158,6 +162,7 @@ public class GrpcWindmillStreamFactory implements
StatusDataProvider {
processHeartbeatResponses,
maxBackOffSupplier,
directStreamingRpcPhysicalStreamHalfCloseAfter,
+ commitWorkStreamRetryTimeout,
scheduledExecutorServiceSupplier);
if (healthCheckIntervalMillis >= 0) {
@@ -200,6 +205,7 @@ public class GrpcWindmillStreamFactory implements
StatusDataProvider {
.setProcessHeartbeatResponses(ignored -> {})
.setDirectStreamingRpcPhysicalStreamHalfCloseAfter(
DEFAULT_DIRECT_STREAMING_RPC_PHYSICAL_STREAM_HALF_CLOSE_AFTER)
+ .setCommitWorkStreamRetryTimeout(java.time.Duration.ZERO)
.setScheduledExecutorServiceSupplier(() -> null);
}
@@ -347,6 +353,7 @@ public class GrpcWindmillStreamFactory implements
StatusDataProvider {
streamIdGenerator,
streamingRpcBatchLimit,
java.time.Duration.ZERO,
+ commitWorkStreamRetryTimeout,
executorForDispatchedStreams("CommitWork"));
}
@@ -363,6 +370,7 @@ public class GrpcWindmillStreamFactory implements
StatusDataProvider {
streamIdGenerator,
streamingRpcBatchLimit,
directStreamingRpcPhysicalStreamHalfCloseAfter,
+ java.time.Duration.ZERO,
executorForDirectStreams(connection.backendWorkerToken(),
"CommitWork"));
}
@@ -426,6 +434,8 @@ public class GrpcWindmillStreamFactory implements
StatusDataProvider {
Builder
setDirectStreamingRpcPhysicalStreamHalfCloseAfter(java.time.Duration timeout);
+ Builder setCommitWorkStreamRetryTimeout(java.time.Duration timeout);
+
Builder setScheduledExecutorServiceSupplier(
Supplier<ScheduledExecutorService> scheduledExecutorServiceSupplier);
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStreamTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStreamTest.java
index e9fd55fa566..9c3d5c9c3ef 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStreamTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStreamTest.java
@@ -38,6 +38,7 @@ import java.util.concurrent.ExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;
import
org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
@@ -1133,6 +1134,87 @@ public class GrpcCommitWorkStreamTest {
assertTrue(commitWorkStream.awaitTermination(10, TimeUnit.SECONDS));
}
+ @Test
+ public void testCommitWorkItem_stopsRetriesAfterDuration() throws Exception {
+ int numCommits = 1;
+ CountDownLatch commitProcessed = new CountDownLatch(numCommits);
+ AtomicReference<Windmill.CommitStatus> commitStatus = new
AtomicReference<>();
+ GrpcCommitWorkStream commitWorkStream =
+ (GrpcCommitWorkStream)
+ GrpcWindmillStreamFactory.of(TEST_JOB_HEADER)
+
.setDirectStreamingRpcPhysicalStreamHalfCloseAfter(Duration.ofMinutes(1))
+ .setCommitWorkStreamRetryTimeout(Duration.ofNanos(1))
+ .build()
+
.createCommitWorkStream(CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel));
+ commitWorkStream.start();
+
+ FakeWindmillGrpcService.CommitStreamInfo streamInfo =
waitForConnectionAndConsumeHeader();
+ try (WindmillStream.CommitWorkStream.RequestBatcher batcher =
commitWorkStream.batcher()) {
+ assertTrue(
+ batcher.commitWorkItem(
+ COMPUTATION_ID,
+ workItemCommitRequest(0),
+ status -> {
+ commitStatus.set(status);
+ commitProcessed.countDown();
+ }));
+ }
+
+ // The next request should have some chunks.
+ assertThat(streamInfo.requests.take().getCommitChunkList()).isNotEmpty();
+
+ // We won't get responses so we will have some pending requests.
+ assertThat(commitProcessed.getCount()).isGreaterThan(0);
+ streamInfo.responseObserver.onError(new IOException("test error"));
+ commitProcessed.await();
+ assertThat(commitStatus.get()).isEqualTo(Windmill.CommitStatus.ABORTED);
+ }
+
+ @Test
+ public void testCommitWorkItem_retriesWithLongerCommitRetryTimeout() throws
Exception {
+ int numCommits = 1;
+ CountDownLatch commitProcessed = new CountDownLatch(numCommits);
+ AtomicReference<Windmill.CommitStatus> commitStatus = new
AtomicReference<>();
+ GrpcCommitWorkStream commitWorkStream =
+ (GrpcCommitWorkStream)
+ GrpcWindmillStreamFactory.of(TEST_JOB_HEADER)
+
.setDirectStreamingRpcPhysicalStreamHalfCloseAfter(Duration.ofMinutes(1))
+ // Verifies that if this is set but is not exceeded that a
retry occurs.
+ .setCommitWorkStreamRetryTimeout(Duration.ofMinutes(10))
+ .build()
+
.createCommitWorkStream(CloudWindmillServiceV1Alpha1Grpc.newStub(inProcessChannel));
+ commitWorkStream.start();
+
+ FakeWindmillGrpcService.CommitStreamInfo streamInfo =
waitForConnectionAndConsumeHeader();
+ try (WindmillStream.CommitWorkStream.RequestBatcher batcher =
commitWorkStream.batcher()) {
+ assertTrue(
+ batcher.commitWorkItem(
+ COMPUTATION_ID,
+ workItemCommitRequest(0),
+ status -> {
+ commitStatus.set(status);
+ commitProcessed.countDown();
+ }));
+ }
+
+ // The next request should have some chunks.
+ assertThat(streamInfo.requests.take().getCommitChunkList()).isNotEmpty();
+
+ // We won't get responses so we will have some pending requests.
+ assertThat(commitProcessed.getCount()).isGreaterThan(0);
+ streamInfo.responseObserver.onError(new IOException("test error"));
+
+ // The stream should reconnect and retry the requests.
+ FakeWindmillGrpcService.CommitStreamInfo reconnectStreamInfo =
+ waitForConnectionAndConsumeHeader();
+ Windmill.StreamingCommitWorkRequest reconnectRequest =
reconnectStreamInfo.requests.take();
+ assertEquals(1, reconnectRequest.getCommitChunkCount());
+ reconnectStreamInfo.responseObserver.onNext(
+ Windmill.StreamingCommitResponse.newBuilder().addRequestId(1).build());
+ commitProcessed.await();
+ assertThat(commitStatus.get()).isEqualTo(Windmill.CommitStatus.OK);
+ }
+
private FakeWindmillGrpcService.CommitStreamInfo
waitForConnectionAndConsumeHeader() {
try {
FakeWindmillGrpcService.CommitStreamInfo info =
fakeService.waitForConnectedCommitStream();
