diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java index 526b67890783..36001c151508 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java @@ -108,6 +108,11 @@ boolean commitWorkItem( Windmill.WorkItemCommitRequest request, Consumer onDone); + boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone); + /** Flushes any pending work items to the wire. */ void flush(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java index b840d22a3434..bbd6cfc9432b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java @@ -17,35 +17,89 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.client.commits; -import com.google.auto.value.AutoValue; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import org.apache.beam.runners.dataflow.worker.streaming.ComputationState; import org.apache.beam.runners.dataflow.worker.streaming.Work; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.MultiKeyWorkItemCommitRequest; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; /** Value class for a queued commit. */ @Internal -@AutoValue -public abstract class Commit { +public class Commit { + + private final ComputationState computationState; + private final ImmutableList workBatch; + private final @Nullable WorkItemCommitRequest singleKeyRequest; + private final @Nullable MultiKeyWorkItemCommitRequest multiKeyRequest; public static Commit create( WorkItemCommitRequest request, ComputationState computationState, Work work) { Preconditions.checkArgument(request.getSerializedSize() > 0); - return new AutoValue_Commit(request, computationState, work); + return new Commit(computationState, ImmutableList.of(work), request, null); + } + + public static Commit createMultiKey( + MultiKeyWorkItemCommitRequest multiKeyRequest, + ComputationState computationState, + ImmutableList workBatch) { + Preconditions.checkArgument(!workBatch.isEmpty()); + return new Commit(computationState, workBatch, null, multiKeyRequest); + } + + private Commit( + ComputationState computationState, + ImmutableList workBatch, + @Nullable WorkItemCommitRequest singleKeyRequest, + @Nullable MultiKeyWorkItemCommitRequest multiKeyRequest) { + this.computationState = computationState; + this.workBatch = workBatch; + this.singleKeyRequest = singleKeyRequest; + this.multiKeyRequest = multiKeyRequest; } public final String computationId() { return computationState().getComputationId(); } - public abstract WorkItemCommitRequest request(); + public @Nullable WorkItemCommitRequest singleKeyRequest() { + return singleKeyRequest; + }; - public abstract ComputationState computationState(); + public ComputationState computationState() { + return computationState; + } + + public @Nullable MultiKeyWorkItemCommitRequest multiKeyRequest() { + return multiKeyRequest; + } - public abstract Work work(); + public ImmutableList workBatch() { + return workBatch; + } + + public final int getSerializedByteSize() { + if (multiKeyRequest() != null) { + return checkStateNotNull(multiKeyRequest()).getSerializedSize(); + } + return checkStateNotNull(singleKeyRequest()).getSerializedSize(); + } - public final int getSize() { - return request().getSerializedSize(); + @Override + public String toString() { + Work work = workBatch.get(0); + return "[computationId=" + + computationId() + + ", shardingKey=" + + work.getShardedKey() + + ", workId=" + + work.id() + + ", workBatchSize=" + + workBatch.size() + + "]"; } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commits.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commits.java index 498e90f78e29..0607baebeb1d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commits.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commits.java @@ -31,6 +31,6 @@ public final class Commits { private Commits() {} public static WeightedSemaphore maxCommitByteSemaphore() { - return WeightedSemaphore.create(MAX_QUEUED_COMMITS_BYTES, Commit::getSize); + return WeightedSemaphore.create(MAX_QUEUED_COMMITS_BYTES, Commit::getSerializedByteSize); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/CompleteCommit.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/CompleteCommit.java index e33e853d3d76..6c0a5a98e2ab 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/CompleteCommit.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/CompleteCommit.java @@ -37,26 +37,11 @@ @AutoValue public abstract class CompleteCommit { - public static CompleteCommit create(Commit commit, CommitStatus commitStatus) { - return new AutoValue_CompleteCommit( - commit.computationId(), - ShardedKey.create(commit.request().getKey(), commit.request().getShardingKey()), - WorkId.builder() - .setWorkToken(commit.request().getWorkToken()) - .setCacheToken(commit.request().getCacheToken()) - .build(), - commitStatus); - } - public static CompleteCommit create( String computationId, ShardedKey shardedKey, WorkId workId, CommitStatus status) { return new AutoValue_CompleteCommit(computationId, shardedKey, workId, status); } - public static CompleteCommit forFailedWork(Commit commit) { - return create(commit, CommitStatus.ABORTED); - } - public abstract String computationId(); public abstract ShardedKey shardedKey(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java index 20b95b0661d0..d627490dfe98 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java @@ -17,6 +17,9 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.client.commits; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; + import java.util.HashMap; import java.util.Map; import java.util.concurrent.ExecutorService; @@ -112,7 +115,8 @@ private void commitLoop() { } while (commit != null) { ComputationState computationState = commit.computationState(); - commit.work().setState(Work.State.COMMITTING); + checkState(commit.workBatch().size() == 1); + commit.workBatch().get(0).setState(Work.State.COMMITTING); Windmill.ComputationCommitWorkRequest.Builder computationRequestBuilder = computationRequestMap.get(computationState); if (computationRequestBuilder == null) { @@ -120,10 +124,10 @@ private void commitLoop() { computationRequestBuilder.setComputationId(computationState.getComputationId()); computationRequestMap.put(computationState, computationRequestBuilder); } - computationRequestBuilder.addRequests(commit.request()); + computationRequestBuilder.addRequests(checkStateNotNull(commit.singleKeyRequest())); // Send the request if we've exceeded the bytes or there is no more // pending work. commitBytes is a long, so this cannot overflow. - commitBytes += commit.getSize(); + commitBytes += commit.getSerializedByteSize(); if (commitBytes >= TARGET_COMMIT_BUNDLE_BYTES) { break; } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java index b68f53121b86..83d0dfc6cda4 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.client.commits; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import com.google.auto.value.AutoBuilder; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -30,6 +32,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.WeightedBoundedQueue; import org.apache.beam.runners.dataflow.worker.streaming.WeightedSemaphore; import org.apache.beam.runners.dataflow.worker.streaming.Work; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; import org.apache.beam.runners.dataflow.worker.windmill.client.CloseableStream; import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.CommitWorkStream; import org.apache.beam.sdk.annotations.Internal; @@ -100,8 +103,8 @@ public void start() { @Override public void commit(Commit commit) { - if (commit.work().isFailed()) { - failCommit(commit); + if (shouldFailCommit(commit)) { + failQueuedCommit(commit); } else { commitQueue.put(commit); } @@ -109,12 +112,7 @@ public void commit(Commit commit) { // Do this check after adding to commitQueue, else commitQueue.put() can race with // drainCommitQueue() in stop() and leave commits orphaned in the queue. if (!this.isRunning.get()) { - LOG.debug( - "Trying to queue commit on shutdown, failing commit=[computationId={}, shardingKey={}," - + " workId={} ].", - commit.computationId(), - commit.work().getShardedKey(), - commit.work().id()); + LOG.debug("Trying to queue commit on shutdown, failing commit={}", commit); drainCommitQueue(); } } @@ -141,14 +139,18 @@ public void stop() { private void drainCommitQueue() { Commit queuedCommit = commitQueue.poll(); while (queuedCommit != null) { - failCommit(queuedCommit); + failQueuedCommit(queuedCommit); queuedCommit = commitQueue.poll(); } } - private void failCommit(Commit commit) { - commit.work().setFailed(); - onCommitComplete.accept(CompleteCommit.forFailedWork(commit)); + private void failQueuedCommit(Commit commit) { + for (Work w : commit.workBatch()) { + w.setFailed(); + onCommitComplete.accept( + CompleteCommit.create( + commit.computationId(), w.getShardedKey(), w.id(), CommitStatus.ABORTED)); + } } @Override @@ -173,8 +175,8 @@ private void streamingCommitLoop() { // take() blocks until a value is available in the commitQueue. Preconditions.checkNotNull(initialCommit); - if (initialCommit.work().isFailed()) { - onCommitComplete.accept(CompleteCommit.forFailedWork(initialCommit)); + if (shouldFailCommit(initialCommit)) { + failQueuedCommit(initialCommit); initialCommit = null; continue; } @@ -194,29 +196,61 @@ private void streamingCommitLoop() { } } finally { if (initialCommit != null) { - failCommit(initialCommit); + failQueuedCommit(initialCommit); + } + } + } + + boolean shouldFailCommit(Commit commit) { + for (Work w : commit.workBatch()) { + if (w.isFailed()) { + return true; } } + return false; } /** Adds the commit to the batch if it fits, returning true if it is consumed. */ private boolean tryAddToCommitBatch(Commit commit, CommitWorkStream.RequestBatcher batcher) { Preconditions.checkNotNull(commit); - commit.work().setState(Work.State.COMMITTING); - activeCommitBytes.addAndGet(commit.getSize()); - boolean isCommitAccepted = - batcher.commitWorkItem( - commit.computationId(), - commit.request(), - commitStatus -> { - onCommitComplete.accept(CompleteCommit.create(commit, commitStatus)); - activeCommitBytes.addAndGet(-commit.getSize()); - }); + for (Work w : commit.workBatch()) { + w.setState(Work.State.COMMITTING); + } + activeCommitBytes.addAndGet(commit.getSerializedByteSize()); + boolean isCommitAccepted; + if (commit.multiKeyRequest() != null) { + isCommitAccepted = + batcher.commitMultiKeyWorkItem( + commit.computationId(), + checkStateNotNull(commit.multiKeyRequest()), + commitStatus -> { + for (Work w : commit.workBatch()) { + onCommitComplete.accept( + CompleteCommit.create( + commit.computationId(), w.getShardedKey(), w.id(), commitStatus)); + } + activeCommitBytes.addAndGet(-commit.getSerializedByteSize()); + }); + } else { + isCommitAccepted = + batcher.commitWorkItem( + commit.computationId(), + checkStateNotNull(commit.singleKeyRequest()), + commitStatus -> { + Work w = commit.workBatch().get(0); + onCommitComplete.accept( + CompleteCommit.create( + commit.computationId(), w.getShardedKey(), w.id(), commitStatus)); + activeCommitBytes.addAndGet(-commit.getSerializedByteSize()); + }); + } // Since the commit was not accepted, revert the changes made above. if (!isCommitAccepted) { - commit.work().setState(Work.State.COMMIT_QUEUED); - activeCommitBytes.addAndGet(-commit.getSize()); + for (Work w : commit.workBatch()) { + w.setState(Work.State.COMMIT_QUEUED); + } + activeCommitBytes.addAndGet(-commit.getSerializedByteSize()); } return isCommitAccepted; @@ -246,8 +280,8 @@ private boolean tryAddToCommitBatch(Commit commit, CommitWorkStream.RequestBatch } // Drop commits for failed work. Such commits will be dropped by Windmill anyway. - if (commit.work().isFailed()) { - onCommitComplete.accept(CompleteCommit.forFailedWork(commit)); + if (shouldFailCommit(commit)) { + failQueuedCommit(commit); continue; } 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 160b0cce0133..e2a54b43cf07 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 @@ -34,6 +34,7 @@ import java.util.function.Function; import javax.annotation.Nullable; import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingCommitRequestChunk; @@ -308,7 +309,7 @@ private void flushInternal(Map requests) if (requests.size() == 1) { Map.Entry elem = requests.entrySet().iterator().next(); - if (elem.getValue().getRequest().getSerializedSize() + if (elem.getValue().serializedCommit().size() > AbstractWindmillStream.RPC_STREAM_CHUNK_SIZE) { issueMultiChunkRequest(elem.getKey(), elem.getValue()); } else { @@ -324,9 +325,10 @@ private void issueSingleRequest(long id, PendingRequest pendingRequest) StreamingCommitWorkRequest.Builder requestBuilder = StreamingCommitWorkRequest.newBuilder(); requestBuilder .addCommitChunkBuilder() - .setComputationId(pendingRequest.getComputationId()) + .setComputationId(pendingRequest.computationId()) .setRequestId(id) .setShardingKey(pendingRequest.shardingKey()) + .setCommitType(pendingRequest.commitType()) .setSerializedWorkItemCommit(pendingRequest.serializedCommit()); StreamingCommitWorkRequest chunk = requestBuilder.build(); synchronized (this) { @@ -349,14 +351,15 @@ private void issueBatchedRequest(Map requests) for (Map.Entry entry : requests.entrySet()) { PendingRequest request = entry.getValue(); StreamingCommitRequestChunk.Builder chunkBuilder = requestBuilder.addCommitChunkBuilder(); - if (lastComputation == null || !lastComputation.equals(request.getComputationId())) { - chunkBuilder.setComputationId(request.getComputationId()); - lastComputation = request.getComputationId(); + if (lastComputation == null || !lastComputation.equals(request.computationId())) { + chunkBuilder.setComputationId(request.computationId()); + lastComputation = request.computationId(); } chunkBuilder .setRequestId(entry.getKey()) .setShardingKey(request.shardingKey()) - .setSerializedWorkItemCommit(request.serializedCommit()); + .setSerializedWorkItemCommit(request.serializedCommit()) + .setCommitType(request.commitType()); } StreamingCommitWorkRequest request = requestBuilder.build(); synchronized (this) { @@ -376,7 +379,7 @@ private void issueBatchedRequest(Map requests) private void issueMultiChunkRequest(long id, PendingRequest pendingRequest) throws WindmillStreamShutdownException { - checkNotNull(pendingRequest.getComputationId(), "Cannot commit WorkItem w/o a computationId."); + checkNotNull(pendingRequest.computationId(), "Cannot commit WorkItem w/o a computationId."); ByteString serializedCommit = pendingRequest.serializedCommit(); synchronized (this) { if (isShutdown) { @@ -397,8 +400,9 @@ private void issueMultiChunkRequest(long id, PendingRequest pendingRequest) StreamingCommitRequestChunk.newBuilder() .setRequestId(id) .setSerializedWorkItemCommit(chunk) - .setComputationId(pendingRequest.getComputationId()) - .setShardingKey(pendingRequest.shardingKey()); + .setComputationId(pendingRequest.computationId()) + .setShardingKey(pendingRequest.shardingKey()) + .setCommitType(pendingRequest.commitType()); int remaining = serializedCommit.size() - end; if (remaining > 0) { chunkBuilder.setRemainingBytesForWorkItem(remaining); @@ -416,24 +420,44 @@ private void issueMultiChunkRequest(long id, PendingRequest pendingRequest) private static class PendingRequest { private final String computationId; - private final WorkItemCommitRequest request; + private final long shardingKey; + private final ByteString serializedCommit; + private final StreamingCommitRequestChunk.CommitType commitType; private final Consumer onDone; private final long startTimeNanos; // System.nanoTime() of when request began. private PendingRequest( - String computationId, WorkItemCommitRequest request, Consumer onDone) { + String computationId, + long shardingKey, + ByteString serializedCommit, + StreamingCommitRequestChunk.CommitType commitType, + Consumer onDone) { this.computationId = computationId; - this.request = request; + this.shardingKey = shardingKey; + this.serializedCommit = serializedCommit; + this.commitType = commitType; this.onDone = onDone; this.startTimeNanos = System.nanoTime(); } - String getComputationId() { + String computationId() { return computationId; } - WorkItemCommitRequest getRequest() { - return request; + long shardingKey() { + return shardingKey; + } + + ByteString serializedCommit() { + return serializedCommit; + } + + StreamingCommitRequestChunk.CommitType commitType() { + return commitType; + } + + Consumer onDone() { + return onDone; } long getStartTimeNanos() { @@ -441,21 +465,13 @@ long getStartTimeNanos() { } private long getBytes() { - return (long) request.getSerializedSize() + computationId.length(); - } - - private ByteString serializedCommit() { - return request.toByteString(); + return (long) serializedCommit.size() + computationId.length(); } private void completeWithStatus(CommitStatus commitStatus) { onDone.accept(commitStatus); } - private long shardingKey() { - return request.getShardingKey(); - } - private void abort() { completeWithStatus(CommitStatus.ABORTED); } @@ -512,7 +528,34 @@ public boolean commitWorkItem( return false; } - PendingRequest request = new PendingRequest(computation, commitRequest, onDone); + PendingRequest request = + new PendingRequest( + computation, + commitRequest.getShardingKey(), + commitRequest.toByteString(), + StreamingCommitRequestChunk.CommitType.COMMIT_TYPE_SINGLE_KEY, + onDone); + add(idGenerator.incrementAndGet(), request); + return true; + } + + @Override + public boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest commitRequest, + Consumer onDone) { + Preconditions.checkArgument(commitRequest.getRequestsCount() > 0); + if (!canAccept(commitRequest.getSerializedSize() + computation.length())) { + return false; + } + PendingRequest request = + new PendingRequest( + computation, + // Any key in the batch for routing + commitRequest.getRequests(0).getShardingKey(), + commitRequest.toByteString(), + StreamingCommitRequestChunk.CommitType.COMMIT_TYPE_MULTI_KEY, + onDone); add(idGenerator.incrementAndGet(), request); return true; } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java index 5be8ec0a6c72..e5d68376a7d5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java @@ -29,7 +29,6 @@ import java.util.ArrayList; import java.util.Collection; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; @@ -37,10 +36,12 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.function.Function; import javax.annotation.concurrent.GuardedBy; @@ -49,6 +50,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.WorkId; import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillMetadataServiceV1Alpha1Grpc; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitWorkResponse; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationCommitWorkRequest; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationGetDataRequest; @@ -87,8 +89,11 @@ public final class FakeWindmillServer extends WindmillServerStub { private final ResponseQueue dataToOffer; private final ResponseQueue commitsToOffer; private final Map streamingCommitsToOffer; + private final AtomicReference multiKeyCommitStatusToOffer; // Keys are work tokens. private final Map commitsReceived; + private final List multiKeyCommitsReceived = + new CopyOnWriteArrayList<>(); private final ArrayList statsReceived; private final LinkedBlockingQueue exceptions; private final AtomicInteger expectedExceptionCount; @@ -118,7 +123,9 @@ public FakeWindmillServer( commitsToOffer = new ResponseQueue() .returnByDefault(CommitWorkResponse.getDefaultInstance()); - streamingCommitsToOffer = new HashMap<>(); + streamingCommitsToOffer = new ConcurrentHashMap<>(); + // Respond multikey commits with ok, unless overridden. + multiKeyCommitStatusToOffer = new AtomicReference<>(CommitStatus.OK); commitsReceived = new ConcurrentHashMap<>(); exceptions = new LinkedBlockingQueue<>(); expectedExceptionCount = new AtomicInteger(); @@ -153,6 +160,11 @@ public Map whenCommitWorkStreamCalled() { return streamingCommitsToOffer; } + /** @param commitStatus status to return to multiKeyCommits */ + public void setMultiKeyCommitStatus(CommitStatus commitStatus) { + this.multiKeyCommitStatusToOffer.set(commitStatus); + } + @Override public Windmill.GetWorkResponse getWork(Windmill.GetWorkRequest request) { LOG.debug("getWorkRequest: {}", request.toString()); @@ -400,6 +412,7 @@ public void shutdown() {} public RequestBatcher batcher() { return new RequestBatcher() { final List requests = new ArrayList<>(); + final List multiKeyRequests = new ArrayList<>(); @Override public boolean commitWorkItem( @@ -423,6 +436,17 @@ public boolean commitWorkItem( return true; } + @Override + public boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone) { + LOG.debug("commitWorkStream::commitMultiKeyWorkItem: {}", request); + multiKeyRequests.add(new MultiKeyRequestAndDone(request, onDone)); + flush(); + return true; + } + @Override public void flush() { for (RequestAndDone elem : requests) { @@ -445,6 +469,24 @@ public void flush() { .orElse(Windmill.CommitStatus.OK)); } requests.clear(); + + for (MultiKeyRequestAndDone elem : multiKeyRequests) { + if (dropStreamingCommits) { + for (WorkItemCommitRequest workRequest : elem.request.getRequestsList()) { + droppedStreamingCommits.put(workRequest.getWorkToken(), elem.onDone); + } + continue; + } + + multiKeyCommitsReceived.add(elem.request); + for (WorkItemCommitRequest workRequest : elem.request.getRequestsList()) { + commitsReceived.put(workRequest.getWorkToken(), workRequest); + } + + Windmill.CommitStatus status = multiKeyCommitStatusToOffer.get(); + elem.onDone.accept(status); + } + multiKeyRequests.clear(); } class RequestAndDone { @@ -456,6 +498,18 @@ class RequestAndDone { this.onDone = onDone; } } + + class MultiKeyRequestAndDone { + final Consumer onDone; + final Windmill.MultiKeyWorkItemCommitRequest request; + + MultiKeyRequestAndDone( + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone) { + this.request = request; + this.onDone = onDone; + } + } }; } @@ -518,6 +572,15 @@ public Map waitForAndGetCommits(int numCommits) { public void clearCommitsReceived() { commitsRequested = 0; commitsReceived.clear(); + multiKeyCommitsReceived.clear(); + } + + public List getMultiKeyCommitsReceived() { + return multiKeyCommitsReceived; + } + + public void clearMultiKeyCommitsReceived() { + multiKeyCommitsReceived.clear(); } public ConcurrentHashMap> waitForDroppedCommits( diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java index 5c3132ae471d..0596210a0270 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java @@ -128,10 +128,11 @@ public void testCommit() { fakeWindmillServer.waitForAndGetCommits(commits.size()); for (Commit commit : commits) { + assertThat(commit.workBatch()).hasSize(1); Windmill.WorkItemCommitRequest request = - committed.get(commit.work().getWorkItem().getWorkToken()); + committed.get(commit.workBatch().get(0).getWorkItem().getWorkToken()); assertNotNull(request); - assertThat(request).isEqualTo(commit.request()); + assertThat(request).isEqualTo(commit.singleKeyRequest()); } assertThat(completeCommits).hasSize(commits.size()); @@ -141,12 +142,14 @@ public void testCommit() { (CompleteCommit completeCommit, Commit commit) -> completeCommit.computationId().equals(commit.computationId()) && completeCommit.status() == Windmill.CommitStatus.OK - && completeCommit.workId().equals(commit.work().id()) + && commit.workBatch().size() == 1 + && completeCommit.workId().equals(commit.workBatch().get(0).id()) && completeCommit .shardedKey() .equals( ShardedKey.create( - commit.request().getKey(), commit.request().getShardingKey())), + commit.singleKeyRequest().getKey(), + commit.singleKeyRequest().getShardingKey())), "expected to equal")) .containsExactlyElementsIn(commits); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java index 01197622c24d..881cf620e8d8 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus.OK; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; @@ -53,6 +54,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.WorkId; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest; import org.apache.beam.runners.dataflow.worker.windmill.client.CloseableStream; @@ -62,6 +64,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.testing.GrpcCleanupRule; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.joda.time.Duration; import org.joda.time.Instant; @@ -134,12 +137,10 @@ private static ComputationState createComputationState(String computationId) { null); } - private static CompleteCommit asCompleteCommit(Commit commit, Windmill.CommitStatus status) { - if (commit.work().isFailed()) { - return CompleteCommit.forFailedWork(commit); - } - - return CompleteCommit.create(commit, status); + private static CompleteCommit asCompleteCommit( + String computationId, Work work, Windmill.CommitStatus status) { + Windmill.CommitStatus finalStatus = work.isFailed() ? Windmill.CommitStatus.ABORTED : status; + return CompleteCommit.create(computationId, work.getShardedKey(), work.id(), finalStatus); } @Before @@ -186,10 +187,15 @@ public void testCommit_sendsCommitsToStreamingEngine() { waitForExpectedSetSize(completeCommits, 5); for (Commit commit : commits) { - WorkItemCommitRequest request = committed.get(commit.work().getWorkItem().getWorkToken()); + assertThat(commit.workBatch()).hasSize(1); + WorkItemCommitRequest request = + committed.get(commit.workBatch().get(0).getWorkItem().getWorkToken()); assertNotNull(request); - assertThat(request).isEqualTo(commit.request()); - assertThat(completeCommits).contains(asCompleteCommit(commit, Windmill.CommitStatus.OK)); + assertThat(request).isEqualTo(commit.singleKeyRequest()); + assertThat(completeCommits) + .contains( + asCompleteCommit( + commit.computationId(), commit.workBatch().get(0), Windmill.CommitStatus.OK)); } workCommitter.stop(); @@ -224,14 +230,24 @@ public void testCommit_handlesFailedCommits() { waitForExpectedSetSize(completeCommits, 10); for (Commit commit : commits) { - if (commit.work().isFailed()) { + assertThat(commit.workBatch()).hasSize(1); + if (commit.workBatch().get(0).isFailed()) { assertThat(completeCommits) - .contains(asCompleteCommit(commit, Windmill.CommitStatus.ABORTED)); - assertThat(committed).doesNotContainKey(commit.work().getWorkItem().getWorkToken()); + .contains( + asCompleteCommit( + commit.computationId(), + commit.workBatch().get(0), + Windmill.CommitStatus.ABORTED)); + assertThat(committed) + .doesNotContainKey(commit.workBatch().get(0).getWorkItem().getWorkToken()); } else { - assertThat(completeCommits).contains(asCompleteCommit(commit, Windmill.CommitStatus.OK)); + assertThat(completeCommits) + .contains( + asCompleteCommit( + commit.computationId(), commit.workBatch().get(0), Windmill.CommitStatus.OK)); assertThat(committed) - .containsEntry(commit.work().getWorkItem().getWorkToken(), commit.request()); + .containsEntry( + commit.workBatch().get(0).getWorkItem().getWorkToken(), commit.singleKeyRequest()); } } @@ -282,11 +298,17 @@ public void testCommit_handlesCompleteCommits_commitStatusNotOK() { waitForExpectedSetSize(completeCommits, commits.size()); for (Commit commit : commits) { - WorkItemCommitRequest request = committed.get(commit.work().getWorkItem().getWorkToken()); + assertThat(commit.workBatch()).hasSize(1); + WorkItemCommitRequest request = + committed.get(commit.workBatch().get(0).getWorkItem().getWorkToken()); assertNotNull(request); - assertThat(request).isEqualTo(commit.request()); + assertThat(request).isEqualTo(commit.singleKeyRequest()); assertThat(completeCommits) - .contains(asCompleteCommit(commit, expectedCommitStatus.get(commit.work().id()))); + .contains( + asCompleteCommit( + commit.computationId(), + commit.workBatch().get(0), + expectedCommitStatus.get(commit.workBatch().get(0).id()))); } workCommitter.stop(); @@ -313,6 +335,14 @@ public boolean commitWorkItem( return false; } + @Override + public boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone) { + return false; + } + @Override public void flush() {} }; @@ -370,7 +400,8 @@ public void shutdown() {} } for (Commit commit : commits) { - assertTrue(commit.work().isFailed()); + assertThat(commit.workBatch()).hasSize(1); + assertTrue(commit.workBatch().get(0).isFailed()); } } @@ -409,10 +440,15 @@ public void testMultipleCommitSendersSingleStream() { waitForExpectedSetSize(completeCommits, commits.size()); for (Commit commit : commits) { - WorkItemCommitRequest request = committed.get(commit.work().getWorkItem().getWorkToken()); + assertThat(commit.workBatch()).hasSize(1); + WorkItemCommitRequest request = + committed.get(commit.workBatch().get(0).getWorkItem().getWorkToken()); assertNotNull(request); - assertThat(request).isEqualTo(commit.request()); - assertThat(completeCommits).contains(asCompleteCommit(commit, Windmill.CommitStatus.OK)); + assertThat(request).isEqualTo(commit.singleKeyRequest()); + assertThat(completeCommits) + .contains( + asCompleteCommit( + commit.computationId(), commit.workBatch().get(0), Windmill.CommitStatus.OK)); } workCommitter.stop(); @@ -474,4 +510,207 @@ public void testStop_drainsCommitQueue_concurrentCommit() waitForExpectedSetSize(completeCommits, sentCommits.intValue()); } + + @Test + public void testCommit_multiKeyCommitSuccess() { + Set completeCommits = Collections.newSetFromMap(new ConcurrentHashMap<>()); + workCommitter = createWorkCommitter(completeCommits::add); + + Work workA = createMockWork(101L); + Work workB = createMockWork(102L); + Work workC = createMockWork(103L); + + Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest = + Windmill.MultiKeyWorkItemCommitRequest.newBuilder() + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workA.getWorkItem().getKey()) + .setShardingKey(workA.getWorkItem().getShardingKey()) + .setWorkToken(workA.getWorkItem().getWorkToken()) + .setCacheToken(workA.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workB.getWorkItem().getKey()) + .setShardingKey(workB.getWorkItem().getShardingKey()) + .setWorkToken(workB.getWorkItem().getWorkToken()) + .setCacheToken(workB.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workC.getWorkItem().getKey()) + .setShardingKey(workC.getWorkItem().getShardingKey()) + .setWorkToken(workC.getWorkItem().getWorkToken()) + .setCacheToken(workC.getWorkItem().getCacheToken()) + .build()) + .build(); + + Commit commit = + Commit.createMultiKey( + multiKeyRequest, + createComputationState("computationId"), + ImmutableList.of(workA, workB, workC)); + + workCommitter.start(); + workCommitter.commit(commit); + + // Wait for the server to receive and process the commits + fakeWindmillServer.waitForAndGetCommits(3); + waitForExpectedSetSize(completeCommits, 3); + + // Verify that FakeWindmillServer received all 3 work requests in multiKeyCommitsReceived + List multiKeyCommits = + fakeWindmillServer.getMultiKeyCommitsReceived(); + assertThat(multiKeyCommits).hasSize(1); + assertThat(multiKeyCommits.get(0)).isEqualTo(multiKeyRequest); + + // Verify all three works are completed successfully + assertThat(completeCommits) + .containsExactly( + CompleteCommit.create( + "computationId", workA.getShardedKey(), workA.id(), CommitStatus.OK), + CompleteCommit.create( + "computationId", workB.getShardedKey(), workB.id(), CommitStatus.OK), + CompleteCommit.create( + "computationId", workC.getShardedKey(), workC.id(), CommitStatus.OK)); + + // There should be no more commits in the queue + assertEquals(0, workCommitter.currentActiveCommitBytes()); + workCommitter.stop(); + } + + @Test + public void testCommit_multiKeyCommitFailedWork() { + Set completeCommits = Collections.newSetFromMap(new ConcurrentHashMap<>()); + workCommitter = createWorkCommitter(completeCommits::add); + + Work workA = createMockWork(101L); + Work workB = createMockWork(102L); + Work workC = createMockWork(103L); + + // Mark non-primary key B as failed + workB.setFailed(); + + Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest = + Windmill.MultiKeyWorkItemCommitRequest.newBuilder() + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workA.getWorkItem().getKey()) + .setShardingKey(workA.getWorkItem().getShardingKey()) + .setWorkToken(workA.getWorkItem().getWorkToken()) + .setCacheToken(workA.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workB.getWorkItem().getKey()) + .setShardingKey(workB.getWorkItem().getShardingKey()) + .setWorkToken(workB.getWorkItem().getWorkToken()) + .setCacheToken(workB.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workC.getWorkItem().getKey()) + .setShardingKey(workC.getWorkItem().getShardingKey()) + .setWorkToken(workC.getWorkItem().getWorkToken()) + .setCacheToken(workC.getWorkItem().getCacheToken()) + .build()) + .build(); + + Commit commit = + Commit.createMultiKey( + multiKeyRequest, + createComputationState("computationId"), + ImmutableList.of(workA, workB, workC)); + + workCommitter.start(); + workCommitter.commit(commit); + + // The entire batch must be aborted immediately without making network calls + waitForExpectedSetSize(completeCommits, 3); + + // Verify all three works are aborted individually + assertThat(completeCommits) + .containsExactly( + CompleteCommit.create( + "computationId", workA.getShardedKey(), workA.id(), CommitStatus.ABORTED), + CompleteCommit.create( + "computationId", workB.getShardedKey(), workB.id(), CommitStatus.ABORTED), + CompleteCommit.create( + "computationId", workC.getShardedKey(), workC.id(), CommitStatus.ABORTED)); + + // There should be no more commits in the queue + assertEquals(0, workCommitter.currentActiveCommitBytes()); + workCommitter.stop(); + } + + @Test + public void testCommit_multiKeyCommitStatusNotOK() { + Set completeCommits = Collections.newSetFromMap(new ConcurrentHashMap<>()); + workCommitter = createWorkCommitter(completeCommits::add); + + Work workA = createMockWork(101L); + Work workB = createMockWork(102L); + Work workC = createMockWork(103L); + + Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest = + Windmill.MultiKeyWorkItemCommitRequest.newBuilder() + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workA.getWorkItem().getKey()) + .setShardingKey(workA.getWorkItem().getShardingKey()) + .setWorkToken(workA.getWorkItem().getWorkToken()) + .setCacheToken(workA.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workB.getWorkItem().getKey()) + .setShardingKey(workB.getWorkItem().getShardingKey()) + .setWorkToken(workB.getWorkItem().getWorkToken()) + .setCacheToken(workB.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workC.getWorkItem().getKey()) + .setShardingKey(workC.getWorkItem().getShardingKey()) + .setWorkToken(workC.getWorkItem().getWorkToken()) + .setCacheToken(workC.getWorkItem().getCacheToken()) + .build()) + .build(); + + Commit commit = + Commit.createMultiKey( + multiKeyRequest, + createComputationState("computationId"), + ImmutableList.of(workA, workB, workC)); + + // Respond to multi key commit with NOT_FOUND status. + fakeWindmillServer.setMultiKeyCommitStatus(CommitStatus.NOT_FOUND); + + workCommitter.start(); + workCommitter.commit(commit); + + // Wait for the server to receive and process the commits + fakeWindmillServer.waitForAndGetCommits(3); + waitForExpectedSetSize(completeCommits, 3); + + // Verify that FakeWindmillServer received the multi-key commit + List multiKeyCommits = + fakeWindmillServer.getMultiKeyCommitsReceived(); + assertThat(multiKeyCommits).hasSize(1); + assertThat(multiKeyCommits.get(0)).isEqualTo(multiKeyRequest); + + // Verify all three works in the multi-key commit are completed with NOT_FOUND status + assertThat(completeCommits) + .containsExactly( + CompleteCommit.create( + "computationId", workA.getShardedKey(), workA.id(), CommitStatus.NOT_FOUND), + CompleteCommit.create( + "computationId", workB.getShardedKey(), workB.id(), CommitStatus.NOT_FOUND), + CompleteCommit.create( + "computationId", workC.getShardedKey(), workC.id(), CommitStatus.NOT_FOUND)); + + // There should be no more commits in the queue + assertEquals(0, workCommitter.currentActiveCommitBytes()); + workCommitter.stop(); + } } 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 9c3d5c9c3ef3..1e995f4047c3 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 @@ -42,6 +42,8 @@ import java.util.function.Supplier; import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingCommitResponse; import org.apache.beam.runners.dataflow.worker.windmill.WindmillConnection; import org.apache.beam.runners.dataflow.worker.windmill.client.TriggeredScheduledExecutorService; import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream; @@ -1134,6 +1136,89 @@ public void testCommitWorkItem_multiplePhysicalStreams_multipleHandovers_halfClo assertTrue(commitWorkStream.awaitTermination(10, TimeUnit.SECONDS)); } + @Test + public void testCommit_multiKeyCommit() throws Exception { + testMultiKeyCommit(CommitStatus.OK); + } + + @Test + public void testCommit_multiKeyCommit_Failure() throws Exception { + testMultiKeyCommit(CommitStatus.NOT_FOUND); + } + + private void testMultiKeyCommit(CommitStatus commitStatus) throws Exception { + GrpcCommitWorkStream commitWorkStream = createCommitWorkStream(); + FakeWindmillGrpcService.CommitStreamInfo streamInfo = waitForConnectionAndConsumeHeader(); + + CompletableFuture commitStatusFuture = new CompletableFuture<>(); + + // 1. Construct two individual WorkItemCommitRequests + long shardingKey1 = 101L; + long workToken1 = 201L; + long cacheToken1 = 301L; + long shardingKey2 = 102L; + long workToken2 = 202L; + long cacheToken2 = 302L; + Windmill.WorkItemCommitRequest request1 = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setShardingKey(shardingKey1) + .setWorkToken(workToken1) + .setCacheToken(cacheToken1) + .build(); + Windmill.WorkItemCommitRequest request2 = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("key2")) + .setShardingKey(shardingKey2) + .setWorkToken(workToken2) + .setCacheToken(cacheToken2) + .build(); + + // 2. Wrap them into a MultiKeyWorkItemCommitRequest + Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest = + Windmill.MultiKeyWorkItemCommitRequest.newBuilder() + .addRequests(request1) + .addRequests(request2) + .build(); + + // 3. Commit the multi-key work item using the request batcher + try (WindmillStream.CommitWorkStream.RequestBatcher batcher = commitWorkStream.batcher()) { + assertTrue( + batcher.commitMultiKeyWorkItem( + COMPUTATION_ID, multiKeyRequest, commitStatusFuture::complete)); + } + + // 4. Receive and assert request properties on FakeWindmillGrpcService + Windmill.StreamingCommitWorkRequest request = streamInfo.requests.take(); + assertThat(request.getCommitChunkCount()).isEqualTo(1); + + Windmill.StreamingCommitRequestChunk chunk = request.getCommitChunk(0); + + // Assert that the commit type is correctly identified as COMMIT_TYPE_MULTI_KEY + assertThat(chunk.getCommitType()) + .isEqualTo(Windmill.StreamingCommitRequestChunk.CommitType.COMMIT_TYPE_MULTI_KEY); + + // Assert that the routing sharding key is mapped to the first request's sharding key + assertThat(chunk.getShardingKey()).isEqualTo(request1.getShardingKey()); + + // Assert that the serialized payload matches the input multiKeyRequest + Windmill.MultiKeyWorkItemCommitRequest parsedRequest = + Windmill.MultiKeyWorkItemCommitRequest.parseFrom(chunk.getSerializedWorkItemCommit()); + assertThat(parsedRequest).isEqualTo(multiKeyRequest); + + // 5. Respond with the generated requestId to complete the commit + long requestId = chunk.getRequestId(); + StreamingCommitResponse.Builder builder = + StreamingCommitResponse.newBuilder().addRequestId(requestId); + if (commitStatus != CommitStatus.OK) { + builder.addStatus(commitStatus); + } + streamInfo.responseObserver.onNext(builder.build()); + + // 6. Verify callback completed with expected sCommitStatus + assertThat(commitStatusFuture.get()).isEqualTo(commitStatus); + } + @Test public void testCommitWorkItem_stopsRetriesAfterDuration() throws Exception { int numCommits = 1; diff --git a/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto b/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto index aaa09c105fc3..a7a99e2ca5a1 100644 --- a/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto +++ b/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto @@ -678,9 +678,24 @@ message WorkItemCommitRequest { reserved 6, 23; } +message MultiKeyWorkItemCommitRequest { + optional Uint128Proto key_group = 7; + + repeated WorkItemCommitRequest requests = 1; + + repeated OutputMessageBundle output_messages = 2; + + repeated PubSubMessageBundle pubsub_messages = 3; + + repeated int64 finalize_ids = 4 [packed = true]; + + reserved 6; +} + message ComputationCommitWorkRequest { required string computation_id = 1; repeated WorkItemCommitRequest requests = 2; + repeated MultiKeyWorkItemCommitRequest multi_key_requests = 3; } message CommitWorkRequest { @@ -906,6 +921,14 @@ message StreamingCommitRequestChunk { // before handing off to the WindmillHost for processing. optional int64 remaining_bytes_for_work_item = 4; optional bytes serialized_work_item_commit = 5; + + enum CommitType { + COMMIT_TYPE_UNSPECIFIED = 0; + COMMIT_TYPE_SINGLE_KEY = 1; + COMMIT_TYPE_MULTI_KEY = 2; + } + + optional CommitType commit_type = 7; } message StreamingCommitResponse {