scwhittle commented on code in PR #38768:
URL: https://github.com/apache/beam/pull/38768#discussion_r3625878499
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java:
##########
@@ -113,8 +114,8 @@ public void commit(Commit commit) {
"Trying to queue commit on shutdown, failing
commit=[computationId={}, shardingKey={},"
+ " workId={} ].",
commit.computationId(),
- commit.work().getShardedKey(),
- commit.work().id());
+ commit.workBatch().get(0).getShardedKey(),
Review Comment:
I was going to say to check the size before blind access here, but maybe we
should just add a toString for the Commit instead and use it here?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java:
##########
@@ -32,20 +35,43 @@ public abstract class Commit {
public static Commit create(
WorkItemCommitRequest request, ComputationState computationState, Work
work) {
Preconditions.checkArgument(request.getSerializedSize() > 0);
- return new AutoValue_Commit(request, computationState, work);
+ return new AutoValue_Commit(
+ Optional.of(request), computationState, Optional.empty(),
ImmutableList.of(work));
+ }
+
+ public static Commit createMultiKey(
+ Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest,
+ ComputationState computationState,
+ ImmutableList<Work> workBatch) {
+ Preconditions.checkArgument(!workBatch.isEmpty());
+ return new AutoValue_Commit(
+ Optional.empty(), computationState, Optional.of(multiKeyRequest),
workBatch);
}
public final String computationId() {
return computationState().getComputationId();
}
- public abstract WorkItemCommitRequest request();
+ public abstract Optional<WorkItemCommitRequest> singleKeyRequest();
public abstract ComputationState computationState();
- public abstract Work work();
+ public abstract Optional<Windmill.MultiKeyWorkItemCommitRequest>
multiKeyRequest();
+
+ public abstract ImmutableList<Work> workBatch();
+
+ public final boolean isFailed() {
+ for (Work w : workBatch()) {
+ if (w.isFailed()) {
+ return true;
+ }
+ }
+ return false;
+ }
public final int getSize() {
Review Comment:
since we now maybe have multiple, size is confusing. How about
serializedByteSize or something more explicit
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java:
##########
@@ -32,20 +35,43 @@ public abstract class Commit {
public static Commit create(
WorkItemCommitRequest request, ComputationState computationState, Work
work) {
Preconditions.checkArgument(request.getSerializedSize() > 0);
- return new AutoValue_Commit(request, computationState, work);
+ return new AutoValue_Commit(
+ Optional.of(request), computationState, Optional.empty(),
ImmutableList.of(work));
+ }
+
+ public static Commit createMultiKey(
+ Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest,
+ ComputationState computationState,
+ ImmutableList<Work> workBatch) {
+ Preconditions.checkArgument(!workBatch.isEmpty());
+ return new AutoValue_Commit(
+ Optional.empty(), computationState, Optional.of(multiKeyRequest),
workBatch);
}
public final String computationId() {
return computationState().getComputationId();
}
- public abstract WorkItemCommitRequest request();
+ public abstract Optional<WorkItemCommitRequest> singleKeyRequest();
Review Comment:
should we remove AutoValue and internally use a boolean and nullable
references instead of Optional?
We can avoid the Optional allocations in hot path and can also enforce that
only one or the other is set.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java:
##########
@@ -147,8 +148,12 @@ private void drainCommitQueue() {
}
private void failCommit(Commit commit) {
Review Comment:
should we name this failQueuedCommit? Might be clearer that we don't need to
update the state or the active bytes
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java:
##########
@@ -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<CommitStatus> onDone) {
+ if (!canAccept(commitRequest.getSerializedSize() +
computation.length())) {
+ return false;
+ }
+ Preconditions.checkArgument(commitRequest.getRequestsCount() > 0);
Review Comment:
put top of function?
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java:
##########
@@ -141,12 +141,13 @@ public void testCommit() {
(CompleteCommit completeCommit, Commit commit) ->
completeCommit.computationId().equals(commit.computationId())
&& completeCommit.status() == Windmill.CommitStatus.OK
- && completeCommit.workId().equals(commit.work().id())
+ &&
completeCommit.workId().equals(commit.workBatch().get(0).id())
Review Comment:
add && commitWorkBatch.size() == 1
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java:
##########
@@ -409,10 +436,14 @@ public void testMultipleCommitSendersSingleStream() {
waitForExpectedSetSize(completeCommits, commits.size());
for (Commit commit : commits) {
- WorkItemCommitRequest request =
committed.get(commit.work().getWorkItem().getWorkToken());
+ WorkItemCommitRequest request =
Review Comment:
verify batch size
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java:
##########
@@ -474,4 +505,201 @@ public void testStop_drainsCommitQueue_concurrentCommit()
waitForExpectedSetSize(completeCommits, sentCommits.intValue());
}
+
+ @Test
+ public void testCommit_multiKeyCommitFailedWork() {
+ Set<CompleteCommit> 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));
+
+ workCommitter.stop();
+ }
+
+ @Test
+ public void testCommit_multiKeyCommitSuccess() {
+ Set<CompleteCommit> 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<Windmill.MultiKeyWorkItemCommitRequest> 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));
+
+ workCommitter.stop();
+ }
+
+ @Test
+ public void testCommit_multiKeyCommitStatusNotOK() {
+ Set<CompleteCommit> 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));
+
+ // Offer NOT_FOUND status for one of the works.
+ fakeWindmillServer.whenCommitWorkStreamCalled().put(workB.id(),
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<Windmill.MultiKeyWorkItemCommitRequest> 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
Review Comment:
this seems more like it is testing the fakeWindmillServer than the
StreamingEngineCommitter, it is the one fanning out from single multicommit to
separate CompleteCommit.
If you want to test StreamingEngineCommitter it seems like it should verify
that the error doesn't result in the committer retrying or something
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java:
##########
@@ -445,6 +460,37 @@ 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);
+ }
+
+ // Determine status for the batch.
+ // Default to OK, but if any of the works in the batch has an
offered status, use it.
+ Windmill.CommitStatus status = Windmill.CommitStatus.OK;
+ for (WorkItemCommitRequest workRequest :
elem.request.getRequestsList()) {
+ Windmill.CommitStatus offeredStatus =
+ streamingCommitsToOffer.remove(
Review Comment:
instead of this merging should we instead have a way to set the status for
multi-key requests?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java:
##########
@@ -32,20 +35,43 @@ public abstract class Commit {
public static Commit create(
WorkItemCommitRequest request, ComputationState computationState, Work
work) {
Preconditions.checkArgument(request.getSerializedSize() > 0);
- return new AutoValue_Commit(request, computationState, work);
+ return new AutoValue_Commit(
+ Optional.of(request), computationState, Optional.empty(),
ImmutableList.of(work));
+ }
+
+ public static Commit createMultiKey(
+ Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest,
+ ComputationState computationState,
+ ImmutableList<Work> workBatch) {
+ Preconditions.checkArgument(!workBatch.isEmpty());
+ return new AutoValue_Commit(
+ Optional.empty(), computationState, Optional.of(multiKeyRequest),
workBatch);
}
public final String computationId() {
return computationState().getComputationId();
}
- public abstract WorkItemCommitRequest request();
+ public abstract Optional<WorkItemCommitRequest> singleKeyRequest();
public abstract ComputationState computationState();
- public abstract Work work();
+ public abstract Optional<Windmill.MultiKeyWorkItemCommitRequest>
multiKeyRequest();
+
+ public abstract ImmutableList<Work> workBatch();
+
+ public final boolean isFailed() {
Review Comment:
Maybe this should be moved to StreamingEngineWorkCommitter just based upon
workbatch?
I think that would be clearer next to the failCommit logic which also works
on the batch.
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java:
##########
@@ -129,9 +129,9 @@ public void testCommit() {
for (Commit commit : commits) {
Windmill.WorkItemCommitRequest request =
- committed.get(commit.work().getWorkItem().getWorkToken());
+
committed.get(commit.workBatch().get(0).getWorkItem().getWorkToken());
Review Comment:
assert workBatch size is 1 also
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java:
##########
@@ -186,10 +186,14 @@ public void testCommit_sendsCommitsToStreamingEngine() {
waitForExpectedSetSize(completeCommits, 5);
for (Commit commit : commits) {
- WorkItemCommitRequest request =
committed.get(commit.work().getWorkItem().getWorkToken());
+ WorkItemCommitRequest request =
+
committed.get(commit.workBatch().get(0).getWorkItem().getWorkToken());
Review Comment:
verify the batch size is 1
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java:
##########
@@ -224,14 +228,24 @@ public void testCommit_handlesFailedCommits() {
waitForExpectedSetSize(completeCommits, 10);
for (Commit commit : commits) {
- if (commit.work().isFailed()) {
+ if (commit.isFailed()) {
Review Comment:
verify the batch size is 1
##########
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;
Review Comment:
is the reserved needed?
what about 5?
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java:
##########
@@ -423,6 +426,18 @@ public boolean commitWorkItem(
return true;
}
+ @Override
+ public boolean commitMultiKeyWorkItem(
+ String computation,
+ Windmill.MultiKeyWorkItemCommitRequest request,
+ Consumer<Windmill.CommitStatus> onDone) {
+ LOG.debug("commitWorkStream::commitMultiKeyWorkItem: {}", request);
+ if (multiKeyRequests.size() > 5) return false;
Review Comment:
is this used by any test? could we have a constant for this magic value if
so?
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java:
##########
@@ -282,11 +296,16 @@ public void
testCommit_handlesCompleteCommits_commitStatusNotOK() {
waitForExpectedSetSize(completeCommits, commits.size());
for (Commit commit : commits) {
- WorkItemCommitRequest request =
committed.get(commit.work().getWorkItem().getWorkToken());
+ WorkItemCommitRequest request =
Review Comment:
ditto
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStreamTest.java:
##########
@@ -1134,6 +1134,76 @@ public void
testCommitWorkItem_multiplePhysicalStreams_multipleHandovers_halfClo
assertTrue(commitWorkStream.awaitTermination(10, TimeUnit.SECONDS));
}
+ @Test
+ public void testCommit_multiKeyCommit() throws Exception {
+ GrpcCommitWorkStream commitWorkStream = createCommitWorkStream();
+ FakeWindmillGrpcService.CommitStreamInfo streamInfo =
waitForConnectionAndConsumeHeader();
+
+ CompletableFuture<Windmill.CommitStatus> 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();
+ streamInfo.responseObserver.onNext(
+
Windmill.StreamingCommitResponse.newBuilder().addRequestId(requestId).build());
+
+ // 6. Verify callback completed successfully with CommitStatus.OK
+ assertThat(commitStatusFuture.get()).isEqualTo(Windmill.CommitStatus.OK);
Review Comment:
similar test but return an error and verify plumbing?
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java:
##########
@@ -474,4 +505,201 @@ public void testStop_drainsCommitQueue_concurrentCommit()
waitForExpectedSetSize(completeCommits, sentCommits.intValue());
}
+
+ @Test
+ public void testCommit_multiKeyCommitFailedWork() {
+ Set<CompleteCommit> 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));
+
+ workCommitter.stop();
+ }
+
+ @Test
+ public void testCommit_multiKeyCommitSuccess() {
Review Comment:
nit: put this before the failed test
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
